fix: Propagate nulls for float methods on all numeric types

pola-rs · Dec 20, 2024 · 9803e45 · 9803e45
1 parent 5d2d550
commit 9803e45
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 7 deletions.
diff --git a/crates/polars-core/src/series/mod.rs b/crates/polars-core/src/series/mod.rs
@@ -590,25 +590,29 @@ impl Series {
         }
     }
 
-    /// Check if float value is NaN (note this is different than missing/ null)
+    /// Check if numeric value is NaN (note this is different than missing/ null)
     pub fn is_nan(&self) -> PolarsResult<BooleanChunked> {
         match self.dtype() {
             DataType::Float32 => Ok(self.f32().unwrap().is_nan()),
             DataType::Float64 => Ok(self.f64().unwrap().is_nan()),
             dt if dt.is_numeric() => {
-                Ok(BooleanChunked::full(self.name().clone(), false, self.len()))
+                let arr = BooleanArray::full(self.len(), false, ArrowDataType::Boolean)
+                    .with_validity(self.rechunk_validity());
+                Ok(BooleanChunked::with_chunk(self.name().clone(), arr))
             },
             _ => polars_bail!(opq = is_nan, self.dtype()),
         }
     }
 
-    /// Check if float value is NaN (note this is different than missing/null)
+    /// Check if numeric value is NaN (note this is different than missing/null)
     pub fn is_not_nan(&self) -> PolarsResult<BooleanChunked> {
         match self.dtype() {
             DataType::Float32 => Ok(self.f32().unwrap().is_not_nan()),
             DataType::Float64 => Ok(self.f64().unwrap().is_not_nan()),
             dt if dt.is_numeric() => {
-                Ok(BooleanChunked::full(self.name().clone(), true, self.len()))
+                let arr = BooleanArray::full(self.len(), true, ArrowDataType::Boolean)
+                    .with_validity(self.rechunk_validity());
+                Ok(BooleanChunked::with_chunk(self.name().clone(), arr))
             },
             _ => polars_bail!(opq = is_not_nan, self.dtype()),
         }
@@ -620,19 +624,23 @@ impl Series {
             DataType::Float32 => Ok(self.f32().unwrap().is_finite()),
             DataType::Float64 => Ok(self.f64().unwrap().is_finite()),
             dt if dt.is_numeric() => {
-                Ok(BooleanChunked::full(self.name().clone(), true, self.len()))
+                let arr = BooleanArray::full(self.len(), true, ArrowDataType::Boolean)
+                    .with_validity(self.rechunk_validity());
+                Ok(BooleanChunked::with_chunk(self.name().clone(), arr))
             },
             _ => polars_bail!(opq = is_finite, self.dtype()),
         }
     }
 
-    /// Check if float value is infinite
+    /// Check if numeric value is infinite
     pub fn is_infinite(&self) -> PolarsResult<BooleanChunked> {
         match self.dtype() {
             DataType::Float32 => Ok(self.f32().unwrap().is_infinite()),
             DataType::Float64 => Ok(self.f64().unwrap().is_infinite()),
             dt if dt.is_numeric() => {
-                Ok(BooleanChunked::full(self.name().clone(), false, self.len()))
+                let arr = BooleanArray::full(self.len(), false, ArrowDataType::Boolean)
+                    .with_validity(self.rechunk_validity());
+                Ok(BooleanChunked::with_chunk(self.name().clone(), arr))
             },
             _ => polars_bail!(opq = is_infinite, self.dtype()),
         }

diff --git a/py-polars/tests/unit/series/test_series.py b/py-polars/tests/unit/series/test_series.py
@@ -1501,6 +1501,15 @@ def test_is_nan_is_not_nan(float_type: PolarsDataType) -> None:
     assert_series_equal(s.drop_nans(), pl.Series([1.0, None], dtype=float_type))
 
 
+def test_float_methods_on_ints() -> None:
+    # these float-specific methods work on non-float numeric types
+    s = pl.Series([1, None], dtype=pl.Int32)
+    assert_series_equal(s.is_finite(), pl.Series([True, None]))
+    assert_series_equal(s.is_infinite(), pl.Series([False, None]))
+    assert_series_equal(s.is_nan(), pl.Series([False, None]))
+    assert_series_equal(s.is_not_nan(), pl.Series([True, None]))
+
+
 def test_dot() -> None:
     s1 = pl.Series("a", [1, 2, 3])
     s2 = pl.Series("b", [4.0, 5.0, 6.0])