Implementing Null Handling(is_finite, is_inf, is_nan)

python/ray/data/expressions.py

@@ -578,6 +578,22 @@ def abs(self) -> "UDFExpr":
         """
         return _create_pyarrow_compute_udf(pc.abs_checked)(self)
 
+    # NULL Handling
+    def is_nan(self) -> "UDFExpr":
+        return _create_pyarrow_compute_udf(pc.is_nan, return_dtype=DataType.bool())(
+            self
+        )
+
+    def is_finite(self) -> "UDFExpr":
+        return _create_pyarrow_compute_udf(pc.is_finite, return_dtype=DataType.bool())(
+            self
+        )
+
+    def is_inf(self) -> "UDFExpr":
+        return _create_pyarrow_compute_udf(pc.is_inf, return_dtype=DataType.bool())(
+            self
+        )
+
     @property
     def arr(self) -> "_ArrayNamespace":
         """Access array operations for this expression."""
@@ -1236,13 +1252,11 @@ def to_arrow(val):
         # Convert inputs to PyArrow and track pandas flags
         args_results = [to_arrow(arg) for arg in args]
         kwargs_results = {k: to_arrow(v) for k, v in kwargs.items()}
-
         converted_args = [v[0] for v in args_results]
         converted_kwargs = {k: v[0] for k, v in kwargs_results.items()}
         input_was_pandas = any(v[1] for v in args_results) or any(
             v[1] for v in kwargs_results.values()
         )
-
         # Call function with converted inputs
         result = fn(*converted_args, **converted_kwargs)
 

python/ray/data/tests/expressions/test_arithmetic.py

@@ -202,6 +202,58 @@ def test_age_group_calculation_with_dataset(self, ray_start_regular_shared):
         expected = pd.DataFrame({"age": [25, 17, 30], "age_group": [20, 10, 30]})
         assert rows_same(result, expected)
 
+    @pytest.mark.parametrize(
+        "test_data, expr_factory, expected_results, test_id",
+        [
+            # Test is_nan
+            pytest.param(
+                [
+                    {"x": float("nan")},
+                    {"x": -3.0},
+                    {"x": 0.0},
+                    {"x": 3.14},
+                    {"x": float("inf")},
+                    {"x": float("-inf")},
+                    {"x": None},
+                ],
+                lambda: col("x").is_nan(),
+                [True, False, False, False, False, False, None],
+                "is_nan",
+            ),
+            # Test is_finite
+            pytest.param(
+                [{"x": float("Inf")}, {"x": -3}, {"x": 0}],
+                lambda: col("x").is_finite(),
+                [False, True, True],
+                "is_finite",
+            ),
+            # Test is_inf
+            pytest.param(
+                [{"x": float("Inf")}, {"x": -3}, {"x": 0}],
+                lambda: col("x").is_inf(),
+                [True, False, False],
+                "is_infinite",
+            ),
+        ],
+    )
+    def test_with_column_null_handling_operations(
+        ray_start_regular_shared,
+        test_data,
+        expr_factory,
+        expected_results,
+        test_id,
+    ):
+        """Test null handling helper expressions."""
+        ds = ray.data.from_items(test_data)
+        expr = expr_factory()
+        result_df = ds.with_column("result", expr).to_pandas()
+
+        # Create expected dataframe
+        expected_df = pd.DataFrame(test_data)
+        expected_df["result"] = expected_results
+
+        assert rows_same(result_df, expected_df)
+
 
 if __name__ == "__main__":
     import sys