fix(python): DataFrame.fill_null include unsigned integers (#5192)

pola-rs · Oct 13, 2022 · 7d45ba1 · 7d45ba1
1 parent fd2e92e
commit 7d45ba1
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 0 deletions.
diff --git a/py-polars/polars/internals/lazyframe/frame.py b/py-polars/polars/internals/lazyframe/frame.py
@@ -21,6 +21,10 @@
     Int64,
     PolarsDataType,
     Schema,
+    UInt8,
+    UInt16,
+    UInt32,
+    UInt64,
     Utf8,
     py_type_to_dtype,
 )
@@ -2166,6 +2170,10 @@ def fill_null(
                     dtypes.append(Int8)
                     dtypes.append(Int16)
                     dtypes.append(Int32)
+                    dtypes.append(UInt8)
+                    dtypes.append(UInt16)
+                    dtypes.append(UInt32)
+                    dtypes.append(UInt64)
                     dtypes.append(Float32)
                     dtypes.append(Float64)
             elif isinstance(value, float):
@@ -2175,6 +2183,10 @@ def fill_null(
                     dtypes.append(Int16)
                     dtypes.append(Int32)
                     dtypes.append(Int64)
+                    dtypes.append(UInt8)
+                    dtypes.append(UInt16)
+                    dtypes.append(UInt32)
+                    dtypes.append(UInt64)
                     dtypes.append(Float32)
                     dtypes.append(Float64)
             elif isinstance(value, str):

diff --git a/py-polars/tests/unit/test_series.py b/py-polars/tests/unit/test_series.py
@@ -742,6 +742,25 @@ def test_fill_null() -> None:
         "str": ["a", "b", "bar"],
         "bool": [True, True, False],
     }
+    df = pl.DataFrame({"a": [1, None, 2, None]})
+
+    out = df.with_columns(
+        [
+            pl.col("a").cast(pl.UInt8).alias("u8"),
+            pl.col("a").cast(pl.UInt16).alias("u16"),
+            pl.col("a").cast(pl.UInt32).alias("u32"),
+            pl.col("a").cast(pl.UInt64).alias("u64"),
+        ]
+    ).fill_null(3)
+
+    assert out.to_dict(False) == {
+        "a": [1, 3, 2, 3],
+        "u8": [1, 3, 2, 3],
+        "u16": [1, 3, 2, 3],
+        "u32": [1, 3, 2, 3],
+        "u64": [1, 3, 2, 3],
+    }
+    assert out.dtypes == [pl.Int64, pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64]
 
 
 def test_fill_nan() -> None: