from __future__ import annotations import functools from abc import ABC, abstractmethod from dataclasses import dataclass, field from enum import Enum from typing import ( TYPE_CHECKING, Any, Callable, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union, ) import pyarrow import pyarrow.compute as pc from ray.data.block import BatchColumn from ray.data.datatype import DataType from ray.util.annotations import DeveloperAPI, PublicAPI if TYPE_CHECKING: from ray.data.namespace_expressions.arr_namespace import _ArrayNamespace from ray.data.namespace_expressions.dt_namespace import _DatetimeNamespace from ray.data.namespace_expressions.list_namespace import _ListNamespace from ray.data.namespace_expressions.string_namespace import _StringNamespace from ray.data.namespace_expressions.struct_namespace import _StructNamespace T = TypeVar("T") UDFCallable = Callable[..., "UDFExpr"] Decorated = Union[UDFCallable, Type[T]] @DeveloperAPI(stability="alpha") class Operation(Enum): """Enumeration of supported operations in expressions. This enum defines all the binary operations that can be performed between expressions, including arithmetic, comparison, and boolean operations. Attributes: ADD: Addition operation (+) SUB: Subtraction operation (-) MUL: Multiplication operation (*) DIV: Division operation (/) FLOORDIV: Floor division operation (//) GT: Greater than comparison (>) LT: Less than comparison (<) GE: Greater than or equal comparison (>=) LE: Less than or equal comparison (<=) EQ: Equality comparison (==) NE: Not equal comparison (!=) AND: Logical AND operation (&) OR: Logical OR operation (|) NOT: Logical NOT operation (~) IS_NULL: Check if value is null IS_NOT_NULL: Check if value is not null IN: Check if value is in a list NOT_IN: Check if value is not in a list """ ADD = "add" SUB = "sub" MUL = "mul" DIV = "div" MOD = "mod" FLOORDIV = "floordiv" GT = "gt" LT = "lt" GE = "ge" LE = "le" EQ = "eq" NE = "ne" AND = "and" OR = "or" NOT = "not" IS_NULL = "is_null" IS_NOT_NULL = "is_not_null" IN = "in" NOT_IN = "not_in" class _ExprVisitor(ABC, Generic[T]): """Base visitor with generic dispatch for Ray Data expressions.""" def visit(self, expr: "Expr") -> T: 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) elif isinstance(expr, StarExpr): return self.visit_star(expr) else: raise TypeError(f"Unsupported expression type for conversion: {type(expr)}") @abstractmethod def visit_column(self, expr: "ColumnExpr") -> T: pass @abstractmethod def visit_literal(self, expr: "LiteralExpr") -> T: pass @abstractmethod def visit_binary(self, expr: "BinaryExpr") -> T: pass @abstractmethod def visit_unary(self, expr: "UnaryExpr") -> T: pass @abstractmethod def visit_alias(self, expr: "AliasExpr") -> T: pass @abstractmethod def visit_udf(self, expr: "UDFExpr") -> T: pass @abstractmethod def visit_star(self, expr: "StarExpr") -> T: pass @abstractmethod def visit_download(self, expr: "DownloadExpr") -> T: pass class _PyArrowExpressionVisitor(_ExprVisitor["pyarrow.compute.Expression"]): """Visitor that converts Ray Data expressions to PyArrow compute expressions.""" def visit_column(self, expr: "ColumnExpr") -> "pyarrow.compute.Expression": return pc.field(expr.name) def visit_literal(self, expr: "LiteralExpr") -> "pyarrow.compute.Expression": return pc.scalar(expr.value) def visit_binary(self, expr: "BinaryExpr") -> "pyarrow.compute.Expression": import pyarrow as pa 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._internal.planner.plan_expression.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._internal.planner.plan_expression.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" ) def visit_star(self, expr: "StarExpr") -> "pyarrow.compute.Expression": raise TypeError("Star expressions cannot be converted to PyArrow expressions") @DeveloperAPI(stability="alpha") @dataclass(frozen=True) class Expr(ABC): """Base class for all expression nodes. This is the abstract base class that all expression types inherit from. It provides operator overloads for building complex expressions using standard Python operators. Expressions form a tree structure where each node represents an operation or value. The tree can be evaluated against data batches to compute results. Example: >>> from ray.data.expressions import col, lit >>> # Create an expression tree: (col("x") + 5) * col("y") >>> expr = (col("x") + lit(5)) * col("y") >>> # This creates a BinaryExpr with operation=MUL >>> # left=BinaryExpr(op=ADD, left=ColumnExpr("x"), right=LiteralExpr(5)) >>> # right=ColumnExpr("y") Note: This class should not be instantiated directly. Use the concrete subclasses like ColumnExpr, LiteralExpr, etc. """ data_type: DataType @property def name(self) -> str | None: """Get the name associated with this expression. Returns: The name for expressions that have one (ColumnExpr, AliasExpr), None otherwise. """ return None @abstractmethod 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 __repr__(self) -> str: """Return a tree-structured string representation of the expression. Returns: A multi-line string showing the expression tree structure using box-drawing characters for visual clarity. Example: >>> from ray.data.expressions import col, lit >>> expr = (col("x") + lit(5)) * col("y") >>> print(expr) MUL ├── left: ADD │ ├── left: COL('x') │ └── right: LIT(5) └── right: COL('y') """ from ray.data._internal.planner.plan_expression.expression_visitors import ( _TreeReprVisitor, ) return _TreeReprVisitor().visit(self) def _bin(self, other: Any, op: Operation) -> "Expr": """Create a binary expression with the given operation. Args: other: The right operand expression or literal value op: The operation to perform Returns: A new BinaryExpr representing the operation Note: If other is not an Expr, it will be automatically converted to a LiteralExpr. """ if not isinstance(other, Expr): other = LiteralExpr(other) return BinaryExpr(op, self, other) # # Arithmetic ops # def __add__(self, other: Any) -> "Expr": """Addition operator (+).""" return self._bin(other, Operation.ADD) def __radd__(self, other: Any) -> "Expr": """Reverse addition operator (for literal + expr).""" return LiteralExpr(other)._bin(self, Operation.ADD) def __sub__(self, other: Any) -> "Expr": """Subtraction operator (-).""" return self._bin(other, Operation.SUB) def __rsub__(self, other: Any) -> "Expr": """Reverse subtraction operator (for literal - expr).""" return LiteralExpr(other)._bin(self, Operation.SUB) def __mul__(self, other: Any) -> "Expr": """Multiplication operator (*).""" return self._bin(other, Operation.MUL) def __rmul__(self, other: Any) -> "Expr": """Reverse multiplication operator (for literal * expr).""" return LiteralExpr(other)._bin(self, Operation.MUL) def __mod__(self, other: Any): """Modulation operator (%).""" return self._bin(other, Operation.MOD) def __rmod__(self, other: Any): """Modulation operator (%).""" return LiteralExpr(other)._bin(self, Operation.MOD) def __truediv__(self, other: Any) -> "Expr": """Division operator (/).""" return self._bin(other, Operation.DIV) def __rtruediv__(self, other: Any) -> "Expr": """Reverse division operator (for literal / expr).""" return LiteralExpr(other)._bin(self, Operation.DIV) def __floordiv__(self, other: Any) -> "Expr": """Floor division operator (//).""" return self._bin(other, Operation.FLOORDIV) def __rfloordiv__(self, other: Any) -> "Expr": """Reverse floor division operator (for literal // expr).""" return LiteralExpr(other)._bin(self, Operation.FLOORDIV) # comparison def __gt__(self, other: Any) -> "Expr": """Greater than operator (>).""" return self._bin(other, Operation.GT) def __lt__(self, other: Any) -> "Expr": """Less than operator (<).""" return self._bin(other, Operation.LT) def __ge__(self, other: Any) -> "Expr": """Greater than or equal operator (>=).""" return self._bin(other, Operation.GE) def __le__(self, other: Any) -> "Expr": """Less than or equal operator (<=).""" return self._bin(other, Operation.LE) def __eq__(self, other: Any) -> "Expr": """Equality operator (==).""" return self._bin(other, Operation.EQ) def __ne__(self, other: Any) -> "Expr": """Not equal operator (!=).""" return self._bin(other, Operation.NE) # boolean def __and__(self, other: Any) -> "Expr": """Logical AND operator (&).""" return self._bin(other, Operation.AND) def __or__(self, other: Any) -> "Expr": """Logical OR operator (|).""" return self._bin(other, Operation.OR) def __invert__(self) -> "Expr": """Logical NOT operator (~).""" return UnaryExpr(Operation.NOT, self) # predicate methods def is_null(self) -> "Expr": """Check if the expression value is null.""" return UnaryExpr(Operation.IS_NULL, self) def is_not_null(self) -> "Expr": """Check if the expression value is not null.""" return UnaryExpr(Operation.IS_NOT_NULL, self) def is_in(self, values: Union[List[Any], "Expr"]) -> "Expr": """Check if the expression value is in a list of values.""" if not isinstance(values, Expr): values = LiteralExpr(values) return self._bin(values, Operation.IN) def not_in(self, values: Union[List[Any], "Expr"]) -> "Expr": """Check if the expression value is not in a list of values.""" if not isinstance(values, Expr): values = LiteralExpr(values) return self._bin(values, Operation.NOT_IN) def alias(self, name: str) -> "Expr": """Rename the expression. This method allows you to assign a new name to an expression result. This is particularly useful when you want to specify the output column name directly within the expression rather than as a separate parameter. Args: name: The new name for the expression Returns: An AliasExpr that wraps this expression with the specified name Example: >>> from ray.data.expressions import col, lit >>> # Create an expression with a new aliased name >>> expr = (col("price") * col("quantity")).alias("total") >>> # Can be used with Dataset operations that support named expressions """ return AliasExpr( data_type=self.data_type, expr=self, _name=name, _is_rename=False ) # rounding helpers def ceil(self) -> "UDFExpr": """Round values up to the nearest integer.""" return _create_pyarrow_compute_udf(pc.ceil)(self) def floor(self) -> "UDFExpr": """Round values down to the nearest integer.""" return _create_pyarrow_compute_udf(pc.floor)(self) def round(self) -> "UDFExpr": """Round values to the nearest integer using PyArrow semantics.""" return _create_pyarrow_compute_udf(pc.round)(self) def trunc(self) -> "UDFExpr": """Truncate fractional values toward zero.""" return _create_pyarrow_compute_udf(pc.trunc)(self) # logarithmic helpers def ln(self) -> "UDFExpr": """Compute the natural logarithm of the expression.""" return _create_pyarrow_compute_udf(pc.ln, return_dtype=DataType.float64())(self) def log10(self) -> "UDFExpr": """Compute the base-10 logarithm of the expression.""" return _create_pyarrow_compute_udf(pc.log10, return_dtype=DataType.float64())( self ) def log2(self) -> "UDFExpr": """Compute the base-2 logarithm of the expression.""" return _create_pyarrow_compute_udf(pc.log2, return_dtype=DataType.float64())( self ) def exp(self) -> "UDFExpr": """Compute the natural exponential of the expression.""" return _create_pyarrow_compute_udf(pc.exp, return_dtype=DataType.float64())( self ) # trigonometric helpers def sin(self) -> "UDFExpr": """Compute the sine of the expression (in radians).""" return _create_pyarrow_compute_udf(pc.sin, return_dtype=DataType.float64())( self ) def cos(self) -> "UDFExpr": """Compute the cosine of the expression (in radians).""" return _create_pyarrow_compute_udf(pc.cos, return_dtype=DataType.float64())( self ) def tan(self) -> "UDFExpr": """Compute the tangent of the expression (in radians).""" return _create_pyarrow_compute_udf(pc.tan, return_dtype=DataType.float64())( self ) def asin(self) -> "UDFExpr": """Compute the arcsine (inverse sine) of the expression, returning radians.""" return _create_pyarrow_compute_udf(pc.asin, return_dtype=DataType.float64())( self ) def acos(self) -> "UDFExpr": """Compute the arccosine (inverse cosine) of the expression, returning radians.""" return _create_pyarrow_compute_udf(pc.acos, return_dtype=DataType.float64())( self ) def atan(self) -> "UDFExpr": """Compute the arctangent (inverse tangent) of the expression, returning radians.""" return _create_pyarrow_compute_udf(pc.atan, return_dtype=DataType.float64())( self ) # arithmetic helpers def negate(self) -> "UDFExpr": """Compute the negation of the expression. Returns: A UDFExpr that computes the negation (multiplies values by -1). Example: >>> from ray.data.expressions import col >>> import ray >>> ds = ray.data.from_items([{"x": 5}, {"x": -3}]) >>> ds = ds.with_column("neg_x", col("x").negate()) >>> # Result: neg_x = [-5, 3] """ return _create_pyarrow_compute_udf(pc.negate_checked)(self) def sign(self) -> "UDFExpr": """Compute the sign of the expression. Returns: A UDFExpr that returns -1 for negative values, 0 for zero, and 1 for positive values. Example: >>> from ray.data.expressions import col >>> import ray >>> ds = ray.data.from_items([{"x": 5}, {"x": -3}, {"x": 0}]) >>> ds = ds.with_column("sign_x", col("x").sign()) >>> # Result: sign_x = [1, -1, 0] """ return _create_pyarrow_compute_udf(pc.sign)(self) def power(self, exponent: Any) -> "UDFExpr": """Raise the expression to the given power. Args: exponent: The exponent to raise the expression to. Returns: A UDFExpr that computes the power operation. Example: >>> from ray.data.expressions import col, lit >>> import ray >>> ds = ray.data.from_items([{"x": 2}, {"x": 3}]) >>> ds = ds.with_column("x_squared", col("x").power(2)) >>> # Result: x_squared = [4, 9] >>> ds = ds.with_column("x_cubed", col("x").power(3)) >>> # Result: x_cubed = [8, 27] """ return _create_pyarrow_compute_udf(pc.power)(self, exponent) def abs(self) -> "UDFExpr": """Compute the absolute value of the expression. Returns: A UDFExpr that computes the absolute value. Example: >>> from ray.data.expressions import col >>> import ray >>> ds = ray.data.from_items([{"x": 5}, {"x": -3}]) >>> ds = ds.with_column("abs_x", col("x").abs()) >>> # Result: abs_x = [5, 3] """ return _create_pyarrow_compute_udf(pc.abs_checked)(self) @property def arr(self) -> "_ArrayNamespace": """Access array operations for this expression.""" from ray.data.namespace_expressions.arr_namespace import _ArrayNamespace return _ArrayNamespace(self) @property def list(self) -> "_ListNamespace": """Access list operations for this expression. Returns: A _ListNamespace that provides list-specific operations for both PyArrow ``List`` and ``FixedSizeList`` columns. Example: >>> from ray.data.expressions import col >>> import ray >>> ds = ray.data.from_items([ ... {"items": [1, 2, 3]}, ... {"items": [4, 5]} ... ]) >>> ds = ds.with_column("num_items", col("items").list.len()) >>> ds = ds.with_column("first_item", col("items").list[0]) >>> ds = ds.with_column("slice", col("items").list[1:3]) """ from ray.data.namespace_expressions.list_namespace import _ListNamespace return _ListNamespace(self) @property def str(self) -> "_StringNamespace": """Access string operations for this expression. Returns: A _StringNamespace that provides string-specific operations. Example: >>> from ray.data.expressions import col >>> import ray >>> ds = ray.data.from_items([ ... {"name": "Alice"}, ... {"name": "Bob"} ... ]) >>> ds = ds.with_column("upper_name", col("name").str.upper()) >>> ds = ds.with_column("name_len", col("name").str.len()) >>> ds = ds.with_column("starts_a", col("name").str.starts_with("A")) """ from ray.data.namespace_expressions.string_namespace import _StringNamespace return _StringNamespace(self) @property def struct(self) -> "_StructNamespace": """Access struct operations for this expression. Returns: A _StructNamespace that provides struct-specific operations. Example: >>> from ray.data.expressions import col >>> import ray >>> import pyarrow as pa >>> ds = ray.data.from_arrow(pa.table({ ... "user": pa.array([ ... {"name": "Alice", "age": 30} ... ], type=pa.struct([ ... pa.field("name", pa.string()), ... pa.field("age", pa.int32()) ... ])) ... })) >>> ds = ds.with_column("age", col("user").struct["age"]) # doctest: +SKIP """ from ray.data.namespace_expressions.struct_namespace import _StructNamespace return _StructNamespace(self) @property def dt(self) -> "_DatetimeNamespace": """Access datetime operations for this expression.""" from ray.data.namespace_expressions.dt_namespace import _DatetimeNamespace return _DatetimeNamespace(self) def _unalias(self) -> "Expr": return self @DeveloperAPI(stability="alpha") @dataclass(frozen=True, eq=False, repr=False) class ColumnExpr(Expr): """Expression that references a column by name. This expression type represents a reference to an existing column in the dataset. When evaluated, it returns the values from the specified column. Args: name: The name of the column to reference Example: >>> from ray.data.expressions import col >>> # Reference the "age" column >>> age_expr = col("age") # Creates ColumnExpr(name="age") """ _name: str data_type: DataType = field(default_factory=lambda: DataType(object), init=False) @property def name(self) -> str: """Get the column name.""" return self._name def _rename(self, name: str): return AliasExpr(self.data_type, self, name, _is_rename=True) def structurally_equals(self, other: Any) -> bool: return isinstance(other, ColumnExpr) and self.name == other.name @DeveloperAPI(stability="alpha") @dataclass(frozen=True, eq=False, repr=False) class LiteralExpr(Expr): """Expression that represents a constant scalar value. This expression type represents a literal value that will be broadcast to all rows when evaluated. The value can be any Python object. Args: value: The constant value to represent Example: >>> from ray.data.expressions import lit >>> import numpy as np >>> # Create a literal value >>> five = lit(5) # Creates LiteralExpr(value=5) >>> name = lit("John") # Creates LiteralExpr(value="John") >>> numpy_val = lit(np.int32(42)) # Creates LiteralExpr with numpy type """ value: Any data_type: DataType = field(init=False) def __post_init__(self): # Infer the type from the value using DataType.infer_dtype inferred_dtype = DataType.infer_dtype(self.value) # Use object.__setattr__ since the dataclass is frozen object.__setattr__(self, "data_type", inferred_dtype) def structurally_equals(self, other: Any) -> bool: return ( isinstance(other, LiteralExpr) and self.value == other.value and type(self.value) is type(other.value) ) @DeveloperAPI(stability="alpha") @dataclass(frozen=True, eq=False, repr=False) class BinaryExpr(Expr): """Expression that represents a binary operation between two expressions. This expression type represents an operation with two operands (left and right). The operation is specified by the `op` field, which must be one of the supported operations from the Operation enum. Args: op: The operation to perform (from Operation enum) left: The left operand expression right: The right operand expression Example: >>> from ray.data.expressions import col, lit, Operation >>> # Manually create a binary expression (usually done via operators) >>> expr = BinaryExpr(Operation.ADD, col("x"), lit(5)) >>> # This is equivalent to: col("x") + lit(5) """ op: Operation left: Expr right: Expr data_type: DataType = field(default_factory=lambda: DataType(object), init=False) def structurally_equals(self, other: Any) -> bool: return ( isinstance(other, BinaryExpr) and self.op is other.op and self.left.structurally_equals(other.left) and self.right.structurally_equals(other.right) ) @DeveloperAPI(stability="alpha") @dataclass(frozen=True, eq=False, repr=False) class UnaryExpr(Expr): """Expression that represents a unary operation on a single expression. This expression type represents an operation with one operand. Common unary operations include logical NOT, IS NULL, IS NOT NULL, etc. Args: op: The operation to perform (from Operation enum) operand: The operand expression Example: >>> from ray.data.expressions import col >>> # Check if a column is null >>> expr = col("age").is_null() # Creates UnaryExpr(IS_NULL, col("age")) >>> # Logical not >>> expr = ~(col("active")) # Creates UnaryExpr(NOT, col("active")) """ op: Operation operand: Expr # Default to bool return dtype for unary operations like is_null() and NOT. # This enables chaining operations such as col("x").is_not_null().alias("valid"), # where downstream expressions (like AliasExpr) need the data type. data_type: DataType = field(default_factory=lambda: DataType.bool(), init=False) def structurally_equals(self, other: Any) -> bool: return ( isinstance(other, UnaryExpr) and self.op is other.op and self.operand.structurally_equals(other.operand) ) @dataclass(frozen=True) class _CallableClassSpec: """Specification for a callable class UDF. This dataclass captures the class type and constructor arguments needed to instantiate a callable class UDF on an actor. It consolidates the callable class metadata that was previously spread across multiple fields. Attributes: cls: The original callable class type args: Positional arguments for the constructor kwargs: Keyword arguments for the constructor _cached_key: Pre-computed key that survives serialization """ cls: type args: Tuple[Any, ...] = () kwargs: Dict[str, Any] = field(default_factory=dict) _cached_key: Optional[Tuple] = field(default=None, compare=False, repr=False) def __post_init__(self): """Pre-compute and cache the key at construction time. This ensures the same key survives serialization, since the cached key tuple (containing the already-computed repr strings) gets pickled and unpickled as-is. """ if self._cached_key is None: class_id = f"{self.cls.__module__}.{self.cls.__qualname__}" try: key = ( class_id, self.args, tuple(sorted(self.kwargs.items())), ) # Verify the key is actually hashable (args may contain lists) hash(key) except TypeError: # Fallback for unhashable args/kwargs - use repr for comparison key = (class_id, repr(self.args), repr(self.kwargs)) # Use object.__setattr__ since dataclass is frozen object.__setattr__(self, "_cached_key", key) def make_key(self) -> Tuple: """Return the pre-computed hashable key for UDF instance lookup. The key uniquely identifies a UDF by its class and constructor arguments. This ensures that the same class with different constructor args (e.g., Multiplier(2) vs Multiplier(3)) are treated as distinct UDFs. Returns: A hashable tuple that uniquely identifies this UDF configuration. """ return self._cached_key class _CallableClassUDF: """A wrapper that makes callable class UDFs appear as regular functions. This class wraps callable class UDFs for use in expressions. It provides an `init()` method that should be called at actor startup via `init_fn` to instantiate the underlying class before any blocks are processed. Key responsibilities: 1. Store the callable class and constructor arguments 2. Provide init() for actor startup initialization 3. Handle async bridging for coroutine/async generator UDFs 4. Reuse the same instance across all calls (actor semantics) Example: >>> @udf(return_dtype=DataType.int32()) ... class AddOffset: ... def __init__(self, offset=1): ... self.offset = offset ... def __call__(self, x): ... return pc.add(x, self.offset) >>> >>> add_five = AddOffset(5) # Creates _CallableClassUDF internally >>> expr = add_five(col("value")) # Creates UDFExpr with fn=_CallableClassUDF """ def __init__( self, cls: type, ctor_args: Tuple[Any, ...], ctor_kwargs: Dict[str, Any], return_dtype: DataType, ): """Initialize the _CallableClassUDF wrapper. Args: cls: The original callable class ctor_args: Constructor positional arguments ctor_kwargs: Constructor keyword arguments return_dtype: The return data type for schema inference """ self._cls = cls self._ctor_args = ctor_args self._ctor_kwargs = ctor_kwargs self._return_dtype = return_dtype # Instance created by init() at actor startup self._instance = None # Cache the spec to avoid creating new instances on each access self._callable_class_spec = _CallableClassSpec( cls=cls, args=ctor_args, kwargs=ctor_kwargs, ) @property def __name__(self) -> str: """Return the original class name for error messages.""" return self._cls.__name__ @property def callable_class_spec(self) -> _CallableClassSpec: """Return the callable class spec for this UDF. Used for deduplication when the same UDF appears multiple times in an expression tree. """ return self._callable_class_spec def init(self) -> None: """Initialize the UDF instance. Called at actor startup via init_fn. This ensures the callable class is instantiated before any blocks are processed, matching the behavior of map_batches callable classes. """ if self._instance is None: self._instance = self._cls(*self._ctor_args, **self._ctor_kwargs) def __call__(self, *args: Any, **kwargs: Any) -> Any: """Call the UDF instance. Args: *args: Evaluated expression arguments (PyArrow arrays, etc.) **kwargs: Evaluated expression keyword arguments Returns: The result of calling the UDF instance Raises: RuntimeError: If init() was not called before __call__ """ if self._instance is None: raise RuntimeError( f"_CallableClassUDF '{self._cls.__name__}' was not initialized. " f"init() must be called before __call__. This typically happens " f"via init_fn at actor startup." ) from ray.data.util.expression_utils import _call_udf_instance_with_async_bridge # Call instance directly, handling async if needed return _call_udf_instance_with_async_bridge(self._instance, *args, **kwargs) @DeveloperAPI(stability="alpha") @dataclass(frozen=True, eq=False, repr=False) class UDFExpr(Expr): """Expression that represents a user-defined function call. This expression type wraps a UDF with schema inference capabilities, allowing UDFs to be used seamlessly within the expression system. UDFs operate on batches of data, where each column argument is passed as a PyArrow Array containing multiple values from that column across the batch. Args: fn: The user-defined function to call. For callable classes, this is an _CallableClassUDF instance that handles lazy instantiation internally. args: List of argument expressions (positional arguments) kwargs: Dictionary of keyword argument expressions Example: >>> from ray.data.expressions import col, udf >>> import pyarrow as pa >>> import pyarrow.compute as pc >>> from ray.data.datatype import DataType >>> >>> @udf(return_dtype=DataType.int32()) ... def add_one(x: pa.Array) -> pa.Array: ... return pc.add(x, 1) >>> >>> # Use in expressions >>> expr = add_one(col("value")) >>> # Callable class example >>> @udf(return_dtype=DataType.int32()) ... class AddOffset: ... def __init__(self, offset=1): ... self.offset = offset ... def __call__(self, x: pa.Array) -> pa.Array: ... return pc.add(x, self.offset) >>> >>> # Use callable class >>> add_five = AddOffset(5) >>> expr = add_five(col("value")) """ fn: Callable[..., BatchColumn] # Can be regular function OR _CallableClassUDF args: List[Expr] kwargs: Dict[str, Expr] @property def callable_class_spec(self) -> Optional[_CallableClassSpec]: """Return callable_class_spec if fn is an _CallableClassUDF, else None. This property maintains backward compatibility with code that checks for callable_class_spec. """ if isinstance(self.fn, _CallableClassUDF): return self.fn.callable_class_spec return None def structurally_equals(self, other: Any) -> bool: if not isinstance(other, UDFExpr): return False # For callable class UDFs (_CallableClassUDF), compare the callable_class_spec. # For regular function UDFs, compare fn directly. if isinstance(self.fn, _CallableClassUDF): if not isinstance(other.fn, _CallableClassUDF): return False if self.fn.callable_class_spec != other.fn.callable_class_spec: return False else: if self.fn != other.fn: return False return ( len(self.args) == len(other.args) and all(a.structurally_equals(b) for a, b in zip(self.args, other.args)) and self.kwargs.keys() == other.kwargs.keys() and all( self.kwargs[k].structurally_equals(other.kwargs[k]) for k in self.kwargs.keys() ) ) def _create_udf_callable( fn: Callable[..., BatchColumn], return_dtype: DataType, ) -> Callable[..., UDFExpr]: """Create a callable that generates UDFExpr when called with expressions. Args: fn: The user-defined function to wrap. Can be a regular function or an _CallableClassUDF instance (for callable classes). return_dtype: The return data type of the UDF Returns: A callable that creates UDFExpr instances when called with expressions """ def udf_callable(*args, **kwargs) -> UDFExpr: # Convert arguments to expressions if they aren't already expr_args = [] for arg in args: if isinstance(arg, Expr): expr_args.append(arg) else: expr_args.append(LiteralExpr(arg)) expr_kwargs = {} for k, v in kwargs.items(): if isinstance(v, Expr): expr_kwargs[k] = v else: expr_kwargs[k] = LiteralExpr(v) return UDFExpr( fn=fn, args=expr_args, kwargs=expr_kwargs, data_type=return_dtype, ) # Preserve original function metadata functools.update_wrapper(udf_callable, fn) # Store the original function for access if needed udf_callable._original_fn = fn return udf_callable @PublicAPI(stability="alpha") def udf(return_dtype: DataType) -> Callable[..., UDFExpr]: """ Decorator to convert a UDF into an expression-compatible function. This decorator allows UDFs to be used seamlessly within the expression system, enabling schema inference and integration with other expressions. IMPORTANT: UDFs operate on batches of data, not individual rows. When your UDF is called, each column argument will be passed as a PyArrow Array containing multiple values from that column across the batch. Under the hood, when working with multiple columns, they get translated to PyArrow arrays (one array per column). Args: return_dtype: The data type of the return value of the UDF Returns: A callable that creates UDFExpr instances when called with expressions Example: >>> from ray.data.expressions import col, udf >>> import pyarrow as pa >>> import pyarrow.compute as pc >>> import ray >>> >>> # UDF that operates on a batch of values (PyArrow Array) >>> @udf(return_dtype=DataType.int32()) ... def add_one(x: pa.Array) -> pa.Array: ... return pc.add(x, 1) # Vectorized operation on the entire Array >>> >>> # UDF that combines multiple columns (each as a PyArrow Array) >>> @udf(return_dtype=DataType.string()) ... def format_name(first: pa.Array, last: pa.Array) -> pa.Array: ... return pc.binary_join_element_wise(first, last, " ") # Vectorized string concatenation >>> >>> # Callable class UDF >>> @udf(return_dtype=DataType.int32()) ... class AddOffset: ... def __init__(self, offset=1): ... self.offset = offset ... def __call__(self, x: pa.Array) -> pa.Array: ... return pc.add(x, self.offset) >>> >>> # Use in dataset operations >>> ds = ray.data.from_items([ ... {"value": 5, "first": "John", "last": "Doe"}, ... {"value": 10, "first": "Jane", "last": "Smith"} ... ]) >>> >>> # Single column transformation (operates on batches) >>> ds_incremented = ds.with_column("value_plus_one", add_one(col("value"))) >>> >>> # Multi-column transformation (each column becomes a PyArrow Array) >>> ds_formatted = ds.with_column("full_name", format_name(col("first"), col("last"))) >>> >>> # Callable class usage >>> add_five = AddOffset(5) >>> ds_with_offset = ds.with_column("value_plus_five", add_five(col("value"))) >>> >>> # Can also be used in complex expressions >>> ds_complex = ds.with_column("doubled_plus_one", add_one(col("value")) * 2) """ def decorator( func_or_class: Union[Callable[..., BatchColumn], Type[T]] ) -> Decorated: # Check if this is a callable class (has __call__ method defined) if isinstance(func_or_class, type) and issubclass(func_or_class, Callable): # Wrapper that delays instantiation and returns expressions instead of executing. # Without this, MyClass(args) would instantiate on the driver and # instance(col(...)) would try to execute rather than building an expression. class ExpressionAwareCallableClass: """Intercepts callable class instantiation to delay until actor execution. Allows natural syntax like: add_five = AddOffset(5) ds.with_column("result", add_five(col("x"))) When instantiated, creates an _CallableClassUDF that is completely self-contained - it handles lazy instantiation and async bridging internally. From the planner's perspective, this is just a regular callable function. """ def __init__(self, *args, **kwargs): # Create an _CallableClassUDF that is self-contained. # It lazily instantiates the class on first call (on the worker) # and handles async bridging internally. self._expr_udf = _CallableClassUDF( cls=func_or_class, ctor_args=args, ctor_kwargs=kwargs, return_dtype=return_dtype, ) def __call__(self, *call_args, **call_kwargs): # Create UDFExpr with fn=_CallableClassUDF # The _CallableClassUDF is self-contained - no external setup needed return _create_udf_callable( self._expr_udf, return_dtype, )(*call_args, **call_kwargs) # Preserve the original class name and module for better error messages ExpressionAwareCallableClass.__name__ = func_or_class.__name__ ExpressionAwareCallableClass.__qualname__ = func_or_class.__qualname__ ExpressionAwareCallableClass.__module__ = func_or_class.__module__ return ExpressionAwareCallableClass else: # Regular function return _create_udf_callable(func_or_class, return_dtype) return decorator def _create_pyarrow_wrapper( fn: Callable[..., BatchColumn] ) -> Callable[..., BatchColumn]: """Wrap a PyArrow compute function to auto-convert inputs to PyArrow format. This wrapper ensures that pandas Series and numpy arrays are converted to PyArrow Arrays before being passed to the function, enabling PyArrow compute functions to work seamlessly with any block format. Args: fn: The PyArrow compute function to wrap Returns: A wrapped function that handles format conversion """ @functools.wraps(fn) def arrow_wrapper(*args, **kwargs): import numpy as np import pandas as pd import pyarrow as pa def to_arrow(val): """Convert a value to PyArrow Array if needed.""" if isinstance(val, (pa.Array, pa.ChunkedArray)): return val, False elif isinstance(val, pd.Series): return pa.Array.from_pandas(val), True elif isinstance(val, np.ndarray): return pa.array(val), False else: return val, False # 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) # Convert result back to pandas if input was pandas if input_was_pandas and isinstance(result, (pa.Array, pa.ChunkedArray)): result = result.to_pandas() return result return arrow_wrapper @PublicAPI(stability="alpha") def pyarrow_udf(return_dtype: DataType) -> Callable[..., UDFExpr]: """Decorator for PyArrow compute functions with automatic format conversion. This decorator wraps PyArrow compute functions to automatically convert pandas Series and numpy arrays to PyArrow Arrays, ensuring the function works seamlessly regardless of the underlying block format (pandas, arrow, or items). Used internally by namespace methods (list, str, struct) that wrap PyArrow compute functions. Args: return_dtype: The data type of the return value Returns: A callable that creates UDFExpr instances with automatic conversion """ def decorator(func: Callable[..., BatchColumn]) -> Callable[..., UDFExpr]: # Wrap the function with PyArrow conversion logic wrapped_fn = _create_pyarrow_wrapper(func) # Create UDFExpr callable using the wrapped function return _create_udf_callable(wrapped_fn, return_dtype) return decorator def _create_pyarrow_compute_udf( pc_func: Callable[..., pyarrow.Array], return_dtype: DataType | None = None, ) -> Callable[..., "UDFExpr"]: """Create an expression UDF backed by a PyArrow compute function.""" def wrapper(expr: "Expr", *positional: Any, **kwargs: Any) -> "UDFExpr": @pyarrow_udf(return_dtype=return_dtype or expr.data_type) def udf(arr: pyarrow.Array) -> pyarrow.Array: return pc_func(arr, *positional, **kwargs) return udf(expr) return wrapper @DeveloperAPI(stability="alpha") @dataclass(frozen=True, eq=False, repr=False) class DownloadExpr(Expr): """Expression that represents a download operation.""" uri_column_name: str filesystem: "pyarrow.fs.FileSystem" = None data_type: DataType = field(default_factory=lambda: DataType.binary(), init=False) def structurally_equals(self, other: Any) -> bool: return ( isinstance(other, DownloadExpr) and self.uri_column_name == other.uri_column_name ) @DeveloperAPI(stability="alpha") @dataclass(frozen=True, eq=False, repr=False) class AliasExpr(Expr): """Expression that represents an alias for an expression.""" expr: Expr _name: str _is_rename: bool @property def name(self) -> str: """Get the alias name.""" return self._name def alias(self, name: str) -> "Expr": # Always unalias before creating new one return AliasExpr( self.expr.data_type, self.expr, _name=name, _is_rename=self._is_rename ) def _unalias(self) -> "Expr": return self.expr def structurally_equals(self, other: Any) -> bool: return ( isinstance(other, AliasExpr) and self.expr.structurally_equals(other.expr) and self.name == other.name and self._is_rename == self._is_rename ) @DeveloperAPI(stability="alpha") @dataclass(frozen=True, eq=False, repr=False) class StarExpr(Expr): """Expression that represents all columns from the input. This is a special expression used in projections to indicate that all existing columns should be preserved at this position in the output. It's typically used internally by operations like with_column() and rename_columns() to maintain existing columns. Example: When with_column("new_col", expr) is called, it creates: Project(exprs=[star(), expr.alias("new_col")]) This means: keep all existing columns, then add/overwrite "new_col" """ # TODO: Add UnresolvedExpr. Both StarExpr and UnresolvedExpr won't have a defined data_type. data_type: DataType = field(default_factory=lambda: DataType(object), init=False) def structurally_equals(self, other: Any) -> bool: return isinstance(other, StarExpr) @PublicAPI(stability="beta") def col(name: str) -> ColumnExpr: """ Reference an existing column by name. This is the primary way to reference columns in expressions. The returned expression will extract values from the specified column when evaluated. Args: name: The name of the column to reference Returns: A ColumnExpr that references the specified column Example: >>> from ray.data.expressions import col >>> # Reference columns in an expression >>> expr = col("price") * col("quantity") >>> >>> # Use with Dataset.with_column() >>> import ray >>> ds = ray.data.from_items([{"price": 10, "quantity": 2}]) >>> ds = ds.with_column("total", col("price") * col("quantity")) """ return ColumnExpr(name) @PublicAPI(stability="beta") def lit(value: Any) -> LiteralExpr: """ Create a literal expression from a constant value. This creates an expression that represents a constant scalar value. The value will be broadcast to all rows when the expression is evaluated. Args: value: The constant value to represent. Can be any Python object (int, float, str, bool, etc.) Returns: A LiteralExpr containing the specified value Example: >>> from ray.data.expressions import col, lit >>> # Create literals of different types >>> five = lit(5) >>> pi = lit(3.14159) >>> name = lit("Alice") >>> flag = lit(True) >>> >>> # Use in expressions >>> expr = col("age") + lit(1) # Add 1 to age column >>> >>> # Use with Dataset.with_column() >>> import ray >>> ds = ray.data.from_items([{"age": 25}, {"age": 30}]) >>> ds = ds.with_column("age_plus_one", col("age") + lit(1)) """ return LiteralExpr(value) # TODO remove @DeveloperAPI(stability="alpha") def star() -> StarExpr: """ References all input columns from the input. This is a special expression used in projections to preserve all existing columns. It's typically used with operations that want to add or modify columns while keeping the rest. Returns: A StarExpr that represents all input columns. """ return StarExpr() @PublicAPI(stability="alpha") def download( uri_column_name: str, *, filesystem: Optional["pyarrow.fs.FileSystem"] = None, ) -> DownloadExpr: """ Create a download expression that downloads content from URIs. This creates an expression that will download bytes from URIs stored in a specified column. When evaluated, it will fetch the content from each URI and return the downloaded bytes. Args: uri_column_name: The name of the column containing URIs to download from filesystem: PyArrow filesystem to use for reading remote files. If None, the filesystem is auto-detected from the path scheme. Returns: A DownloadExpr that will download content from the specified URI column Example: >>> from ray.data.expressions import download >>> import ray >>> # Create dataset with URIs >>> ds = ray.data.from_items([ ... {"uri": "s3://bucket/file1.jpg", "id": "1"}, ... {"uri": "s3://bucket/file2.jpg", "id": "2"} ... ]) >>> # Add downloaded bytes column >>> ds_with_bytes = ds.with_column("bytes", download("uri")) """ return DownloadExpr(uri_column_name=uri_column_name, filesystem=filesystem) # ────────────────────────────────────── # Public API for evaluation # ────────────────────────────────────── # Note: Implementation details are in _expression_evaluator.py # Re-export eval_expr for public use __all__ = [ "Operation", "Expr", "ColumnExpr", "LiteralExpr", "BinaryExpr", "UnaryExpr", "UDFExpr", "DownloadExpr", "AliasExpr", "StarExpr", "pyarrow_udf", "udf", "col", "lit", "download", "star", "_ArrayNamespace", "_ListNamespace", "_StringNamespace", "_StructNamespace", "_DatetimeNamespace", ] def __getattr__(name: str): """Lazy import of namespace classes to avoid circular imports.""" if name == "_ArrayNamespace": from ray.data.namespace_expressions.arr_namespace import _ArrayNamespace return _ArrayNamespace elif name == "_ListNamespace": from ray.data.namespace_expressions.list_namespace import _ListNamespace return _ListNamespace elif name == "_StringNamespace": from ray.data.namespace_expressions.string_namespace import _StringNamespace return _StringNamespace elif name == "_StructNamespace": from ray.data.namespace_expressions.struct_namespace import _StructNamespace return _StructNamespace elif name == "_DatetimeNamespace": from ray.data.namespace_expressions.dt_namespace import _DatetimeNamespace return _DatetimeNamespace raise AttributeError(f"module {__name__!r} has no attribute {name!r}")