test: enhance null DataFrame tests to include date32 and date64 columns

Author: kosiew

python/tests/test_dataframe.py

@@ -17,7 +17,7 @@
 import os
 import re
 from typing import Any
-
+import datetime
 import pyarrow as pa
 import pyarrow.parquet as pq
 import pytest
@@ -128,7 +128,9 @@ def null_df():
         pa.array([4.5, 6.7, None, None], type=pa.float64()),
         pa.array(["a", None, "c", None], type=pa.string()),
         pa.array([True, None, False, None], type=pa.bool_()),
-    ], names=["int_col", "float_col", "str_col", "bool_col"])
+        pa.array([10957, None, 18993, None], type=pa.date32()),  # 2000-01-01, null, 2022-01-01, null
+        pa.array([946684800000, None, 1640995200000, None], type=pa.date64()),  # 2000-01-01, null, 2022-01-01, null
+    ], names=["int_col", "float_col", "str_col", "bool_col", "date32_col", "date64_col"])
 
     return ctx.create_dataframe([[batch]])
 
@@ -1524,7 +1526,7 @@ def add_with_parameter(df_internal, value: Any) -> DataFrame:
 def test_dataframe_repr_html_structure(df) -> None:
     """Test that DataFrame._repr_html_ produces expected HTML output structure."""
     import re
-
+    
     output = df._repr_html_()
 
     # Since we've added a fair bit of processing to the html output, lets just verify
@@ -1658,14 +1660,12 @@ def test_html_formatter_manual_format_html(clean_formatter_state):
     local_formatter = DataFrameHtmlFormatter(use_shared_styles=False)
 
     # Both calls should include styles
-    
     local_html_1 = local_formatter.format_html([batch], batch.schema)
     local_html_2 = local_formatter.format_html([batch], batch.schema)
 
     assert "<style>" in local_html_1
     assert "<style>" in local_html_2
-
-
+    
 def test_fill_null_basic(null_df):
     """Test basic fill_null functionality with a single value."""
     # Fill all nulls with 0
@@ -1674,12 +1674,12 @@ def test_fill_null_basic(null_df):
     result = filled_df.collect()[0]
 
     # Check that nulls were filled with 0 (or equivalent)
-    assert result.column(0).to_pylist() == [1, 0, 3, 0]
-    assert result.column(1).to_pylist() == [4.5, 6.7, 0.0, 0.0]
+    assert result.column(0) == pa.array([1, 0, 3, 0])
+    assert result.column(1) == pa.array([4.5, 6.7, 0.0, 0.0])
     # String column should be filled with "0"
-    assert result.column(2).to_pylist() == ["a", "0", "c", "0"]
+    assert result.column(2) == pa.array(["a", "0", "c", "0"])
     # Boolean column should be filled with False (0 converted to bool)
-    assert result.column(3).to_pylist() == [True, False, False, False]
+    assert result.column(3) == pa.array([True, False, False, False])
 
 
 def test_fill_null_subset(null_df):
@@ -1690,11 +1690,131 @@ def test_fill_null_subset(null_df):
     result = filled_df.collect()[0]
 
     # Check that nulls were filled only in specified columns
-    assert result.column(0).to_pylist() == [1, 0, 3, 0]
-    assert result.column(1).to_pylist() == [4.5, 6.7, 0.0, 0.0]
+    assert result.column(0) == pa.array([1, 0, 3, 0])
+    assert result.column(1) == pa.array([4.5, 6.7, 0.0, 0.0])
     # These should still have nulls
     assert None in result.column(2).to_pylist()
     assert None in result.column(3).to_pylist()
