fix: Fix Decimal type fill_null (#19981)

pola-rs · Nov 25, 2024 · f0d087d · f0d087d
1 parent 7e78aa9
commit f0d087d
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/py-polars/polars/lazyframe/frame.py b/py-polars/polars/lazyframe/frame.py
@@ -5680,6 +5680,8 @@ def fill_null(
         │ 4   ┆ 13.0 │
         └─────┴──────┘
         """
+        from polars import Decimal
+
         dtypes: Sequence[PolarsDataType] | None
 
         if value is not None:
@@ -5703,6 +5705,7 @@ def infer_dtype(value: Any) -> PolarsDataType:
                     UInt64,
                     Float32,
                     Float64,
+                    Decimal,
                 ]
             elif isinstance(value, int):
                 dtypes = [Int64]

diff --git a/py-polars/tests/unit/operations/test_fill_null.py b/py-polars/tests/unit/operations/test_fill_null.py
@@ -31,6 +31,18 @@ def test_fill_null_static_schema_4843() -> None:
     assert df3.collect_schema() == {"a": pl.Int64, "b": pl.Int64}
 
 
+def test_fill_null_non_lit() -> None:
+    df = pl.DataFrame(
+        {
+            "a": pl.Series([1, None], dtype=pl.Int32),
+            "b": pl.Series([None, 2], dtype=pl.UInt32),
+            "c": pl.Series([None, 2], dtype=pl.Int64),
+            "d": pl.Series([None, 2], dtype=pl.Decimal),
+        }
+    )
+    assert df.fill_null(0).select(pl.all().null_count()).transpose().sum().item() == 0
+
+
 def test_fill_null_f32_with_lit() -> None:
     # ensure the literal integer does not upcast the f32 to an f64
     df = pl.DataFrame({"a": [1.1, 1.2]}, schema=[("a", pl.Float32)])