chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,940 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import logging
|
||||
import operator
|
||||
from typing import Any, Callable, Dict, List, Optional, TypeVar, Union
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pyarrow.compute as pc
|
||||
import pyarrow.dataset as ds
|
||||
|
||||
from ray.data._internal.execution.interfaces.task_context import TaskContext
|
||||
from ray.data._internal.logical.rules.projection_pushdown import (
|
||||
_extract_input_columns_renaming_mapping,
|
||||
)
|
||||
from ray.data.block import Block, BlockAccessor, BlockColumn, BlockType
|
||||
from ray.data.expressions import (
|
||||
AliasExpr,
|
||||
BinaryExpr,
|
||||
ColumnExpr,
|
||||
DownloadExpr,
|
||||
Expr,
|
||||
LiteralExpr,
|
||||
MonotonicallyIncreasingIdExpr,
|
||||
Operation,
|
||||
RandomExpr,
|
||||
StarExpr,
|
||||
UDFExpr,
|
||||
UnaryExpr,
|
||||
UUIDExpr,
|
||||
_ExprVisitor,
|
||||
col,
|
||||
is_rename_expr,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _pa_is_in(left: Any, right: Any) -> Any:
|
||||
if not isinstance(right, (pa.Array, pa.ChunkedArray)):
|
||||
right = pa.array(right.as_py() if isinstance(right, pa.Scalar) else right)
|
||||
return pc.is_in(left, right)
|
||||
|
||||
|
||||
_PANDAS_EXPR_OPS_MAP: Dict[Operation, Callable[..., Any]] = {
|
||||
Operation.ADD: operator.add,
|
||||
Operation.SUB: operator.sub,
|
||||
Operation.MUL: operator.mul,
|
||||
Operation.DIV: operator.truediv,
|
||||
Operation.MOD: operator.mod,
|
||||
Operation.FLOORDIV: operator.floordiv,
|
||||
Operation.GT: operator.gt,
|
||||
Operation.LT: operator.lt,
|
||||
Operation.GE: operator.ge,
|
||||
Operation.LE: operator.le,
|
||||
Operation.EQ: operator.eq,
|
||||
Operation.NE: operator.ne,
|
||||
Operation.AND: operator.and_,
|
||||
Operation.OR: operator.or_,
|
||||
Operation.NOT: operator.invert,
|
||||
Operation.IS_NULL: pd.isna,
|
||||
Operation.IS_NOT_NULL: pd.notna,
|
||||
Operation.IN: lambda left, right: left.isin(right),
|
||||
Operation.NOT_IN: lambda left, right: ~left.isin(right),
|
||||
}
|
||||
|
||||
|
||||
def _is_pa_string_type(t: pa.DataType) -> bool:
|
||||
return pa.types.is_string(t) or pa.types.is_large_string(t)
|
||||
|
||||
|
||||
def _is_pa_string_like(x: Union[pa.Array, pa.ChunkedArray]) -> bool:
|
||||
t = x.type
|
||||
if pa.types.is_dictionary(t):
|
||||
t = t.value_type
|
||||
return _is_pa_string_type(t)
|
||||
|
||||
|
||||
def _pa_decode_dict_string_array(x: Union[pa.Array, pa.ChunkedArray]) -> Any:
|
||||
"""Convert Arrow dictionary-encoded string arrays to regular string arrays.
|
||||
|
||||
Dictionary encoding stores strings as indices into a dictionary of unique values.
|
||||
This function converts them back to regular string arrays for string operations.
|
||||
|
||||
Example:
|
||||
# Input: pa.array(['a', 'b']).dictionary_encode()
|
||||
# -- dictionary: ["a", "b"]
|
||||
# -- indices: [0, 1]
|
||||
# Output: regular string array ["a", "b"]
|
||||
Args:
|
||||
x: The input array to convert.
|
||||
Returns:
|
||||
The converted string array.
|
||||
"""
|
||||
if pa.types.is_dictionary(x.type) and _is_pa_string_type(x.type.value_type):
|
||||
return pc.cast(x, pa.string())
|
||||
return x
|
||||
|
||||
|
||||
def _to_pa_string_input(x: Any) -> Any:
|
||||
if isinstance(x, str):
|
||||
return pa.scalar(x)
|
||||
if isinstance(x, (pa.Array, pa.ChunkedArray)) and _is_pa_string_like(x):
|
||||
return _pa_decode_dict_string_array(x)
|
||||
actual_type = (
|
||||
str(x.type) if isinstance(x, (pa.Array, pa.ChunkedArray)) else type(x).__name__
|
||||
)
|
||||
raise TypeError(
|
||||
"Expected string or string-like pyarrow Array/ChunkedArray for string "
|
||||
f"concatenation, got {actual_type}."
|
||||
)
|
||||
|
||||
|
||||
def _pa_add_or_concat(left: Any, right: Any) -> Any:
|
||||
if isinstance(left, pa.Scalar):
|
||||
left = left.as_py()
|
||||
if isinstance(right, pa.Scalar):
|
||||
right = right.as_py()
|
||||
# If either side is string-like, perform string concatenation.
|
||||
if (
|
||||
isinstance(left, str)
|
||||
or isinstance(right, str)
|
||||
or (isinstance(left, (pa.Array, pa.ChunkedArray)) and _is_pa_string_like(left))
|
||||
or (
|
||||
isinstance(right, (pa.Array, pa.ChunkedArray)) and _is_pa_string_like(right)
|
||||
)
|
||||
):
|
||||
left_input = _to_pa_string_input(left)
|
||||
right_input = _to_pa_string_input(right)
|
||||
return pc.binary_join_element_wise(left_input, right_input, "")
|
||||
return pc.add(left, right)
|
||||
|
||||
|
||||
_ARROW_EXPR_OPS_MAP: Dict[Operation, Callable[..., Any]] = {
|
||||
Operation.ADD: _pa_add_or_concat,
|
||||
Operation.SUB: pc.subtract,
|
||||
Operation.MUL: pc.multiply,
|
||||
Operation.DIV: pc.divide,
|
||||
Operation.MOD: lambda left, right: (
|
||||
# Modulo op is essentially:
|
||||
# r = N - floor(N/M) * M
|
||||
pc.subtract(left, pc.multiply(pc.floor(pc.divide(left, right)), right))
|
||||
),
|
||||
Operation.FLOORDIV: lambda left, right: pc.floor(pc.divide(left, right)),
|
||||
Operation.GT: pc.greater,
|
||||
Operation.LT: pc.less,
|
||||
Operation.GE: pc.greater_equal,
|
||||
Operation.LE: pc.less_equal,
|
||||
Operation.EQ: pc.equal,
|
||||
Operation.NE: pc.not_equal,
|
||||
Operation.AND: pc.and_kleene,
|
||||
Operation.OR: pc.or_kleene,
|
||||
Operation.NOT: pc.invert,
|
||||
Operation.IS_NULL: pc.is_null,
|
||||
Operation.IS_NOT_NULL: pc.is_valid,
|
||||
Operation.IN: _pa_is_in,
|
||||
Operation.NOT_IN: lambda left, right: pc.invert(_pa_is_in(left, right)),
|
||||
}
|
||||
|
||||
|
||||
# NOTE: (srinathk) There are 3 distinct stages of handling passed in exprs:
|
||||
# 1. Parsing it (as text)
|
||||
# 2. Resolving unbound names (to schema)
|
||||
# 3. Converting resolved expressions to PA ones
|
||||
# Need to break up the abstraction provided by ExpressionEvaluator.
|
||||
|
||||
ScalarType = TypeVar("ScalarType")
|
||||
|
||||
|
||||
class ExpressionEvaluator:
|
||||
@staticmethod
|
||||
def get_filters(expression: str) -> ds.Expression:
|
||||
"""Parse and evaluate the expression to generate a filter condition.
|
||||
|
||||
Args:
|
||||
expression: A string representing the filter expression to parse.
|
||||
|
||||
Returns:
|
||||
A PyArrow compute expression for filtering data.
|
||||
|
||||
"""
|
||||
try:
|
||||
tree = ast.parse(expression, mode="eval")
|
||||
return _ConvertToArrowExpressionVisitor().visit(tree.body)
|
||||
except SyntaxError as e:
|
||||
raise ValueError(f"Invalid syntax in the expression: {expression}") from e
|
||||
except Exception as e:
|
||||
logger.exception(f"Error processing expression: {e}")
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
def parse_native_expression(expression: str) -> "Expr":
|
||||
"""Parse and evaluate the expression to generate a Ray Data expression.
|
||||
|
||||
Args:
|
||||
expression: A string representing the filter expression to parse.
|
||||
|
||||
Returns:
|
||||
A Ray Data Expr object for filtering data.
|
||||
|
||||
"""
|
||||
try:
|
||||
tree = ast.parse(expression, mode="eval")
|
||||
return _ConvertToNativeExpressionVisitor().visit(tree.body)
|
||||
except SyntaxError as e:
|
||||
raise ValueError(f"Invalid syntax in the expression: {expression}") from e
|
||||
except Exception as e:
|
||||
logger.exception(f"Error processing expression: {e}")
|
||||
raise
|
||||
|
||||
|
||||
class _ConvertToArrowExpressionVisitor(ast.NodeVisitor):
|
||||
# TODO: Deprecate this visitor after we remove string support in filter API.
|
||||
def visit_Compare(self, node: ast.Compare) -> ds.Expression:
|
||||
"""Handle comparison operations (e.g., a == b, a < b, a in b).
|
||||
|
||||
Args:
|
||||
node: The AST node representing a comparison operation.
|
||||
|
||||
Returns:
|
||||
An expression representing the comparison.
|
||||
"""
|
||||
# Handle left operand
|
||||
# TODO Validate columns
|
||||
if isinstance(node.left, ast.Attribute):
|
||||
# Visit and handle attributes
|
||||
left_expr = self.visit(node.left)
|
||||
elif isinstance(node.left, ast.Name):
|
||||
# Treat as a simple field
|
||||
left_expr = self.visit(node.left)
|
||||
elif isinstance(node.left, ast.Constant):
|
||||
# Constant values are used directly
|
||||
left_expr = node.left.value
|
||||
else:
|
||||
raise ValueError(f"Unsupported left operand type: {type(node.left)}")
|
||||
|
||||
comparators = [self.visit(comp) for comp in node.comparators]
|
||||
|
||||
op = node.ops[0]
|
||||
if isinstance(op, ast.In):
|
||||
return pc.is_in(left_expr, comparators[0])
|
||||
elif isinstance(op, ast.NotIn):
|
||||
return ~pc.is_in(left_expr, comparators[0])
|
||||
elif isinstance(op, ast.Eq):
|
||||
return left_expr == comparators[0]
|
||||
elif isinstance(op, ast.NotEq):
|
||||
return left_expr != comparators[0]
|
||||
elif isinstance(op, ast.Lt):
|
||||
return left_expr < comparators[0]
|
||||
elif isinstance(op, ast.LtE):
|
||||
return left_expr <= comparators[0]
|
||||
elif isinstance(op, ast.Gt):
|
||||
return left_expr > comparators[0]
|
||||
elif isinstance(op, ast.GtE):
|
||||
return left_expr >= comparators[0]
|
||||
else:
|
||||
raise ValueError(f"Unsupported operator type: {op}")
|
||||
|
||||
def visit_BoolOp(self, node: ast.BoolOp) -> ds.Expression:
|
||||
"""Handle logical operations (e.g., a and b, a or b).
|
||||
|
||||
Args:
|
||||
node: The AST node representing a boolean operation.
|
||||
|
||||
Returns:
|
||||
An expression representing the logical operation.
|
||||
"""
|
||||
conditions = [self.visit(value) for value in node.values]
|
||||
combined_expr = conditions[0]
|
||||
|
||||
for condition in conditions[1:]:
|
||||
if isinstance(node.op, ast.And):
|
||||
# Combine conditions with logical AND
|
||||
combined_expr &= condition
|
||||
elif isinstance(node.op, ast.Or):
|
||||
# Combine conditions with logical OR
|
||||
combined_expr |= condition
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported logical operator: {type(node.op).__name__}"
|
||||
)
|
||||
|
||||
return combined_expr
|
||||
|
||||
def visit_Name(self, node: ast.Name) -> ds.Expression:
|
||||
"""Handle variable (name) nodes and return them as pa.dataset.Expression.
|
||||
|
||||
Even if the name contains periods, it's treated as a single string.
|
||||
|
||||
Args:
|
||||
node: The AST node representing a variable.
|
||||
|
||||
Returns:
|
||||
The variable wrapped as a pa.dataset.Expression.
|
||||
"""
|
||||
# Directly use the field name as a string (even if it contains periods)
|
||||
field_name = node.id
|
||||
return pc.field(field_name)
|
||||
|
||||
def visit_Attribute(self, node: ast.Attribute) -> object:
|
||||
"""Handle attribute access (e.g., np.nan).
|
||||
|
||||
Args:
|
||||
node: The AST node representing an attribute access.
|
||||
|
||||
Returns:
|
||||
object: The attribute value.
|
||||
|
||||
Raises:
|
||||
ValueError: If the attribute is unsupported.
|
||||
"""
|
||||
# Recursively visit the left side (base object or previous attribute)
|
||||
if isinstance(node.value, ast.Attribute):
|
||||
# If the value is an attribute, recursively resolve it
|
||||
left_expr = self.visit(node.value)
|
||||
return pc.field(f"{left_expr}.{node.attr}")
|
||||
|
||||
elif isinstance(node.value, ast.Name):
|
||||
# If the value is a name (e.g., "foo"), we can directly return the field
|
||||
left_name = node.value.id # The base name, e.g., "foo"
|
||||
return pc.field(f"{left_name}.{node.attr}")
|
||||
|
||||
raise ValueError(f"Unsupported attribute: {node.attr}")
|
||||
|
||||
def visit_List(self, node: ast.List) -> ds.Expression:
|
||||
"""Handle list literals.
|
||||
|
||||
Args:
|
||||
node: The AST node representing a list.
|
||||
|
||||
Returns:
|
||||
The list of elements wrapped as a pa.dataset.Expression.
|
||||
"""
|
||||
elements = [self.visit(elt) for elt in node.elts]
|
||||
return pa.array(elements)
|
||||
|
||||
def visit_UnaryOp(self, node: ast.UnaryOp) -> ds.Expression:
|
||||
"""Handle case where comparator is UnaryOP (e.g., a == -1).
|
||||
|
||||
AST for this expression will be Compare(left=Name(id='a'), ops=[Eq()],
|
||||
comparators=[UnaryOp(op=USub(), operand=Constant(value=1))])
|
||||
|
||||
Args:
|
||||
node: The constant value.
|
||||
|
||||
Returns:
|
||||
A PyArrow scalar expression representing the unary operation result.
|
||||
"""
|
||||
op = node.op
|
||||
if isinstance(op, ast.USub):
|
||||
return pc.scalar(-node.operand.value)
|
||||
else:
|
||||
raise ValueError(f"Unsupported unary operator: {op}")
|
||||
|
||||
# TODO (srinathk) Note that visit_Constant does not return pa.dataset.Expression
|
||||
# because to support function in() which takes in a List, the elements in the List
|
||||
# needs to values instead of pa.dataset.Expression per pyarrow.dataset.Expression
|
||||
# specification. May be down the road, we can update it as Arrow relaxes this
|
||||
# constraint.
|
||||
def visit_Constant(self, node: ast.Constant) -> object:
|
||||
"""Handle constant values (e.g., numbers, strings).
|
||||
|
||||
Args:
|
||||
node: The AST node representing a constant value.
|
||||
|
||||
Returns:
|
||||
object: The constant value itself (e.g., number, string, or boolean).
|
||||
"""
|
||||
return node.value # Return the constant value directly.
|
||||
|
||||
def visit_Call(self, node: ast.Call) -> ds.Expression:
|
||||
"""Handle function calls (e.g., is_nan(a), is_valid(b)).
|
||||
|
||||
Args:
|
||||
node: The AST node representing a function call.
|
||||
|
||||
Returns:
|
||||
The corresponding expression based on the function called.
|
||||
|
||||
Raises:
|
||||
ValueError: If the function is unsupported or has incorrect arguments.
|
||||
"""
|
||||
func_name = node.func.id
|
||||
function_map = {
|
||||
"is_nan": lambda arg: arg.is_nan(),
|
||||
"is_null": lambda arg, nan_is_null=False: arg.is_null(
|
||||
nan_is_null=nan_is_null
|
||||
),
|
||||
"is_valid": lambda arg: arg.is_valid(),
|
||||
"is_in": lambda arg1, arg2: pc.is_in(arg1, arg2),
|
||||
}
|
||||
|
||||
if func_name in function_map:
|
||||
# Visit all arguments of the function call
|
||||
args = [self.visit(arg) for arg in node.args]
|
||||
# Handle the "is_null" function with one or two arguments
|
||||
if func_name == "is_null":
|
||||
if len(args) == 1:
|
||||
return function_map[func_name](args[0])
|
||||
elif len(args) == 2:
|
||||
return function_map[func_name](args[0], args[1])
|
||||
else:
|
||||
raise ValueError("is_null function requires one or two arguments.")
|
||||
# Handle the "is_in" function with exactly two arguments
|
||||
elif func_name == "is_in" and len(args) != 2:
|
||||
raise ValueError("is_in function requires two arguments.")
|
||||
# Ensure the function has one argument (for functions like is_valid)
|
||||
elif func_name != "is_in" and len(args) != 1:
|
||||
raise ValueError(f"{func_name} function requires exactly one argument.")
|
||||
# Call the corresponding function with the arguments
|
||||
return function_map[func_name](*args)
|
||||
else:
|
||||
raise ValueError(f"Unsupported function: {func_name}")
|
||||
|
||||
|
||||
class _ConvertToNativeExpressionVisitor(ast.NodeVisitor):
|
||||
"""AST visitor that converts string expressions to Ray Data expressions."""
|
||||
|
||||
def visit_Compare(self, node: ast.Compare) -> "Expr":
|
||||
"""Handle comparison operations (e.g., a == b, a < b, a in b)."""
|
||||
from ray.data.expressions import BinaryExpr, Operation
|
||||
|
||||
if len(node.ops) != 1 or len(node.comparators) != 1:
|
||||
raise ValueError("Only simple binary comparisons are supported")
|
||||
|
||||
left = self.visit(node.left)
|
||||
right = self.visit(node.comparators[0])
|
||||
op = node.ops[0]
|
||||
|
||||
# Map AST comparison operators to Ray Data operations
|
||||
op_map = {
|
||||
ast.Eq: Operation.EQ,
|
||||
ast.NotEq: Operation.NE,
|
||||
ast.Lt: Operation.LT,
|
||||
ast.LtE: Operation.LE,
|
||||
ast.Gt: Operation.GT,
|
||||
ast.GtE: Operation.GE,
|
||||
ast.In: Operation.IN,
|
||||
ast.NotIn: Operation.NOT_IN,
|
||||
}
|
||||
|
||||
if type(op) not in op_map:
|
||||
raise ValueError(f"Unsupported comparison operator: {type(op).__name__}")
|
||||
|
||||
return BinaryExpr(op_map[type(op)], left, right)
|
||||
|
||||
def visit_BoolOp(self, node: ast.BoolOp) -> "Expr":
|
||||
"""Handle logical operations (e.g., a and b, a or b)."""
|
||||
from ray.data.expressions import BinaryExpr, Operation
|
||||
|
||||
conditions = [self.visit(value) for value in node.values]
|
||||
combined_expr = conditions[0]
|
||||
|
||||
for condition in conditions[1:]:
|
||||
if isinstance(node.op, ast.And):
|
||||
combined_expr = BinaryExpr(Operation.AND, combined_expr, condition)
|
||||
elif isinstance(node.op, ast.Or):
|
||||
combined_expr = BinaryExpr(Operation.OR, combined_expr, condition)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported logical operator: {type(node.op).__name__}"
|
||||
)
|
||||
|
||||
return combined_expr
|
||||
|
||||
def visit_UnaryOp(self, node: ast.UnaryOp) -> "Expr":
|
||||
"""Handle unary operations (e.g., not a, -5)."""
|
||||
from ray.data.expressions import Operation, UnaryExpr, lit
|
||||
|
||||
if isinstance(node.op, ast.Not):
|
||||
operand = self.visit(node.operand)
|
||||
return UnaryExpr(Operation.NOT, operand)
|
||||
elif isinstance(node.op, ast.USub):
|
||||
operand = self.visit(node.operand)
|
||||
return operand * lit(-1)
|
||||
else:
|
||||
raise ValueError(f"Unsupported unary operator: {type(node.op).__name__}")
|
||||
|
||||
def visit_Name(self, node: ast.Name) -> "Expr":
|
||||
"""Handle variable names (column references)."""
|
||||
from ray.data.expressions import col
|
||||
|
||||
return col(node.id)
|
||||
|
||||
def visit_Constant(self, node: ast.Constant) -> "Expr":
|
||||
"""Handle constant values (numbers, strings, booleans)."""
|
||||
from ray.data.expressions import lit
|
||||
|
||||
return lit(node.value)
|
||||
|
||||
def visit_List(self, node: ast.List) -> "Expr":
|
||||
"""Handle list literals."""
|
||||
from ray.data.expressions import LiteralExpr, lit
|
||||
|
||||
# Visit all elements first
|
||||
visited_elements = [self.visit(elt) for elt in node.elts]
|
||||
|
||||
# Try to extract constant values for literal list
|
||||
elements = []
|
||||
for elem in visited_elements:
|
||||
if isinstance(elem, LiteralExpr):
|
||||
elements.append(elem.value)
|
||||
else:
|
||||
# For compatibility with Arrow visitor, we need to support non-literals
|
||||
# but Ray Data expressions may have limitations here
|
||||
raise ValueError(
|
||||
"List contains non-constant expressions. Ray Data expressions "
|
||||
"currently only support lists of constant values."
|
||||
)
|
||||
|
||||
return lit(elements)
|
||||
|
||||
def visit_Attribute(self, node: ast.Attribute) -> "Expr":
|
||||
"""Handle attribute access (e.g., for nested column names)."""
|
||||
from ray.data.expressions import col
|
||||
|
||||
# For nested column names like "user.age", combine them with dots
|
||||
if isinstance(node.value, ast.Name):
|
||||
return col(f"{node.value.id}.{node.attr}")
|
||||
elif isinstance(node.value, ast.Attribute):
|
||||
# Recursively handle nested attributes
|
||||
left_expr = self.visit(node.value)
|
||||
if isinstance(left_expr, ColumnExpr):
|
||||
return col(f"{left_expr._name}.{node.attr}")
|
||||
|
||||
raise ValueError(
|
||||
f"Unsupported attribute access: {node.attr}. Node details: {ast.dump(node)}"
|
||||
)
|
||||
|
||||
def visit_Call(self, node: ast.Call) -> "Expr":
|
||||
"""Handle function calls for operations like is_null, is_not_null, is_nan, random."""
|
||||
from ray.data.expressions import (
|
||||
BinaryExpr,
|
||||
Operation,
|
||||
UnaryExpr,
|
||||
)
|
||||
|
||||
func_name = node.func.id if isinstance(node.func, ast.Name) else str(node.func)
|
||||
|
||||
if func_name == "is_null":
|
||||
if len(node.args) != 1:
|
||||
raise ValueError("is_null() expects exactly one argument")
|
||||
operand = self.visit(node.args[0])
|
||||
return UnaryExpr(Operation.IS_NULL, operand)
|
||||
# Adding this conditional to keep it consistent with the current implementation,
|
||||
# of carrying Pyarrow's semantic of `is_valid`
|
||||
elif func_name == "is_valid" or func_name == "is_not_null":
|
||||
if len(node.args) != 1:
|
||||
raise ValueError(f"{func_name}() expects exactly one argument")
|
||||
operand = self.visit(node.args[0])
|
||||
return UnaryExpr(Operation.IS_NOT_NULL, operand)
|
||||
elif func_name == "is_nan":
|
||||
if len(node.args) != 1:
|
||||
raise ValueError("is_nan() expects exactly one argument")
|
||||
operand = self.visit(node.args[0])
|
||||
# Use x != x pattern for NaN detection (NaN != NaN is True)
|
||||
return BinaryExpr(Operation.NE, operand, operand)
|
||||
elif func_name == "is_in":
|
||||
if len(node.args) != 2:
|
||||
raise ValueError("is_in() expects exactly two arguments")
|
||||
left = self.visit(node.args[0])
|
||||
right = self.visit(node.args[1])
|
||||
return BinaryExpr(Operation.IN, left, right)
|
||||
elif func_name == "random":
|
||||
raise ValueError(
|
||||
"random() is not supported in string expressions. "
|
||||
"String expressions are deprecated. Please use the expression API instead: "
|
||||
"from ray.data.expressions import random; ds.filter(expr=(random(seed=42)>0.5))"
|
||||
)
|
||||
elif func_name == "uuid":
|
||||
raise ValueError(
|
||||
"uuid() is not supported in string expressions. "
|
||||
"String expressions are deprecated. Please use the expression API instead: "
|
||||
"ds.filter(expr=uuid().str.starts_with('a'))"
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported function: {func_name}")
|
||||
|
||||
|
||||
class NativeExpressionEvaluator(_ExprVisitor[Union[BlockColumn, ScalarType]]):
|
||||
"""Visitor-based expression evaluator that uses Block and BlockColumns
|
||||
|
||||
This evaluator implements the visitor pattern to traverse expression trees
|
||||
and evaluate them against Block data structures. It maintains operation
|
||||
mappings in shared state and returns consistent BlockColumn types.
|
||||
"""
|
||||
|
||||
def __init__(self, block: Block):
|
||||
"""Initialize the evaluator with a block and operation mappings.
|
||||
|
||||
Args:
|
||||
block: The Block to evaluate expressions against.
|
||||
"""
|
||||
self.block = block
|
||||
self.block_accessor = BlockAccessor.for_block(block)
|
||||
|
||||
# Use BlockAccessor to determine operation mappings
|
||||
block_type = self.block_accessor.block_type()
|
||||
if block_type == BlockType.PANDAS:
|
||||
self.ops = _PANDAS_EXPR_OPS_MAP
|
||||
elif block_type == BlockType.ARROW:
|
||||
self.ops = _ARROW_EXPR_OPS_MAP
|
||||
else:
|
||||
raise TypeError(f"Unsupported block type: {block_type}")
|
||||
|
||||
def visit_column(self, expr: ColumnExpr) -> Union[BlockColumn, ScalarType]:
|
||||
"""Visit a column expression and return the column data.
|
||||
|
||||
Args:
|
||||
expr: The column expression.
|
||||
|
||||
Returns:
|
||||
The column data as a BlockColumn.
|
||||
"""
|
||||
return self.block[expr.name]
|
||||
|
||||
def visit_literal(self, expr: LiteralExpr) -> Union[BlockColumn, ScalarType]:
|
||||
"""Visit a literal expression and return the literal value.
|
||||
|
||||
Args:
|
||||
expr: The literal expression.
|
||||
|
||||
Returns:
|
||||
The literal value.
|
||||
"""
|
||||
# Given that expressions support pandas blocks, we need to return the value as is.
|
||||
# Pandas has multiple dtype_backends, so there's no guarantee on the return type.
|
||||
return expr.value
|
||||
|
||||
def visit_binary(self, expr: BinaryExpr) -> Union[BlockColumn, ScalarType]:
|
||||
"""Visit a binary expression and return the result of the operation.
|
||||
|
||||
Args:
|
||||
expr: The binary expression.
|
||||
|
||||
Returns:
|
||||
The result of the binary operation as a BlockColumn.
|
||||
"""
|
||||
left_result = self.visit(expr.left)
|
||||
right_result = self.visit(expr.right)
|
||||
|
||||
return self.ops[expr.op](left_result, right_result)
|
||||
|
||||
def visit_unary(self, expr: UnaryExpr) -> Union[BlockColumn, ScalarType]:
|
||||
"""Visit a unary expression and return the result of the operation.
|
||||
|
||||
Args:
|
||||
expr: The unary expression.
|
||||
|
||||
Returns:
|
||||
The result of the unary operation as a BlockColumn.
|
||||
"""
|
||||
operand_result = self.visit(expr.operand)
|
||||
return self.ops[expr.op](operand_result)
|
||||
|
||||
def visit_udf(self, expr: UDFExpr) -> Union[BlockColumn, ScalarType]:
|
||||
"""Visit a UDF expression and return the result of the function call.
|
||||
|
||||
Args:
|
||||
expr: The UDF expression.
|
||||
|
||||
Returns:
|
||||
The result of the UDF call as a BlockColumn.
|
||||
"""
|
||||
args = [self.visit(arg) for arg in expr.args]
|
||||
kwargs = {k: self.visit(v) for k, v in expr.kwargs.items()}
|
||||
|
||||
result = expr.fn(*args, **kwargs)
|
||||
|
||||
if not isinstance(result, (pd.Series, np.ndarray, pa.Array, pa.ChunkedArray)):
|
||||
function_name = expr.fn.__name__
|
||||
raise TypeError(
|
||||
f"UDF '{function_name}' returned invalid type {type(result).__name__}. "
|
||||
f"Expected type (pandas.Series, numpy.ndarray, pyarrow.Array, "
|
||||
f"pyarrow.ChunkedArray)"
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
def visit_alias(self, expr: AliasExpr) -> Union[BlockColumn, ScalarType]:
|
||||
"""Visit an alias expression and return the renamed result.
|
||||
|
||||
Args:
|
||||
expr: The alias expression.
|
||||
|
||||
Returns:
|
||||
A Block with the data from the inner expression.
|
||||
"""
|
||||
# Evaluate the inner expression
|
||||
return self.visit(expr.expr)
|
||||
|
||||
def visit_star(self, expr: StarExpr) -> Union[BlockColumn, ScalarType]:
|
||||
"""Visit a star expression.
|
||||
|
||||
Args:
|
||||
expr: The star expression.
|
||||
|
||||
Returns:
|
||||
TypeError: StarExpr cannot be evaluated as a regular expression.
|
||||
"""
|
||||
# star() should not be evaluated directly - it's handled at Project level
|
||||
raise TypeError(
|
||||
"StarExpr cannot be evaluated as a regular expression. "
|
||||
"It should only be used in Project operations."
|
||||
)
|
||||
|
||||
def visit_download(self, expr: DownloadExpr) -> Union[BlockColumn, ScalarType]:
|
||||
"""Visit a download expression.
|
||||
|
||||
Args:
|
||||
expr: The download expression.
|
||||
|
||||
Returns:
|
||||
TypeError: DownloadExpr evaluation not yet implemented.
|
||||
"""
|
||||
raise TypeError(
|
||||
"DownloadExpr evaluation is not yet implemented in NativeExpressionEvaluator."
|
||||
)
|
||||
|
||||
def visit_monotonically_increasing_id(
|
||||
self, expr: MonotonicallyIncreasingIdExpr
|
||||
) -> Union[BlockColumn, ScalarType]:
|
||||
"""Visit a monotonically_increasing_id expression.
|
||||
|
||||
Args:
|
||||
expr: The monotonically_increasing_id expression.
|
||||
|
||||
Returns:
|
||||
The result of the monotonically_increasing_id expression as a BlockColumn.
|
||||
"""
|
||||
ctx = TaskContext.get_current()
|
||||
assert (
|
||||
ctx is not None
|
||||
), "TaskContext is required for monotonically_increasing_id()"
|
||||
|
||||
# Key the counter by expression instance ID so that multiple expressions
|
||||
# in the same projection will have isolated row count state.
|
||||
# This is required because a single task may process multiple blocks if
|
||||
# the upstream data source does not compress the data into a single block.
|
||||
counter_key = f"_mono_id_{expr._instance_id}_counter"
|
||||
|
||||
start_idx = ctx.kwargs.get(counter_key, 0)
|
||||
num_rows = self.block_accessor.num_rows()
|
||||
end_idx = start_idx + num_rows
|
||||
ctx.kwargs[counter_key] = end_idx
|
||||
|
||||
# int64 (signed): upper 30 bits = task ID, lower 33 bits = row number.
|
||||
# Note end_idx is an exclusive upper bound, as the max row ID is end_idx - 1.
|
||||
ROW_BITS = 33
|
||||
TASK_BITS = 30
|
||||
if end_idx > (1 << ROW_BITS):
|
||||
raise ValueError(
|
||||
f"Cannot generate monotonically increasing IDs: row count for this task exceeds the maximum allowed value of {(1 << ROW_BITS) - 1}"
|
||||
)
|
||||
if ctx.task_idx >= (1 << TASK_BITS):
|
||||
raise ValueError(
|
||||
f"Cannot generate monotonically increasing IDs: number of tasks exceeds the maximum allowed value of {(1 << TASK_BITS) - 1}"
|
||||
)
|
||||
|
||||
partition_mask = ctx.task_idx << ROW_BITS
|
||||
ids = partition_mask + np.arange(start_idx, end_idx, dtype=np.int64)
|
||||
|
||||
block_type = self.block_accessor.block_type()
|
||||
if block_type == BlockType.PANDAS:
|
||||
return pd.Series(ids)
|
||||
elif block_type == BlockType.ARROW:
|
||||
return pa.array(ids)
|
||||
else:
|
||||
raise TypeError(f"Unsupported block type: {block_type}")
|
||||
|
||||
def visit_random(self, expr: RandomExpr) -> Union[BlockColumn, ScalarType]:
|
||||
"""Visit a random expression and return the result of the operation.
|
||||
|
||||
Args:
|
||||
expr: The random expression.
|
||||
|
||||
Returns:
|
||||
The result of the random operation as a BlockColumn.
|
||||
"""
|
||||
from ray.data._internal.planner.plan_expression.synthetic_impl import (
|
||||
eval_random,
|
||||
)
|
||||
|
||||
return eval_random(
|
||||
self.block_accessor.num_rows(),
|
||||
self.block_accessor.block_type(),
|
||||
seed=expr.seed,
|
||||
reseed_after_execution=expr.reseed_after_execution,
|
||||
instance_id=expr._instance_id,
|
||||
)
|
||||
|
||||
def visit_uuid(self, expr: UUIDExpr) -> Union[BlockColumn, ScalarType]:
|
||||
"""Visit a uuid expression and return the result of the operation.
|
||||
|
||||
Args:
|
||||
expr: The uuid expression.
|
||||
|
||||
Returns:
|
||||
The result of the uuid operation as a BlockColumn.
|
||||
"""
|
||||
from ray.data._internal.planner.plan_expression.synthetic_impl import eval_uuid
|
||||
|
||||
return eval_uuid(
|
||||
self.block_accessor.num_rows(), self.block_accessor.block_type()
|
||||
)
|
||||
|
||||
|
||||
def eval_expr(expr: Expr, block: Block) -> Union[BlockColumn, ScalarType]:
|
||||
"""Evaluate an expression against a block using the visitor pattern.
|
||||
|
||||
Args:
|
||||
expr: The expression to evaluate.
|
||||
block: The Block to evaluate against.
|
||||
|
||||
Returns:
|
||||
The evaluated result as a BlockColumn or a scalar value.
|
||||
"""
|
||||
evaluator = NativeExpressionEvaluator(block)
|
||||
return evaluator.visit(expr)
|
||||
|
||||
|
||||
def _eval_projection_without_cse(projection_exprs: List[Expr], block: Block) -> Block:
|
||||
"""
|
||||
Evaluate a projection (list of expressions) against a block.
|
||||
|
||||
Handles projection semantics including:
|
||||
- Empty projections
|
||||
- Star() expressions for preserving existing columns
|
||||
- Rename detection
|
||||
- Column ordering
|
||||
|
||||
Args:
|
||||
projection_exprs: List of expressions to evaluate (may include StarExpr)
|
||||
block: The block to project
|
||||
|
||||
Returns:
|
||||
A new block with the projected schema
|
||||
"""
|
||||
block_accessor = BlockAccessor.for_block(block)
|
||||
|
||||
# Skip projection only for schema-less empty blocks.
|
||||
if block_accessor.num_rows() == 0 and len(block_accessor.column_names()) == 0:
|
||||
return block
|
||||
|
||||
# Handle simple cases early.
|
||||
if len(projection_exprs) == 0:
|
||||
return block_accessor.select([])
|
||||
|
||||
input_column_names = list(block_accessor.column_names())
|
||||
# Collect input column rename map from the projection list
|
||||
input_column_rename_map = _extract_input_columns_renaming_mapping(projection_exprs)
|
||||
|
||||
# Expand star expr (if any). ``Project.__post_init__`` eagerly expands
|
||||
# ``StarExpr`` to explicit ``col()`` refs whenever the
|
||||
# input schema is known, so this runtime branch is hit only on the
|
||||
# UDF-fallback path (Project on top of an opaque-schema input).
|
||||
if isinstance(projection_exprs[0], StarExpr):
|
||||
# Bucket the trailing exprs: rename ``AliasExpr``s of an input
|
||||
# column get placed into the original column's position (so the
|
||||
# output preserves on-disk column order); anything else (e.g.
|
||||
# ``with_column`` computed expressions) is appended afterwards.
|
||||
rename_exprs_by_source: Dict[str, Expr] = {}
|
||||
extra_exprs: List[Expr] = []
|
||||
for expr in projection_exprs[1:]:
|
||||
# e.g. ``col(source)._rename(new_name)`` — bucket by ``source`` for column order.
|
||||
# ``rename_exprs_by_source``: input column name -> that rename ``AliasExpr``.
|
||||
if is_rename_expr(expr) and expr.expr.name in input_column_rename_map:
|
||||
rename_exprs_by_source[expr.expr.name] = expr
|
||||
else:
|
||||
extra_exprs.append(expr)
|
||||
|
||||
ordered_exprs: List[Expr] = []
|
||||
for c in input_column_names:
|
||||
if c in rename_exprs_by_source:
|
||||
ordered_exprs.append(rename_exprs_by_source.pop(c))
|
||||
elif c not in input_column_rename_map:
|
||||
ordered_exprs.append(col(c))
|
||||
|
||||
# Any rename whose source column isn't in the block falls through to
|
||||
# ``extra_exprs`` so evaluation raises a "column not found" error
|
||||
# instead of silently dropping the expression.
|
||||
extra_exprs = list(rename_exprs_by_source.values()) + extra_exprs
|
||||
|
||||
projection_exprs = ordered_exprs + extra_exprs
|
||||
|
||||
names, output_cols = zip(*[(e.name, eval_expr(e, block)) for e in projection_exprs])
|
||||
|
||||
# This clumsy workaround is necessary to be able to fill in Pyarrow tables
|
||||
# that has to be "seeded" from existing table with N rows, and couldn't be
|
||||
# started from a truly empty table.
|
||||
#
|
||||
# TODO fix
|
||||
new_block = BlockAccessor.for_block(block).fill_column("__stub__", None)
|
||||
new_block = BlockAccessor.for_block(new_block).drop(input_column_names)
|
||||
|
||||
for name, output_col in zip(names, output_cols):
|
||||
new_block = BlockAccessor.for_block(new_block).fill_column(name, output_col)
|
||||
|
||||
return BlockAccessor.for_block(new_block).drop(["__stub__"])
|
||||
|
||||
|
||||
def _drop_cse_temp_columns(block: Block, temp_columns: List[str]) -> Block:
|
||||
block_accessor = BlockAccessor.for_block(block)
|
||||
drop_columns = [
|
||||
name for name in temp_columns if name in block_accessor.column_names()
|
||||
]
|
||||
if not drop_columns:
|
||||
return block
|
||||
return block_accessor.drop(drop_columns)
|
||||
|
||||
|
||||
def eval_projection(
|
||||
projection_exprs: List[Expr],
|
||||
block: Block,
|
||||
*,
|
||||
common_sub_exprs: Optional[List[Expr]] = None,
|
||||
) -> Block:
|
||||
"""
|
||||
Evaluate a projection (list of expressions) against a block.
|
||||
|
||||
If CSE common expressions are provided, they are evaluated first into
|
||||
temporary columns on a working block. Visible projection expressions are
|
||||
then evaluated against that working block.
|
||||
"""
|
||||
if not common_sub_exprs:
|
||||
return _eval_projection_without_cse(projection_exprs, block)
|
||||
|
||||
working_block = block
|
||||
for common_expr in common_sub_exprs:
|
||||
assert common_expr.name is not None
|
||||
working_block = BlockAccessor.for_block(working_block).fill_column(
|
||||
common_expr.name,
|
||||
eval_expr(common_expr, working_block),
|
||||
)
|
||||
|
||||
output_block = _eval_projection_without_cse(projection_exprs, working_block)
|
||||
temp_columns = [expr.name for expr in common_sub_exprs]
|
||||
return _drop_cse_temp_columns(output_block, temp_columns)
|
||||
@@ -0,0 +1,843 @@
|
||||
from collections import Counter
|
||||
from dataclasses import dataclass, replace
|
||||
from typing import Dict, Hashable, List, TypeVar
|
||||
|
||||
from ray.data.expressions import (
|
||||
AliasExpr,
|
||||
BinaryExpr,
|
||||
ColumnExpr,
|
||||
DownloadExpr,
|
||||
Expr,
|
||||
LiteralExpr,
|
||||
MonotonicallyIncreasingIdExpr,
|
||||
Operation,
|
||||
RandomExpr,
|
||||
StarExpr,
|
||||
UDFExpr,
|
||||
UnaryExpr,
|
||||
UUIDExpr,
|
||||
_CallableClassUDF,
|
||||
_ExprVisitor,
|
||||
)
|
||||
from ray.data.util.expression_utils import (
|
||||
_alias_fingerprint_key,
|
||||
_binary_fingerprint_key,
|
||||
_column_fingerprint_key,
|
||||
_download_fingerprint_key,
|
||||
_literal_fingerprint_key,
|
||||
_monotonically_increasing_id_fingerprint_key,
|
||||
_random_fingerprint_key,
|
||||
_star_fingerprint_key,
|
||||
_udf_fingerprint_key,
|
||||
_unary_fingerprint_key,
|
||||
_uuid_fingerprint_key,
|
||||
)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
# Mapping of operations to their string symbols for inline representation
|
||||
_INLINE_OP_SYMBOLS = {
|
||||
Operation.ADD: "+",
|
||||
Operation.SUB: "-",
|
||||
Operation.MUL: "*",
|
||||
Operation.DIV: "/",
|
||||
Operation.MOD: "%",
|
||||
Operation.FLOORDIV: "//",
|
||||
Operation.GT: ">",
|
||||
Operation.LT: "<",
|
||||
Operation.GE: ">=",
|
||||
Operation.LE: "<=",
|
||||
Operation.EQ: "==",
|
||||
Operation.NE: "!=",
|
||||
Operation.AND: "&",
|
||||
Operation.OR: "|",
|
||||
Operation.IN: "in",
|
||||
Operation.NOT_IN: "not in",
|
||||
}
|
||||
|
||||
|
||||
class _ExprVisitorBase(_ExprVisitor[None]):
|
||||
"""Base visitor that provides automatic recursive traversal.
|
||||
|
||||
This class extends _ExprVisitor and provides default implementations
|
||||
for composite nodes that automatically traverse child expressions.
|
||||
"""
|
||||
|
||||
def visit_binary(self, expr: "BinaryExpr") -> None:
|
||||
"""Default implementation: recursively visit both operands."""
|
||||
super().visit(expr.left)
|
||||
super().visit(expr.right)
|
||||
|
||||
def visit_unary(self, expr: "UnaryExpr") -> None:
|
||||
"""Default implementation: recursively visit the operand."""
|
||||
super().visit(expr.operand)
|
||||
|
||||
def visit_alias(self, expr: "AliasExpr") -> None:
|
||||
"""Default implementation: recursively visit the inner expression."""
|
||||
super().visit(expr.expr)
|
||||
|
||||
def visit_udf(self, expr: "UDFExpr") -> None:
|
||||
"""Default implementation: recursively visit all arguments."""
|
||||
for arg in expr.args:
|
||||
super().visit(arg)
|
||||
for value in expr.kwargs.values():
|
||||
super().visit(value)
|
||||
|
||||
def visit_literal(self, expr: LiteralExpr) -> None:
|
||||
"""Visit a literal expression (no columns to collect)."""
|
||||
pass
|
||||
|
||||
def visit_star(self, expr: StarExpr) -> None:
|
||||
"""Visit a star expression (no columns to collect)."""
|
||||
pass
|
||||
|
||||
def visit_download(self, expr: "Expr") -> None:
|
||||
"""Visit a download expression (no columns to collect)."""
|
||||
pass
|
||||
|
||||
def visit_monotonically_increasing_id(
|
||||
self, expr: "MonotonicallyIncreasingIdExpr"
|
||||
) -> None:
|
||||
"""Visit a monotonically_increasing_id expression (no columns to collect)."""
|
||||
pass
|
||||
|
||||
def visit_random(self, expr: "RandomExpr") -> None:
|
||||
"""Visit a synthetic expression (no columns to collect)."""
|
||||
pass
|
||||
|
||||
def visit_uuid(self, expr: "UUIDExpr") -> None:
|
||||
"""Visit a uuid expression (no columns to collect)."""
|
||||
pass
|
||||
|
||||
|
||||
class _ColumnReferenceCollector(_ExprVisitorBase):
|
||||
"""Visitor that collects all column references from expression trees.
|
||||
|
||||
Backed by a ``Counter`` so callers can take either:
|
||||
- ``get_column_refs()`` -> ordered, de-duplicated column names, or
|
||||
- ``get_counts()`` -> per-name reference multiplicity, counting repeats
|
||||
*within* a single expression (``x + x`` -> ``{"x": 2}``).
|
||||
|
||||
``Counter`` preserves first-insertion order, so ``get_column_refs()`` returns the
|
||||
same ordered, de-duplicated list as a plain insertion-ordered ``dict`` would.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize with an empty reference counter."""
|
||||
self._col_refs: Counter = Counter()
|
||||
|
||||
def get_column_refs(self) -> List[str]:
|
||||
return list(self._col_refs.keys())
|
||||
|
||||
def get_counts(self) -> Counter:
|
||||
return self._col_refs
|
||||
|
||||
def visit_column(self, expr: ColumnExpr) -> None:
|
||||
"""Visit a column expression and count its name.
|
||||
|
||||
Args:
|
||||
expr: The column expression.
|
||||
|
||||
Returns:
|
||||
None (only counts columns as a side effect).
|
||||
"""
|
||||
self._col_refs[expr.name] += 1
|
||||
|
||||
def visit_alias(self, expr: AliasExpr) -> None:
|
||||
"""Visit an alias expression and collect from its inner expression.
|
||||
|
||||
Args:
|
||||
expr: The alias expression.
|
||||
|
||||
Returns:
|
||||
None (only collects columns as a side effect).
|
||||
"""
|
||||
self.visit(expr.expr)
|
||||
|
||||
|
||||
class _IdempotencyVisitor(_ExprVisitor[bool]):
|
||||
"""Reports whether an expression is safe to duplicate, reorder, or move.
|
||||
|
||||
Returns ``True`` only when every node in the tree is idempotent. The three
|
||||
non-idempotent leaf types (``RandomExpr``, ``UUIDExpr``,
|
||||
``MonotonicallyIncreasingIdExpr``) return ``False`` and propagate upward: a
|
||||
composite is idempotent iff all of its children are.
|
||||
|
||||
Optimizer rules consult this (via :func:`is_idempotent`) before any rewrite that
|
||||
would change an expression's evaluation count, row set, or position.
|
||||
"""
|
||||
|
||||
# --- non-idempotent leaves ---
|
||||
def visit_random(self, expr: RandomExpr) -> bool:
|
||||
# Conservatively non-idempotent even when seeded: CSE matches structurally and
|
||||
# ignores ``_instance_id``, while the runtime RNG counter keys on it, so a
|
||||
# seeded RandomExpr cannot be safely de-duplicated in general.
|
||||
return False
|
||||
|
||||
def visit_uuid(self, expr: UUIDExpr) -> bool:
|
||||
return False
|
||||
|
||||
def visit_monotonically_increasing_id(
|
||||
self, expr: MonotonicallyIncreasingIdExpr
|
||||
) -> bool:
|
||||
return False
|
||||
|
||||
# --- idempotent leaves ---
|
||||
def visit_column(self, expr: ColumnExpr) -> bool:
|
||||
return True
|
||||
|
||||
def visit_literal(self, expr: LiteralExpr) -> bool:
|
||||
return True
|
||||
|
||||
def visit_star(self, expr: StarExpr) -> bool:
|
||||
return True
|
||||
|
||||
def visit_download(self, expr: DownloadExpr) -> bool:
|
||||
# ``DownloadExpr`` is a leaf with no Expr children. It is idempotent (same URI
|
||||
# yields the same bytes); CSE avoids re-fetching it for *cost* reasons, which
|
||||
# is a separate concern from this correctness contract.
|
||||
return True
|
||||
|
||||
# --- composites: idempotent iff all children are ---
|
||||
#
|
||||
# Children are visited via ``child.is_idempotent()`` (not ``self.visit(child)``)
|
||||
# so each node's result is read from / written to its per-instance cache. This
|
||||
# keeps an all-nodes query (e.g. CSE visiting every occurrence) linear overall
|
||||
# instead of re-walking each subtree.
|
||||
def visit_alias(self, expr: AliasExpr) -> bool:
|
||||
return expr.expr.is_idempotent()
|
||||
|
||||
def visit_unary(self, expr: UnaryExpr) -> bool:
|
||||
return expr.operand.is_idempotent()
|
||||
|
||||
def visit_binary(self, expr: BinaryExpr) -> bool:
|
||||
return expr.left.is_idempotent() and expr.right.is_idempotent()
|
||||
|
||||
def visit_udf(self, expr: UDFExpr) -> bool:
|
||||
# FUTURE EXTENSION POINT: today UDFs are assumed idempotent and we only recurse
|
||||
# into their argument expressions. When per-UDF non-determinism is supported,
|
||||
# gate this on the UDF's declared determinism as well.
|
||||
return all(arg.is_idempotent() for arg in expr.args) and all(
|
||||
value.is_idempotent() for value in expr.kwargs.values()
|
||||
)
|
||||
|
||||
|
||||
# Stateless singleton: ``Expr.is_idempotent`` reuses this rather than allocating a
|
||||
# visitor per node during the initial (uncached) computation.
|
||||
_IDEMPOTENCY_VISITOR = _IdempotencyVisitor()
|
||||
|
||||
|
||||
class _CallableClassUDFCollector(_ExprVisitorBase):
|
||||
"""Visitor that collects all callable class UDFs from expression trees.
|
||||
|
||||
This visitor traverses expression trees and collects _CallableClassUDF instances
|
||||
that wrap callable classes (as opposed to regular functions).
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize with an empty list of _CallableClassUDF instances."""
|
||||
self._expr_udfs: List[_CallableClassUDF] = []
|
||||
|
||||
def get_callable_class_udfs(self) -> List[_CallableClassUDF]:
|
||||
"""Get the list of collected _CallableClassUDF instances.
|
||||
|
||||
Returns:
|
||||
List of _CallableClassUDF instances that wrap callable classes.
|
||||
"""
|
||||
return self._expr_udfs
|
||||
|
||||
def visit_column(self, expr: ColumnExpr) -> None:
|
||||
"""Visit a column expression (no UDFs to collect)."""
|
||||
pass
|
||||
|
||||
def visit_udf(self, expr: UDFExpr) -> None:
|
||||
"""Visit a UDF expression and collect it if it's a callable class.
|
||||
|
||||
Args:
|
||||
expr: The UDF expression.
|
||||
|
||||
Returns:
|
||||
None (only collects UDFs as a side effect).
|
||||
"""
|
||||
# Check if fn is an _CallableClassUDF (indicates callable class)
|
||||
if isinstance(expr.fn, _CallableClassUDF):
|
||||
self._expr_udfs.append(expr.fn)
|
||||
|
||||
# Continue visiting child expressions
|
||||
super().visit_udf(expr)
|
||||
|
||||
|
||||
class _ColumnSubstitutionVisitor(_ExprVisitor[Expr]):
|
||||
"""Visitor rebinding column references in ``Expression``s.
|
||||
|
||||
This visitor traverses given ``Expression`` trees and substitutes column references
|
||||
according to a provided substitution map.
|
||||
"""
|
||||
|
||||
def __init__(self, column_ref_substitutions: Dict[str, Expr]):
|
||||
"""Initialize with a column substitution map.
|
||||
|
||||
Args:
|
||||
column_ref_substitutions: Mapping from column names to replacement expressions.
|
||||
"""
|
||||
self._col_ref_substitutions = column_ref_substitutions
|
||||
|
||||
def visit_column(self, expr: ColumnExpr) -> Expr:
|
||||
"""Visit a column expression and substitute it.
|
||||
|
||||
Args:
|
||||
expr: The column expression.
|
||||
|
||||
Returns:
|
||||
The substituted expression or the original if no substitution exists.
|
||||
"""
|
||||
substitution = self._col_ref_substitutions.get(expr.name)
|
||||
|
||||
return substitution if substitution is not None else expr
|
||||
|
||||
def visit_literal(self, expr: LiteralExpr) -> Expr:
|
||||
"""Visit a literal expression (no rewriting needed).
|
||||
|
||||
Args:
|
||||
expr: The literal expression.
|
||||
|
||||
Returns:
|
||||
The original literal expression.
|
||||
"""
|
||||
return expr
|
||||
|
||||
def visit_binary(self, expr: BinaryExpr) -> Expr:
|
||||
"""Visit a binary expression and rewrite its operands.
|
||||
|
||||
Args:
|
||||
expr: The binary expression.
|
||||
|
||||
Returns:
|
||||
A new binary expression with rewritten operands.
|
||||
"""
|
||||
return BinaryExpr(
|
||||
expr.op,
|
||||
self.visit(expr.left),
|
||||
self.visit(expr.right),
|
||||
)
|
||||
|
||||
def visit_unary(self, expr: UnaryExpr) -> Expr:
|
||||
"""Visit a unary expression and rewrite its operand.
|
||||
|
||||
Args:
|
||||
expr: The unary expression.
|
||||
|
||||
Returns:
|
||||
A new unary expression with rewritten operand.
|
||||
"""
|
||||
return UnaryExpr(expr.op, self.visit(expr.operand))
|
||||
|
||||
def visit_udf(self, expr: UDFExpr) -> Expr:
|
||||
"""Visit a UDF expression and rewrite its arguments.
|
||||
|
||||
Args:
|
||||
expr: The UDF expression.
|
||||
|
||||
Returns:
|
||||
A new UDF expression with rewritten arguments.
|
||||
"""
|
||||
new_args = [self.visit(arg) for arg in expr.args]
|
||||
new_kwargs = {key: self.visit(value) for key, value in expr.kwargs.items()}
|
||||
return replace(expr, args=new_args, kwargs=new_kwargs)
|
||||
|
||||
def visit_alias(self, expr: AliasExpr) -> Expr:
|
||||
"""Visit an alias expression and rewrite its inner expression.
|
||||
|
||||
Args:
|
||||
expr: The alias expression.
|
||||
|
||||
Returns:
|
||||
A new alias expression with rewritten inner expression and preserved name.
|
||||
"""
|
||||
# We unalias returned expression to avoid nested aliasing
|
||||
visited = self.visit(expr.expr)._unalias()
|
||||
# NOTE: We're carrying over all of the other aspects of the alias
|
||||
# only replacing inner expre
|
||||
return replace(
|
||||
expr,
|
||||
expr=visited,
|
||||
# Alias expression will remain a renaming one (ie replacing source column)
|
||||
# so long as it's referencing another column (and not otherwise)
|
||||
#
|
||||
# TODO replace w/ standalone rename expr
|
||||
_is_rename=expr._is_rename and _is_col_expr(visited),
|
||||
)
|
||||
|
||||
def visit_download(self, expr: "Expr") -> Expr:
|
||||
"""Visit a download expression (no rewriting needed).
|
||||
|
||||
Args:
|
||||
expr: The download expression.
|
||||
|
||||
Returns:
|
||||
The original download expression.
|
||||
"""
|
||||
return expr
|
||||
|
||||
def visit_star(self, expr: StarExpr) -> Expr:
|
||||
"""Visit a star expression (no rewriting needed).
|
||||
|
||||
Args:
|
||||
expr: The star expression.
|
||||
|
||||
Returns:
|
||||
The original star expression.
|
||||
"""
|
||||
return expr
|
||||
|
||||
def visit_monotonically_increasing_id(
|
||||
self, expr: MonotonicallyIncreasingIdExpr
|
||||
) -> Expr:
|
||||
"""Visit a monotonically_increasing_id expression (no rewriting needed).
|
||||
|
||||
Args:
|
||||
expr: The monotonically_increasing_id expression.
|
||||
|
||||
Returns:
|
||||
The original expression.
|
||||
"""
|
||||
return expr
|
||||
|
||||
def visit_random(self, expr: "RandomExpr") -> Expr:
|
||||
"""Visit a random expression (no rewriting needed).
|
||||
|
||||
Args:
|
||||
expr: The random expression.
|
||||
|
||||
Returns:
|
||||
The original random expression.
|
||||
"""
|
||||
return expr
|
||||
|
||||
def visit_uuid(self, expr: "UUIDExpr") -> Expr:
|
||||
"""Visit a uuid expression (no rewriting needed).
|
||||
|
||||
Args:
|
||||
expr: The uuid expression.
|
||||
|
||||
Returns:
|
||||
The original uuid expression.
|
||||
"""
|
||||
return expr
|
||||
|
||||
|
||||
def _is_col_expr(expr: Expr) -> bool:
|
||||
return isinstance(expr, ColumnExpr) or (
|
||||
isinstance(expr, AliasExpr) and isinstance(expr.expr, ColumnExpr)
|
||||
)
|
||||
|
||||
|
||||
class _TreeReprVisitor(_ExprVisitor[str]):
|
||||
"""Visitor that generates a readable tree representation of expressions. Returns in pre-order traversal."""
|
||||
|
||||
def __init__(self, prefix: str = "", is_last: bool = True):
|
||||
"""
|
||||
Initialize the tree representation visitor.
|
||||
|
||||
Args:
|
||||
prefix: The prefix string for indentation (accumulated from parent nodes)
|
||||
is_last: Whether this node is the last child of its parent
|
||||
"""
|
||||
self.prefix = prefix
|
||||
self.is_last = is_last
|
||||
self._max_length = 50 # Maximum length of the node label
|
||||
|
||||
def _make_tree_lines(
|
||||
self,
|
||||
node_label: str,
|
||||
children: List[tuple[str, "Expr"]] = None,
|
||||
expr: "Expr" = None,
|
||||
) -> str:
|
||||
"""
|
||||
Format a node and its children with tree box-drawing characters.
|
||||
|
||||
Args:
|
||||
node_label: The label for this node (e.g., "ADD")
|
||||
children: List of (label, child_expr) tuples to render as children
|
||||
expr: The expression node (used to extract datatype)
|
||||
|
||||
Returns:
|
||||
Multi-line string representation of the tree
|
||||
"""
|
||||
lines = [node_label]
|
||||
|
||||
if children:
|
||||
for i, (label, child_expr) in enumerate(children):
|
||||
is_last_child = i == len(children) - 1
|
||||
|
||||
# Build prefix for the child based on whether current node is last
|
||||
child_prefix = self.prefix + (" " if self.is_last else "│ ")
|
||||
|
||||
# Choose connector: └── for last child, ├── for others
|
||||
connector = "└── " if is_last_child else "├── "
|
||||
|
||||
# Recursively visit the child with updated prefix
|
||||
child_visitor = _TreeReprVisitor(child_prefix, is_last_child)
|
||||
child_lines = child_visitor.visit(child_expr).split("\n")
|
||||
|
||||
# Add the first line with label and connector
|
||||
if label:
|
||||
lines.append(f"{child_prefix}{connector}{label}: {child_lines[0]}")
|
||||
else:
|
||||
lines.append(f"{child_prefix}{connector}{child_lines[0]}")
|
||||
|
||||
# Add remaining lines from child with proper indentation
|
||||
for line in child_lines[1:]:
|
||||
lines.append(line)
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
def visit_column(self, expr: "ColumnExpr") -> str:
|
||||
return self._make_tree_lines(f"COL({expr.name!r})", expr=expr)
|
||||
|
||||
def visit_literal(self, expr: "LiteralExpr") -> str:
|
||||
# Truncate long values for readability
|
||||
value_repr = repr(expr.value)
|
||||
if len(value_repr) > self._max_length:
|
||||
value_repr = value_repr[: self._max_length - 3] + "..."
|
||||
return self._make_tree_lines(f"LIT({value_repr})", expr=expr)
|
||||
|
||||
def visit_binary(self, expr: "BinaryExpr") -> str:
|
||||
return self._make_tree_lines(
|
||||
f"{expr.op.name}",
|
||||
children=[
|
||||
("left", expr.left),
|
||||
("right", expr.right),
|
||||
],
|
||||
expr=expr,
|
||||
)
|
||||
|
||||
def visit_unary(self, expr: "UnaryExpr") -> str:
|
||||
return self._make_tree_lines(
|
||||
f"{expr.op.name}",
|
||||
children=[("operand", expr.operand)],
|
||||
expr=expr,
|
||||
)
|
||||
|
||||
def visit_alias(self, expr: "AliasExpr") -> str:
|
||||
rename_marker = " [rename]" if expr._is_rename else ""
|
||||
return self._make_tree_lines(
|
||||
f"ALIAS({expr.name!r}){rename_marker}",
|
||||
children=[("", expr.expr)],
|
||||
expr=expr,
|
||||
)
|
||||
|
||||
def visit_udf(self, expr: "UDFExpr") -> str:
|
||||
# Get function name for better readability
|
||||
fn_name = getattr(expr.fn, "__name__", str(expr.fn))
|
||||
|
||||
children = []
|
||||
# Add positional arguments
|
||||
for i, arg in enumerate(expr.args):
|
||||
children.append((f"arg[{i}]", arg))
|
||||
|
||||
# Add keyword arguments
|
||||
for key, value in expr.kwargs.items():
|
||||
children.append((f"kwarg[{key!r}]", value))
|
||||
|
||||
return self._make_tree_lines(
|
||||
f"UDF({fn_name})",
|
||||
children=children if children else None,
|
||||
expr=expr,
|
||||
)
|
||||
|
||||
def visit_download(self, expr: "DownloadExpr") -> str:
|
||||
return self._make_tree_lines(f"DOWNLOAD({expr.uri_column_name!r})", expr=expr)
|
||||
|
||||
def visit_star(self, expr: "StarExpr") -> str:
|
||||
return self._make_tree_lines("COL(*)", expr=expr)
|
||||
|
||||
def visit_monotonically_increasing_id(
|
||||
self, expr: "MonotonicallyIncreasingIdExpr"
|
||||
) -> str:
|
||||
return self._make_tree_lines("MONOTONICALLY_INCREASING_ID()", expr=expr)
|
||||
|
||||
def visit_random(self, expr: "RandomExpr") -> str:
|
||||
if expr.seed is None:
|
||||
label = "RANDOM()"
|
||||
else:
|
||||
label = f"RANDOM(seed={expr.seed}, reseed_after_execution={expr.reseed_after_execution})"
|
||||
return self._make_tree_lines(label, expr=expr)
|
||||
|
||||
def visit_uuid(self, expr: "UUIDExpr") -> str:
|
||||
return self._make_tree_lines("UUID()", expr=expr)
|
||||
|
||||
|
||||
class _InlineExprReprVisitor(_ExprVisitor[str]):
|
||||
"""Visitor that generates concise inline string representations of expressions.
|
||||
|
||||
This visitor creates single-line string representations suitable for displaying
|
||||
in operator names, log messages, etc. It aims to be human-readable while keeping
|
||||
the representation compact.
|
||||
"""
|
||||
|
||||
def __init__(self, max_literal_length: int = 20):
|
||||
"""Initialize the inline representation visitor.
|
||||
|
||||
Args:
|
||||
max_literal_length: Maximum length for literal value representations
|
||||
"""
|
||||
self._max_literal_length = max_literal_length
|
||||
|
||||
def visit_column(self, expr: "ColumnExpr") -> str:
|
||||
"""Visit a column expression and return its inline representation."""
|
||||
return f"col({expr.name!r})"
|
||||
|
||||
def visit_literal(self, expr: "LiteralExpr") -> str:
|
||||
"""Visit a literal expression and return its inline representation."""
|
||||
value_repr = repr(expr.value)
|
||||
if len(value_repr) > self._max_literal_length:
|
||||
value_repr = value_repr[: self._max_literal_length - 3] + "..."
|
||||
return value_repr
|
||||
|
||||
def visit_binary(self, expr: "BinaryExpr") -> str:
|
||||
"""Visit a binary expression and return its inline representation."""
|
||||
left_str = self.visit(expr.left)
|
||||
right_str = self.visit(expr.right)
|
||||
|
||||
# Add parentheses around child binary expressions to avoid ambiguity
|
||||
if isinstance(expr.left, BinaryExpr):
|
||||
left_str = f"({left_str})"
|
||||
if isinstance(expr.right, BinaryExpr):
|
||||
right_str = f"({right_str})"
|
||||
|
||||
op_str = _INLINE_OP_SYMBOLS.get(expr.op, expr.op.name.lower())
|
||||
return f"{left_str} {op_str} {right_str}"
|
||||
|
||||
def visit_unary(self, expr: "UnaryExpr") -> str:
|
||||
"""Visit a unary expression and return its inline representation."""
|
||||
operand_str = self.visit(expr.operand)
|
||||
|
||||
# Add parentheses around binary expression operands to avoid ambiguity
|
||||
if isinstance(expr.operand, BinaryExpr):
|
||||
operand_str = f"({operand_str})"
|
||||
|
||||
# Map operations to symbols/functions
|
||||
if expr.op == Operation.NOT:
|
||||
return f"~{operand_str}"
|
||||
elif expr.op == Operation.IS_NULL:
|
||||
return f"{operand_str}.is_null()"
|
||||
elif expr.op == Operation.IS_NOT_NULL:
|
||||
return f"{operand_str}.is_not_null()"
|
||||
else:
|
||||
return f"{expr.op.name.lower()}({operand_str})"
|
||||
|
||||
def visit_alias(self, expr: "AliasExpr") -> str:
|
||||
"""Visit an alias expression and return its inline representation."""
|
||||
inner_str = self.visit(expr.expr)
|
||||
return f"{inner_str}.alias({expr.name!r})"
|
||||
|
||||
def visit_udf(self, expr: "UDFExpr") -> str:
|
||||
"""Visit a UDF expression and return its inline representation."""
|
||||
# Get function name for better readability
|
||||
# For callable objects (instances with __call__), use the class name
|
||||
fn_name = getattr(expr.fn, "__name__", expr.fn.__class__.__name__)
|
||||
|
||||
# Build argument list
|
||||
args_str = []
|
||||
for arg in expr.args:
|
||||
args_str.append(self.visit(arg))
|
||||
for key, value in expr.kwargs.items():
|
||||
args_str.append(f"{key}={self.visit(value)}")
|
||||
|
||||
args_repr = ", ".join(args_str) if args_str else ""
|
||||
return f"{fn_name}({args_repr})"
|
||||
|
||||
def visit_download(self, expr: "DownloadExpr") -> str:
|
||||
"""Visit a download expression and return its inline representation."""
|
||||
return f"download({expr.uri_column_name!r})"
|
||||
|
||||
def visit_star(self, expr: "StarExpr") -> str:
|
||||
"""Visit a star expression and return its inline representation."""
|
||||
return "col(*)"
|
||||
|
||||
def visit_monotonically_increasing_id(
|
||||
self, expr: "MonotonicallyIncreasingIdExpr"
|
||||
) -> str:
|
||||
"""Visit a monotonically_increasing_id expression and return its inline representation."""
|
||||
return "monotonically_increasing_id()"
|
||||
|
||||
def visit_random(self, expr: "RandomExpr") -> str:
|
||||
"""Visit a random expression and return its inline representation."""
|
||||
return "random()"
|
||||
|
||||
def visit_uuid(self, expr: "UUIDExpr") -> str:
|
||||
"""Visit a uuid expression and return its inline representation."""
|
||||
return "uuid()"
|
||||
|
||||
|
||||
class _StructuralFingerprintVisitor(_ExprVisitor[Hashable]):
|
||||
"""Visitor that computes a hashable structural fingerprint for an expression.
|
||||
|
||||
Two expressions that are structurally equivalent produce equal fingerprints,
|
||||
so the fingerprint can be used as a cheap bucketing key before falling back to
|
||||
full ``structurally_equals`` comparison (e.g. for common sub-expression
|
||||
elimination).
|
||||
"""
|
||||
|
||||
def visit_column(self, expr: ColumnExpr) -> Hashable:
|
||||
return _column_fingerprint_key(expr)
|
||||
|
||||
def visit_literal(self, expr: LiteralExpr) -> Hashable:
|
||||
return _literal_fingerprint_key(expr)
|
||||
|
||||
def visit_binary(self, expr: BinaryExpr) -> Hashable:
|
||||
return _binary_fingerprint_key(
|
||||
expr,
|
||||
self.visit(expr.left),
|
||||
self.visit(expr.right),
|
||||
)
|
||||
|
||||
def visit_unary(self, expr: UnaryExpr) -> Hashable:
|
||||
return _unary_fingerprint_key(expr, self.visit(expr.operand))
|
||||
|
||||
def visit_udf(self, expr: UDFExpr) -> Hashable:
|
||||
return _udf_fingerprint_key(
|
||||
expr,
|
||||
tuple(self.visit(arg) for arg in expr.args),
|
||||
tuple(
|
||||
(k, self.visit(v))
|
||||
for k, v in sorted(expr.kwargs.items(), key=lambda item: item[0])
|
||||
),
|
||||
)
|
||||
|
||||
def visit_alias(self, expr: AliasExpr) -> Hashable:
|
||||
return _alias_fingerprint_key(expr, self.visit(expr.expr))
|
||||
|
||||
def visit_download(self, expr: DownloadExpr) -> Hashable:
|
||||
return _download_fingerprint_key(expr)
|
||||
|
||||
def visit_star(self, expr: StarExpr) -> Hashable:
|
||||
return _star_fingerprint_key()
|
||||
|
||||
def visit_monotonically_increasing_id(
|
||||
self, expr: MonotonicallyIncreasingIdExpr
|
||||
) -> Hashable:
|
||||
return _monotonically_increasing_id_fingerprint_key(expr)
|
||||
|
||||
def visit_random(self, expr: RandomExpr) -> Hashable:
|
||||
return _random_fingerprint_key(expr)
|
||||
|
||||
def visit_uuid(self, expr: UUIDExpr) -> Hashable:
|
||||
return _uuid_fingerprint_key(expr)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _ExpressionOccurrence:
|
||||
expr: Expr
|
||||
key: Hashable
|
||||
depth: int
|
||||
|
||||
|
||||
class _StructuralFingerprintOccurrenceCollector(_ExprVisitor[Hashable]):
|
||||
"""Collect expression occurrences while computing structural keys bottom-up."""
|
||||
|
||||
def __init__(self):
|
||||
self._occurrences: List[_ExpressionOccurrence] = []
|
||||
self._depth = 0
|
||||
|
||||
def get_occurrences(self) -> List[_ExpressionOccurrence]:
|
||||
return self._occurrences
|
||||
|
||||
def _visit_child(self, expr: Expr) -> Hashable:
|
||||
self._depth += 1
|
||||
try:
|
||||
return self.visit(expr)
|
||||
finally:
|
||||
self._depth -= 1
|
||||
|
||||
def _record(self, expr: Expr, key: Hashable) -> Hashable:
|
||||
self._occurrences.append(
|
||||
_ExpressionOccurrence(
|
||||
expr=expr,
|
||||
key=key,
|
||||
depth=self._depth,
|
||||
)
|
||||
)
|
||||
return key
|
||||
|
||||
def visit_column(self, expr: ColumnExpr) -> Hashable:
|
||||
return self._record(expr, _column_fingerprint_key(expr))
|
||||
|
||||
def visit_literal(self, expr: LiteralExpr) -> Hashable:
|
||||
return self._record(expr, _literal_fingerprint_key(expr))
|
||||
|
||||
def visit_binary(self, expr: BinaryExpr) -> Hashable:
|
||||
return self._record(
|
||||
expr,
|
||||
_binary_fingerprint_key(
|
||||
expr,
|
||||
self._visit_child(expr.left),
|
||||
self._visit_child(expr.right),
|
||||
),
|
||||
)
|
||||
|
||||
def visit_unary(self, expr: UnaryExpr) -> Hashable:
|
||||
return self._record(
|
||||
expr,
|
||||
_unary_fingerprint_key(expr, self._visit_child(expr.operand)),
|
||||
)
|
||||
|
||||
def visit_udf(self, expr: UDFExpr) -> Hashable:
|
||||
return self._record(
|
||||
expr,
|
||||
_udf_fingerprint_key(
|
||||
expr,
|
||||
tuple(self._visit_child(arg) for arg in expr.args),
|
||||
tuple(
|
||||
(k, self._visit_child(v))
|
||||
for k, v in sorted(expr.kwargs.items(), key=lambda item: item[0])
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
def visit_alias(self, expr: AliasExpr) -> Hashable:
|
||||
return self._record(
|
||||
expr,
|
||||
_alias_fingerprint_key(expr, self._visit_child(expr.expr)),
|
||||
)
|
||||
|
||||
def visit_download(self, expr: DownloadExpr) -> Hashable:
|
||||
return self._record(expr, _download_fingerprint_key(expr))
|
||||
|
||||
def visit_star(self, expr: StarExpr) -> Hashable:
|
||||
return self._record(expr, _star_fingerprint_key())
|
||||
|
||||
def visit_monotonically_increasing_id(
|
||||
self, expr: MonotonicallyIncreasingIdExpr
|
||||
) -> Hashable:
|
||||
return self._record(expr, _monotonically_increasing_id_fingerprint_key(expr))
|
||||
|
||||
def visit_random(self, expr: RandomExpr) -> Hashable:
|
||||
return self._record(expr, _random_fingerprint_key(expr))
|
||||
|
||||
def visit_uuid(self, expr: UUIDExpr) -> Hashable:
|
||||
return self._record(expr, _uuid_fingerprint_key(expr))
|
||||
|
||||
|
||||
def get_column_references(expr: Expr) -> List[str]:
|
||||
"""Extract all column references from an expression.
|
||||
|
||||
This is a convenience function that creates a _ColumnReferenceCollector,
|
||||
visits the expression tree, and returns the list of referenced column names.
|
||||
|
||||
Args:
|
||||
expr: The expression to extract column references from.
|
||||
|
||||
Returns:
|
||||
List of column names referenced in the expression, in order of appearance.
|
||||
|
||||
Example:
|
||||
>>> from ray.data.expressions import col
|
||||
>>> expr = (col("a") > 5) & (col("b") == "test")
|
||||
>>> get_column_references(expr)
|
||||
['a', 'b']
|
||||
"""
|
||||
collector = _ColumnReferenceCollector()
|
||||
collector.visit(expr)
|
||||
return collector.get_column_refs()
|
||||
@@ -0,0 +1,124 @@
|
||||
import uuid
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
|
||||
from ray.data._internal.execution.interfaces.task_context import TaskContext
|
||||
from ray.data.block import BlockColumn, BlockType
|
||||
|
||||
|
||||
def eval_random(
|
||||
num_rows: int,
|
||||
block_type: BlockType,
|
||||
*,
|
||||
seed: int | None = None,
|
||||
reseed_after_execution: bool = True,
|
||||
instance_id: str | None = None,
|
||||
) -> BlockColumn:
|
||||
"""Implementation of the random expression.
|
||||
|
||||
Args:
|
||||
num_rows: The number of rows to generate random values for.
|
||||
block_type: The type of block to generate random values for.
|
||||
seed: The seed to use for the random number generator.
|
||||
reseed_after_execution: Whether to reseed the random number generator after each execution.
|
||||
instance_id: Unique identifier for the random expression instance, used to isolate
|
||||
block count state when a single task processes multiple blocks.
|
||||
|
||||
Returns:
|
||||
A BlockColumn containing the random values.
|
||||
|
||||
Raises:
|
||||
TypeError: If the block type is not supported.
|
||||
"""
|
||||
|
||||
if seed is not None:
|
||||
# Numpy allows using a seed sequence (list of integers) to initialize
|
||||
# a random number generator. This allows us to maintain reproduciblity while
|
||||
# ensuring randomness in parallel execution.
|
||||
# See https://numpy.org/doc/2.2/reference/random/parallel.html#sequence-of-integer-seeds
|
||||
# Below we uses four components to create a seed sequence (fastest changing component first):
|
||||
# 1. A per-block counter within the task (to differentiate blocks in the same task)
|
||||
# 2. An index based on the remote task in Ray Data
|
||||
# 3. An incrementing index of Ray Dataset execution (e.g., multiple training epochs)
|
||||
# 4. A base seed fixed by the user
|
||||
|
||||
ctx = TaskContext.get_current()
|
||||
|
||||
if ctx is None:
|
||||
import warnings
|
||||
|
||||
warnings.warn(
|
||||
"TaskContext is not available for random() expression with seed. "
|
||||
"Falling back to task_idx=0 for all tasks, which reduces the parallelism "
|
||||
"benefits of random number generation. If you see this warning, please "
|
||||
"report it as it may indicate an execution context issue.",
|
||||
stacklevel=2,
|
||||
)
|
||||
task_idx = 0
|
||||
block_idx = 0
|
||||
else:
|
||||
task_idx = ctx.task_idx
|
||||
|
||||
# Key the counter by expression instance ID so that multiple expressions
|
||||
# in the same projection will have isolated block count state.
|
||||
# This is required because a single task may process multiple blocks if
|
||||
# the upstream data source does not compress the data into a single block.
|
||||
if instance_id is not None:
|
||||
counter_key = f"_random_{instance_id}_counter"
|
||||
block_idx = ctx.kwargs.get(counter_key, 0)
|
||||
ctx.kwargs[counter_key] = block_idx + 1
|
||||
else:
|
||||
block_idx = 0
|
||||
|
||||
if reseed_after_execution:
|
||||
from ray.data.context import DataContext
|
||||
|
||||
data_context = (
|
||||
DataContext.get_current()
|
||||
) # get or create DataContext, never None
|
||||
execution_idx = data_context._execution_idx
|
||||
else:
|
||||
execution_idx = 0
|
||||
|
||||
# Numpy recommends fastest changing component to be the first element
|
||||
block_seed = [block_idx, task_idx, execution_idx, seed]
|
||||
else:
|
||||
block_seed = None
|
||||
|
||||
rng = np.random.default_rng(block_seed)
|
||||
random_values = rng.random(num_rows)
|
||||
|
||||
# Convert to appropriate format based on block type
|
||||
if block_type == BlockType.PANDAS:
|
||||
return pd.Series(random_values, dtype=np.float64)
|
||||
elif block_type == BlockType.ARROW:
|
||||
return pa.array(random_values, type=pa.float64())
|
||||
|
||||
raise TypeError(f"Unsupported block type: {block_type}")
|
||||
|
||||
|
||||
def eval_uuid(
|
||||
num_rows: int,
|
||||
block_type: BlockType,
|
||||
) -> BlockColumn:
|
||||
"""Implementation of the uuid expression.
|
||||
|
||||
Args:
|
||||
num_rows: The number of rows to generate uuid values for.
|
||||
block_type: The type of block to generate uuid values for.
|
||||
|
||||
Returns:
|
||||
A BlockColumn containing the uuid values.
|
||||
|
||||
Raises:
|
||||
TypeError: If the block type is not supported.
|
||||
"""
|
||||
arr = [str(uuid.uuid4()) for _ in range(num_rows)]
|
||||
if block_type == BlockType.PANDAS:
|
||||
return pd.Series(arr, dtype=pd.StringDtype())
|
||||
elif block_type == BlockType.ARROW:
|
||||
return pa.array(arr, type=pa.string())
|
||||
|
||||
raise TypeError(f"Unsupported block type: {block_type}")
|
||||
Reference in New Issue
Block a user