[Data] - Add to_pyarrow() to Expr (ray-project#57271)

goutamvenkat-anyscale · Future-Outlier · commit 1eed30a7b289 · 2025-12-07T16:50:53.000+08:00
## Why are these changes needed? Ray Data Expr to Pyarrow Compute Expression converter. ## Related issue number  ## Checks - [x] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [x] I've run pre-commit jobs to lint the changes in this PR. ([pre-commit setup](https://docs.ray.io/en/latest/ray-contribute/getting-involved.html#lint-and-formatting)) - [ ] I've included any doc changes needed for https://docs.ray.io/en/master/. - [ ] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [x] Unit tests - [ ] Release tests - [ ] This PR is not tested :( --------- Signed-off-by: Goutam V. <goutam@anyscale.com> Signed-off-by: Goutam <goutam@anyscale.com> Signed-off-by: Future-Outlier <eric901201@gmail.com>
diff --git a/python/ray/data/expressions.py b/python/ray/data/expressions.py
@@ -6,6 +6,8 @@
 from enum import Enum
 from typing import Any, Callable, Dict, List, Union
 
+import pyarrow
+
 from ray.data.block import BatchColumn
 from ray.data.datatype import DataType
 from ray.util.annotations import DeveloperAPI, PublicAPI
@@ -59,6 +61,118 @@ class Operation(Enum):
     NOT_IN = "not_in"
 
 
+class _ExprVisitor(ABC):
+    """Base visitor with generic dispatch for Ray Data expressions."""
+
+    def visit(self, expr: "Expr") -> Any:
+        if isinstance(expr, ColumnExpr):
+            return self.visit_column(expr)
+        elif isinstance(expr, LiteralExpr):
+            return self.visit_literal(expr)
+        elif isinstance(expr, BinaryExpr):
+            return self.visit_binary(expr)
+        elif isinstance(expr, UnaryExpr):
+            return self.visit_unary(expr)
+        elif isinstance(expr, AliasExpr):
+            return self.visit_alias(expr)
+        elif isinstance(expr, UDFExpr):
+            return self.visit_udf(expr)
+        elif isinstance(expr, DownloadExpr):
+            return self.visit_download(expr)
+        else:
+            raise TypeError(f"Unsupported expression type for conversion: {type(expr)}")
+
+    @abstractmethod
+    def visit_column(self, expr: "ColumnExpr") -> Any:
+        pass
+
+    @abstractmethod
+    def visit_literal(self, expr: "LiteralExpr") -> Any:
+        pass
+
+    @abstractmethod
+    def visit_binary(self, expr: "BinaryExpr") -> Any:
+        pass
+
+    @abstractmethod
+    def visit_unary(self, expr: "UnaryExpr") -> Any:
+        pass
+
+    @abstractmethod
+    def visit_alias(self, expr: "AliasExpr") -> Any:
+        pass
+
+    @abstractmethod
+    def visit_udf(self, expr: "UDFExpr") -> Any:
+        pass
+
+    @abstractmethod
+    def visit_download(self, expr: "DownloadExpr") -> Any:
+        pass
+
+
+class _PyArrowExpressionVisitor(_ExprVisitor):
+    """Visitor that converts Ray Data expressions to PyArrow compute expressions."""
+
+    def visit_column(self, expr: "ColumnExpr") -> "pyarrow.compute.Expression":
+        import pyarrow.compute as pc
+
+        return pc.field(expr.name)
+
+    def visit_literal(self, expr: "LiteralExpr") -> "pyarrow.compute.Expression":
+        import pyarrow.compute as pc
+
+        return pc.scalar(expr.value)
+
+    def visit_binary(self, expr: "BinaryExpr") -> "pyarrow.compute.Expression":
+        import pyarrow as pa
+        import pyarrow.compute as pc
+
+        if expr.op in (Operation.IN, Operation.NOT_IN):
+            left = self.visit(expr.left)
+            if isinstance(expr.right, LiteralExpr):
+                right_value = expr.right.value
+                right = (
+                    pa.array(right_value)
+                    if isinstance(right_value, list)
+                    else pa.array([right_value])
+                )
+            else:
+                raise ValueError(
+                    f"is_in/not_in operations require the right operand to be a "
+                    f"literal list, got {type(expr.right).__name__}."
+                )
+            result = pc.is_in(left, right)
+            return pc.invert(result) if expr.op == Operation.NOT_IN else result
+
+        left = self.visit(expr.left)
+        right = self.visit(expr.right)
+        from ray.data._expression_evaluator import _ARROW_EXPR_OPS_MAP
+
+        if expr.op in _ARROW_EXPR_OPS_MAP:
+            return _ARROW_EXPR_OPS_MAP[expr.op](left, right)
+        raise ValueError(f"Unsupported binary operation for PyArrow: {expr.op}")
+
+    def visit_unary(self, expr: "UnaryExpr") -> "pyarrow.compute.Expression":
+        operand = self.visit(expr.operand)
+        from ray.data._expression_evaluator import _ARROW_EXPR_OPS_MAP
+
+        if expr.op in _ARROW_EXPR_OPS_MAP:
+            return _ARROW_EXPR_OPS_MAP[expr.op](operand)
+        raise ValueError(f"Unsupported unary operation for PyArrow: {expr.op}")
+
+    def visit_alias(self, expr: "AliasExpr") -> "pyarrow.compute.Expression":
+        return self.visit(expr.expr)
+
+    def visit_udf(self, expr: "UDFExpr") -> "pyarrow.compute.Expression":
+        raise TypeError("UDF expressions cannot be converted to PyArrow expressions")
+
+    def visit_download(self, expr: "DownloadExpr") -> "pyarrow.compute.Expression":
+        raise TypeError(
+            "Download expressions cannot be converted to PyArrow expressions"
+        )
+
+
 @DeveloperAPI(stability="alpha")
 @dataclass(frozen=True)
 class Expr(ABC):
@@ -101,6 +215,18 @@ def structurally_equals(self, other: Any) -> bool:
         """Compare two expression ASTs for structural equality."""
         raise NotImplementedError
 
+    def to_pyarrow(self) -> "pyarrow.compute.Expression":
+        """Convert this Ray Data expression to a PyArrow compute expression.
+
+        Returns:
+            A PyArrow compute expression equivalent to this Ray Data expression.
+
+        Raises:
+            ValueError: If the expression contains operations not supported by PyArrow.
+            TypeError: If the expression type cannot be converted to PyArrow.
+        """
+        return _PyArrowExpressionVisitor().visit(self)
+
     def _bin(self, other: Any, op: Operation) -> "Expr":
         """Create a binary expression with the given operation.
 
diff --git a/python/ray/data/tests/test_expressions.py b/python/ray/data/tests/test_expressions.py
@@ -1,3 +1,5 @@
+import pyarrow as pa
+import pyarrow.compute as pc
 import pytest
 
 from ray.data.expressions import (
@@ -262,6 +264,160 @@ def test_complex_boolean_expressions(self):
         assert very_complex.op == Operation.AND
 
 
+class TestToPyArrow:
+    """Test conversion of Ray Data expressions to PyArrow compute expressions."""
+
+    @pytest.mark.parametrize(
+        "ray_expr, equivalent_pyarrow_expr, description",
+        [
+            # Basic expressions
+            (col("age"), lambda: pc.field("age"), "column reference"),
+            (lit(42), lambda: pc.scalar(42), "integer literal"),
+            (lit("hello"), lambda: pc.scalar("hello"), "string literal"),
+            # Arithmetic operations
+            (
+                col("x") + 5,
+                lambda: pc.add(pc.field("x"), pc.scalar(5)),
+                "addition",
+            ),
+            (
+                col("x") * 2,
+                lambda: pc.multiply(pc.field("x"), pc.scalar(2)),
+                "multiplication",
+            ),
+            # Comparison operations
+            (
+                col("age") > 18,
+                lambda: pc.greater(pc.field("age"), pc.scalar(18)),
+                "greater than",
+            ),
+            (
+                col("status") == "active",
+                lambda: pc.equal(pc.field("status"), pc.scalar("active")),
+                "equality",
+            ),
+            # Boolean operations
+            (
+                (col("age") > 18) & (col("age") < 65),
+                lambda: pc.and_kleene(
+                    pc.greater(pc.field("age"), pc.scalar(18)),
+                    pc.less(pc.field("age"), pc.scalar(65)),
+                ),
+                "logical AND",
+            ),
+            (
+                ~(col("active")),
+                lambda: pc.invert(pc.field("active")),
+                "logical NOT",
+            ),
+            # Unary operations
+            (
+                col("value").is_null(),
+                lambda: pc.is_null(pc.field("value")),
+                "is_null check",
+            ),
+            # In operations
+            (
+                col("status").is_in(["active", "pending"]),
+                lambda: pc.is_in(pc.field("status"), pa.array(["active", "pending"])),
+                "is_in with list",
+            ),
+            # Complex nested expressions
+            (
+                (col("price") * col("quantity")) + col("tax"),
+                lambda: pc.add(
+                    pc.multiply(pc.field("price"), pc.field("quantity")),
+                    pc.field("tax"),
+                ),
+                "nested arithmetic",
+            ),
+            # Alias expressions (should unwrap to inner expression)
+            (
+                (col("x") + 5).alias("result"),
+                lambda: pc.add(pc.field("x"), pc.scalar(5)),
+                "aliased expression",
+            ),
+        ],
+        ids=[
+            "col",
+            "int_lit",
+            "str_lit",
+            "add",
+            "mul",
+            "gt",
+            "eq",
+            "and",
+            "not",
+            "is_null",
+            "is_in",
+            "nested",
+            "alias",
+        ],
+    )
+    def test_to_pyarrow_equivalence(
+        self, ray_expr, equivalent_pyarrow_expr, description
+    ):
+        """Test that Ray Data expressions convert to equivalent PyArrow expressions.
+
+        This test documents the expected PyArrow expression for each Ray Data expression
+        and verifies correctness by comparing results on sample data.
+        """
+        import pyarrow.dataset as ds
+
+        # Convert Ray expression to PyArrow
+        converted = ray_expr.to_pyarrow()
+        expected = equivalent_pyarrow_expr()
+
+        # Both should be PyArrow expressions
+        assert isinstance(converted, pc.Expression)
+        assert isinstance(expected, pc.Expression)
+
+        # Verify they produce the same results on sample data
+        test_data = pa.table(
+            {
+                "age": [15, 25, 45, 70],
+                "x": [1, 2, 3, 4],
+                "price": [10.0, 20.0, 30.0, 40.0],
+                "quantity": [2, 3, 1, 5],
+                "tax": [1.0, 2.0, 3.0, 4.0],
+                "status": ["active", "pending", "inactive", "active"],
+                "value": [1, None, 3, None],
+                "active": [True, False, True, False],
+            }
+        )
+
+        dataset = ds.dataset(test_data)
+
+        try:
+            # For boolean expressions, compare filter results
+            result_converted = dataset.scanner(filter=converted).to_table()
+            result_expected = dataset.scanner(filter=expected).to_table()
+            assert result_converted.equals(
+                result_expected
+            ), f"Expressions produce different results for {description}"
+        except (TypeError, pa.lib.ArrowInvalid, pa.lib.ArrowNotImplementedError):
+            # For non-boolean expressions, just verify both are valid
+            pass
+
+    def test_to_pyarrow_unsupported_expressions(self):
+        """Test that unsupported expression types raise appropriate errors."""
+        from ray.data.datatype import DataType
+        from ray.data.expressions import UDFExpr
+
+        def dummy_fn(x):
+            return x
+
+        udf_expr = UDFExpr(
+            fn=dummy_fn,
+            args=[col("x")],
+            kwargs={},
+            data_type=DataType(int),
+        )
+
+        with pytest.raises(TypeError, match="UDF expressions cannot be converted"):
+            udf_expr.to_pyarrow()
+
+
 if __name__ == "__main__":
     import sys
 