+    
+    def test_fill_null_str_column(null_df):
+        """Test filling nulls in string columns with different values."""
+        # Fill string nulls with a replacement string
+        filled_df = null_df.fill_null("N/A", subset=["str_col"])
+        
+        result = filled_df.collect()[0]
+        
+        # Check that string nulls were filled with "N/A"
+        assert result.column(2).to_pylist() == ["a", "N/A", "c", "N/A"]
+        
+        # Other columns should be unchanged
+        assert None in result.column(0).to_pylist()
+        assert None in result.column(1).to_pylist()
+        assert None in result.column(3).to_pylist()
+        
+        # Fill with an empty string
+        filled_df = null_df.fill_null("", subset=["str_col"])
+        result = filled_df.collect()[0]
+        assert result.column(2).to_pylist() == ["a", "", "c", ""]
+
+
+    def test_fill_null_bool_column(null_df):
+        """Test filling nulls in boolean columns with different values."""
+        # Fill bool nulls with True
+        filled_df = null_df.fill_null(True, subset=["bool_col"])
+        
+        result = filled_df.collect()[0]
+        
+        # Check that bool nulls were filled with True
+        assert result.column(3).to_pylist() == [True, True, False, True]
+        
+        # Other columns should be unchanged
+        assert None in result.column(0).to_pylist()
+        
+        # Fill bool nulls with False
+        filled_df = null_df.fill_null(False, subset=["bool_col"])
+        result = filled_df.collect()[0]
+        assert result.column(3).to_pylist() == [True, False, False, False]
+
+
+    def test_fill_null_date32_column(null_df):
+        """Test filling nulls in date32 columns."""
+        
+        # Fill date32 nulls with a specific date (1970-01-01)
+        epoch_date = datetime.date(1970, 1, 1)
+        filled_df = null_df.fill_null(epoch_date, subset=["date32_col"])
+        
+        result = filled_df.collect()[0]
+        
+        # Check that date32 nulls were filled with epoch date
+        dates = result.column(4).to_pylist()
+        assert dates[0] == datetime.date(2000, 1, 1)  # Original value
+        assert dates[1] == epoch_date                 # Filled value
+        assert dates[2] == datetime.date(2022, 1, 1)  # Original value
+        assert dates[3] == epoch_date                 # Filled value
+        
+        # Other date column should be unchanged
+        assert None in result.column(5).to_pylist()
+
+
+    def test_fill_null_date64_column(null_df):
+        """Test filling nulls in date64 columns."""
+        
+        # Fill date64 nulls with a specific date (1970-01-01)
+        epoch_date = datetime.date(1970, 1, 1)
+        filled_df = null_df.fill_null(epoch_date, subset=["date64_col"])
+        
+        result = filled_df.collect()[0]
+        
+        # Check that date64 nulls were filled with epoch date
+        dates = result.column(5).to_pylist()
+        assert dates[0] == datetime.date(2000, 1, 1)  # Original value
+        assert dates[1] == epoch_date                 # Filled value
+        assert dates[2] == datetime.date(2022, 1, 1)  # Original value
+        assert dates[3] == epoch_date                 # Filled value
+        
+        # Other date column should be unchanged
+        assert None in result.column(4).to_pylist()
+
+
+    def test_fill_null_type_coercion(null_df):
+        """Test type coercion when filling nulls with values of different types."""
+        # Try to fill string nulls with a number
+        filled_df = null_df.fill_null(42, subset=["str_col"])
+        
+        result = filled_df.collect()[0]
+        
+        # String nulls should be filled with string representation of the number
+        assert result.column(2).to_pylist() == ["a", "42", "c", "42"]
+        
+        # Try to fill bool nulls with a string that converts to True
+        filled_df = null_df.fill_null("true", subset=["bool_col"])
+        result = filled_df.collect()[0]
+        
+        # This behavior depends on the implementation - check it works without error
+        # but don't make assertions about exact conversion behavior
+        assert None not in result.column(3).to_pylist()
+
+
+    def test_fill_null_multiple_date_columns(null_df):
+        """Test filling nulls in both date column types simultaneously."""
+        
+        # Fill both date column types with the same date
+        test_date = datetime.date(2023, 12, 31)
+        filled_df = null_df.fill_null(test_date, subset=["date32_col", "date64_col"])
+        
+        result = filled_df.collect()[0]
+        
+        # Check both date columns were filled correctly
+        date32_vals = result.column(4).to_pylist()
+        date64_vals = result.column(5).to_pylist()
+        
+        assert None not in date32_vals
+        assert None not in date64_vals
+        
+        assert date32_vals[1] == test_date
+        assert date32_vals[3] == test_date
+        assert date64_vals[1] == test_date
+        assert date64_vals[3] == test_date
 
 
 def test_fill_null_specific_types(null_df):
@@ -1705,10 +1825,13 @@ def test_fill_null_specific_types(null_df):
     result = filled_df.collect()[0]
 
     # Check that nulls were filled appropriately by type
-    assert result.column(0).to_pylist() == [1, 0, 3, 0]  # Int gets 0 from "missing" conversion
-    assert result.column(1).to_pylist() == [4.5, 6.7, 0.0, 0.0]  # Float gets 0.0
-    assert result.column(2).to_pylist() == ["a", "missing", "c", "missing"]  # String gets "missing"
-    assert result.column(3).to_pylist() == [True, False, False, False]  # Bool gets False
+    
+    assert result.column(0).to_pylist() == [1, None, 3, None]  
+    assert result.column(1).to_pylist() == [4.5, 6.7, None, None] 
+    assert result.column(2).to_pylist() == ["a", "missing", "c", "missing"]
+    assert result.column(3).to_pylist() == [True, None, False, None]  # Bool gets False
+    assert result.column(4).to_pylist() == [datetime.date(2000, 1, 1), None, datetime.date(2022, 1, 1), None]
+    assert result.column(5).to_pylist() == [datetime.date(2000, 1, 1), None, datetime.date(2022, 1, 1), None]
 
 
 def test_fill_null_immutability(null_df):
@@ -1763,7 +1886,3 @@ def test_fill_null_all_null_column(ctx):
     # Check that all nulls were filled
     result = filled_df.collect()[0]
     assert result.column(1).to_pylist() == ["filled", "filled", "filled"]
-    
-    # Original should be unchanged
-    original = all_null_df.collect()[0]
-    assert original.column(1).null_count == 3

src/dataframe.rs

@@ -94,7 +94,7 @@ impl PyDataFrame {
 
 #[pymethods]
 impl PyDataFrame {
-    /// Enable selection for `df[col]`, `df[col1, col2, col2]`, and `df[[col1, col2, col3]]`
+    /// Enable selection for `df[col]`, `df[col1, col2, col3]`, and `df[[col1, col2, col3]]`
     fn __getitem__(&self, key: Bound<'_, PyAny>) -> PyDataFusionResult<Self> {
         if let Ok(key) = key.extract::<PyBackedStr>() {
             // df[col]