chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,446 @@
|
||||
import copy
|
||||
from dataclasses import dataclass, is_dataclass, replace
|
||||
from typing import List, Optional
|
||||
|
||||
from ray.data._internal.logical.interfaces import (
|
||||
LogicalOperator,
|
||||
LogicalOperatorSupportsPredicatePassThrough,
|
||||
LogicalOperatorSupportsPredicatePushdown,
|
||||
LogicalPlan,
|
||||
PredicatePassThroughBehavior,
|
||||
Rule,
|
||||
)
|
||||
from ray.data._internal.logical.operators import (
|
||||
AbstractAllToAll,
|
||||
AbstractMap,
|
||||
Filter,
|
||||
Join,
|
||||
Limit,
|
||||
Project,
|
||||
RandomShuffle,
|
||||
Repartition,
|
||||
Union,
|
||||
)
|
||||
from ray.data._internal.planner.plan_expression.expression_visitors import (
|
||||
_ColumnSubstitutionVisitor,
|
||||
)
|
||||
from ray.data.expressions import BinaryExpr, Expr, Operation, col
|
||||
|
||||
__all__ = [
|
||||
"PredicatePushdown",
|
||||
]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _ConvertibilitySplit:
|
||||
"""Result of splitting a predicate by PyArrow convertibility.
|
||||
|
||||
Attributes:
|
||||
convertible: The conjuncts that can be lowered to PyArrow and pushed
|
||||
into the datasource. ``None`` if nothing is convertible.
|
||||
residual: The conjuncts that cannot be lowered and must remain as a
|
||||
``Filter``. ``None`` if everything is convertible.
|
||||
"""
|
||||
|
||||
convertible: Optional[Expr]
|
||||
residual: Optional[Expr]
|
||||
|
||||
|
||||
class PredicatePushdown(Rule):
|
||||
"""Pushes down predicates across the graph.
|
||||
|
||||
This rule performs the following optimizations:
|
||||
1. Combines chained Filter operators with compatible expressions
|
||||
2. Pushes filter expressions through eligible operators using trait-based rules
|
||||
3. Pushes filters into data sources that support predicate pushdown
|
||||
|
||||
Eligibility is determined by the LogicalOperatorSupportsPredicatePassThrough trait, which operators
|
||||
implement to declare their pushdown behavior:
|
||||
- PASSTHROUGH: Filter passes through unchanged (Sort, Repartition, Shuffle, Limit)
|
||||
- PASSTHROUGH_WITH_SUBSTITUTION: Filter passes through with column rebinding (Project)
|
||||
- PUSH_INTO_BRANCHES: Filter is pushed into each branch (Union)
|
||||
- CONDITIONAL: Filter may be pushed based on analysis (Join - analyzes which side
|
||||
the predicate references and pushes to that side if safe for the join type)
|
||||
"""
|
||||
|
||||
def apply(self, plan: LogicalPlan) -> LogicalPlan:
|
||||
"""Apply predicate pushdown optimization to the logical plan."""
|
||||
dag = plan.dag
|
||||
new_dag = dag._apply_transform(self._try_fuse_filters)
|
||||
new_dag = new_dag._apply_transform(self._try_push_down_predicate)
|
||||
return LogicalPlan(new_dag, plan.context) if dag is not new_dag else plan
|
||||
|
||||
@classmethod
|
||||
def _is_valid_filter_operator(cls, op: LogicalOperator) -> bool:
|
||||
return isinstance(op, Filter) and op.is_expression_based()
|
||||
|
||||
@classmethod
|
||||
def _try_fuse_filters(cls, op: LogicalOperator) -> LogicalOperator:
|
||||
"""Fuse consecutive Filter operators with compatible expressions."""
|
||||
if not cls._is_valid_filter_operator(op):
|
||||
return op
|
||||
|
||||
input_op = op.input_dependencies[0]
|
||||
if not cls._is_valid_filter_operator(input_op):
|
||||
return op
|
||||
|
||||
# Do not fuse across a non-idempotent predicate (random/uuid/
|
||||
# monotonically_increasing_id): combining moves where a predicate is
|
||||
# evaluated (and thus which rows it sees), changing results.
|
||||
if (
|
||||
not op.predicate_expr.is_idempotent()
|
||||
or not input_op.predicate_expr.is_idempotent()
|
||||
):
|
||||
return op
|
||||
|
||||
# Combine predicates
|
||||
combined_predicate = op.predicate_expr & input_op.predicate_expr
|
||||
|
||||
# Create new filter on the input of the lower filter
|
||||
return Filter(
|
||||
predicate_expr=combined_predicate,
|
||||
input_dependencies=[input_op.input_dependencies[0]],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _can_push_filter_through_projection(
|
||||
cls, filter_op: "Filter", projection_op: Project
|
||||
) -> bool:
|
||||
"""Check if a filter can be pushed through a projection operator.
|
||||
|
||||
Returns False (blocks pushdown) if filter references:
|
||||
- Columns removed by select: select(['a']).filter(col('b'))
|
||||
- Computed columns: with_column('d', 4).filter(col('d'))
|
||||
- Old column names after rename: rename({'b': 'B'}).filter(col('b'))
|
||||
|
||||
Returns True (allows pushdown) for:
|
||||
- Columns present in output: select(['a', 'b']).filter(col('a'))
|
||||
- New column names after rename: rename({'b': 'B'}).filter(col('B'))
|
||||
- Rename chains with name reuse: rename({'a': 'b', 'b': 'c'}).filter(col('b'))
|
||||
(where 'b' is valid output created by a->b)
|
||||
"""
|
||||
from ray.data._internal.planner.plan_expression.expression_visitors import (
|
||||
_ColumnReferenceCollector,
|
||||
)
|
||||
from ray.data.expressions import AliasExpr, is_rename_expr
|
||||
|
||||
# Do not push a filter below a projection that produces a non-idempotent
|
||||
# column (random/uuid/monotonically_increasing_id): reordering changes the row
|
||||
# set / position the expression is evaluated over (e.g.
|
||||
# monotonically_increasing_id reassigned over the filtered subset).
|
||||
if not projection_op.is_idempotent():
|
||||
return False
|
||||
|
||||
collector = _ColumnReferenceCollector()
|
||||
collector.visit(filter_op.predicate_expr)
|
||||
predicate_columns = set(collector.get_column_refs() or [])
|
||||
|
||||
output_columns = set()
|
||||
new_names = set()
|
||||
original_columns_being_renamed = set()
|
||||
|
||||
for expr in projection_op.exprs:
|
||||
if expr.name is not None:
|
||||
# Collect output column names
|
||||
output_columns.add(expr.name)
|
||||
|
||||
# Process AliasExpr (computed columns or renames)
|
||||
if isinstance(expr, AliasExpr):
|
||||
new_names.add(expr.name)
|
||||
|
||||
# Check computed column: with_column('d', 4) creates AliasExpr(lit(4), 'd')
|
||||
if expr.name in predicate_columns and not is_rename_expr(expr):
|
||||
return False # Computed column
|
||||
|
||||
# Track old names being renamed for later check
|
||||
if is_rename_expr(expr):
|
||||
original_columns_being_renamed.add(expr.expr.name)
|
||||
|
||||
# Check if filter references columns removed by explicit select.
|
||||
# Valid if: projection includes all columns (star, UDF-fallback path)
|
||||
# OR predicate columns exist in the explicit output set (typed path,
|
||||
# where ``StarExpr`` is expanded into explicit ``col()`` refs in
|
||||
# ``Project.__post_init__`` when the input schema is known).
|
||||
has_required_columns = (
|
||||
projection_op.has_star_expr() or predicate_columns.issubset(output_columns)
|
||||
)
|
||||
if not has_required_columns:
|
||||
return False
|
||||
|
||||
# Find old names that are:
|
||||
# 1. Being renamed away (in original_columns_being_renamed), AND
|
||||
# 2. Referenced in predicate (in predicate_columns), AND
|
||||
# 3. NOT recreated as new names (not in new_names)
|
||||
#
|
||||
# Examples:
|
||||
# rename({'b': 'B'}).filter(col('b'))
|
||||
# → {'b'} & {'b'} - {'B'} = {'b'} → BLOCKS (old name 'b' no longer exists)
|
||||
#
|
||||
# rename({'a': 'b', 'b': 'c'}).filter(col('b'))
|
||||
# → {'a','b'} & {'b'} - {'b','c'} = {} → ALLOWS (new 'b' created by a->b)
|
||||
#
|
||||
# rename({'b': 'B'}).filter(col('B'))
|
||||
# → {'b'} & {'B'} - {'B'} = {} → ALLOWS (using new name 'B')
|
||||
invalid_old_names = (
|
||||
original_columns_being_renamed & predicate_columns
|
||||
) - new_names
|
||||
if invalid_old_names:
|
||||
return False # Old name after rename
|
||||
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def _substitute_predicate_columns(
|
||||
cls, predicate_expr: Expr, column_rename_map: dict[str, str]
|
||||
) -> Expr:
|
||||
"""Rebind column references in a predicate expression.
|
||||
|
||||
When pushing a predicate through a projection with column renames,
|
||||
we need to rewrite column references from new names to old names.
|
||||
|
||||
Args:
|
||||
predicate_expr: The predicate with new column names
|
||||
column_rename_map: Mapping from old_name -> new_name
|
||||
|
||||
Returns:
|
||||
The predicate rewritten to use old column names
|
||||
"""
|
||||
# Invert the mapping: new_name -> old_name (as col expression)
|
||||
# This is because the predicate uses new names and we need to map
|
||||
# them back to old names
|
||||
column_mapping = {
|
||||
new_col: col(old_col) for old_col, new_col in column_rename_map.items()
|
||||
}
|
||||
|
||||
visitor = _ColumnSubstitutionVisitor(column_mapping)
|
||||
return visitor.visit(predicate_expr)
|
||||
|
||||
@classmethod
|
||||
def _combine_with_and(
|
||||
cls, left: Optional[Expr], right: Optional[Expr]
|
||||
) -> Optional[Expr]:
|
||||
"""Combine two optional predicates with ``AND``, ignoring ``None``."""
|
||||
if left is not None and right is not None:
|
||||
return left & right
|
||||
return left if left is not None else right
|
||||
|
||||
@classmethod
|
||||
def _split_by_convertibility(cls, predicate: Expr) -> _ConvertibilitySplit:
|
||||
"""Split a predicate into PyArrow convertible and residual parts.
|
||||
|
||||
Walks the top level ``AND`` chain and buckets each conjunct by whether
|
||||
it can be lowered to PyArrow. The convertible part can be pushed into
|
||||
the datasource while the residual part stays as a ``Filter`` above it.
|
||||
|
||||
Args:
|
||||
predicate: The predicate expression to split.
|
||||
|
||||
Returns:
|
||||
A ``_ConvertibilitySplit`` whose ``convertible`` and ``residual``
|
||||
fields hold the two parts. Both are optional.
|
||||
"""
|
||||
if isinstance(predicate, BinaryExpr) and predicate.op == Operation.AND:
|
||||
left = cls._split_by_convertibility(predicate.left)
|
||||
right = cls._split_by_convertibility(predicate.right)
|
||||
return _ConvertibilitySplit(
|
||||
convertible=cls._combine_with_and(left.convertible, right.convertible),
|
||||
residual=cls._combine_with_and(left.residual, right.residual),
|
||||
)
|
||||
|
||||
if predicate._is_pyarrow_convertible():
|
||||
return _ConvertibilitySplit(convertible=predicate, residual=None)
|
||||
return _ConvertibilitySplit(convertible=None, residual=predicate)
|
||||
|
||||
@classmethod
|
||||
def _try_push_down_predicate(cls, op: LogicalOperator) -> LogicalOperator:
|
||||
"""Push Filter down through the operator tree."""
|
||||
if not cls._is_valid_filter_operator(op):
|
||||
return op
|
||||
filter_op: Filter = op
|
||||
input_op = filter_op.input_dependencies[0]
|
||||
predicate_expr = filter_op.predicate_expr
|
||||
|
||||
# Case 1: Check if operator supports predicate pushdown (e.g., Read).
|
||||
# The read stage never renames columns (renaming is always carried
|
||||
# by an ``AliasExpr`` in a ``Project`` operator above the read), so
|
||||
# the predicate above the read is already in the same column
|
||||
# namespace the scanner sees — no rebinding is required here.
|
||||
if (
|
||||
isinstance(input_op, LogicalOperatorSupportsPredicatePushdown)
|
||||
and input_op.supports_predicate_pushdown()
|
||||
):
|
||||
# Datasources evaluate pushed predicates via PyArrow. A predicate
|
||||
# that can't be lowered to PyArrow (e.g. it contains a UDF) must
|
||||
# stay as a Filter. Split the top level AND chain so the convertible
|
||||
# conjuncts can still be pushed while the residual ones are kept as
|
||||
# a Filter above the read.
|
||||
split = cls._split_by_convertibility(predicate_expr)
|
||||
|
||||
if split.convertible is None:
|
||||
return filter_op
|
||||
|
||||
result_op = input_op.apply_predicate(split.convertible)
|
||||
|
||||
# If the operator is unchanged (e.g., predicate references partition columns
|
||||
# that can't be pushed down), keep the Filter operator
|
||||
if result_op is input_op:
|
||||
return filter_op
|
||||
|
||||
# Convertible conjuncts were pushed into the read. Re-apply any
|
||||
# residual (non-convertible) conjuncts as a Filter above it.
|
||||
if split.residual is None:
|
||||
return result_op
|
||||
return Filter(predicate_expr=split.residual, input_dependencies=[result_op])
|
||||
|
||||
# Datasource pushdown (Case 1) only lowers PyArrow-convertible (hence
|
||||
# idempotent) conjuncts. Beyond that, do not relocate a filter whose predicate
|
||||
# is non-idempotent (random/uuid/monotonically_increasing_id): pushing it
|
||||
# through a pass-through operator, into Union branches, or to a Join side
|
||||
# changes the row set / position the expression is evaluated over.
|
||||
if not predicate_expr.is_idempotent():
|
||||
return filter_op
|
||||
|
||||
# Case 2: Check if operator allows predicates to pass through
|
||||
if isinstance(input_op, LogicalOperatorSupportsPredicatePassThrough):
|
||||
behavior = input_op.predicate_passthrough_behavior()
|
||||
|
||||
if behavior in (
|
||||
PredicatePassThroughBehavior.PASSTHROUGH,
|
||||
PredicatePassThroughBehavior.PASSTHROUGH_WITH_SUBSTITUTION,
|
||||
):
|
||||
# Both cases push through a single input with optional column rebinding
|
||||
assert len(input_op.input_dependencies) == 1, (
|
||||
f"{behavior.value} operators must have exactly 1 input, "
|
||||
f"got {len(input_op.input_dependencies)}"
|
||||
)
|
||||
|
||||
# Apply column substitution if needed
|
||||
if (
|
||||
behavior
|
||||
== PredicatePassThroughBehavior.PASSTHROUGH_WITH_SUBSTITUTION
|
||||
):
|
||||
# Check if we can safely push the filter through this projection
|
||||
if isinstance(
|
||||
input_op, Project
|
||||
) and not cls._can_push_filter_through_projection(
|
||||
filter_op, input_op
|
||||
):
|
||||
return filter_op
|
||||
|
||||
rename_map = input_op.get_column_substitutions()
|
||||
if rename_map:
|
||||
predicate_expr = cls._substitute_predicate_columns(
|
||||
predicate_expr, rename_map
|
||||
)
|
||||
|
||||
# Push filter through and recursively try to push further
|
||||
new_filter = Filter(
|
||||
predicate_expr=predicate_expr,
|
||||
input_dependencies=[input_op.input_dependencies[0]],
|
||||
)
|
||||
pushed_filter = cls._try_push_down_predicate(new_filter)
|
||||
|
||||
# Return input_op with the pushed filter as its input
|
||||
return cls._clone_op_with_new_inputs(input_op, [pushed_filter])
|
||||
|
||||
elif behavior == PredicatePassThroughBehavior.PUSH_INTO_BRANCHES:
|
||||
# Push into each branch (e.g., Union)
|
||||
# Apply filter to each branch and recursively push down
|
||||
new_inputs = []
|
||||
for branch_op in input_op.input_dependencies:
|
||||
branch_filter = Filter(
|
||||
predicate_expr=predicate_expr, input_dependencies=[branch_op]
|
||||
)
|
||||
pushed_branch = cls._try_push_down_predicate(branch_filter)
|
||||
new_inputs.append(pushed_branch)
|
||||
|
||||
# Return operator with filtered branches
|
||||
return cls._clone_op_with_new_inputs(input_op, new_inputs)
|
||||
|
||||
elif behavior == PredicatePassThroughBehavior.CONDITIONAL:
|
||||
# Handle conditional pushdown (e.g., Join)
|
||||
return cls._push_filter_through_conditionally(filter_op, input_op)
|
||||
|
||||
return filter_op
|
||||
|
||||
@classmethod
|
||||
def _push_filter_through_conditionally(
|
||||
cls, filter_op: Filter, conditional_op: LogicalOperator
|
||||
) -> LogicalOperator:
|
||||
"""Handle conditional pushdown for operators like Join.
|
||||
|
||||
For operators with multiple inputs, we can push predicates that reference
|
||||
only one side down to that side, when semantically safe.
|
||||
"""
|
||||
# Check if operator supports conditional pushdown by having the required method
|
||||
if not hasattr(conditional_op, "which_side_to_push_predicate"):
|
||||
return filter_op
|
||||
|
||||
push_side = conditional_op.which_side_to_push_predicate(
|
||||
filter_op.predicate_expr
|
||||
)
|
||||
|
||||
if push_side is None:
|
||||
# Cannot push through
|
||||
return filter_op
|
||||
|
||||
# Use the enum value directly as branch index
|
||||
branch_idx = push_side.value
|
||||
|
||||
# Push to the appropriate branch
|
||||
new_inputs = list(conditional_op.input_dependencies)
|
||||
branch_filter = Filter(
|
||||
predicate_expr=filter_op.predicate_expr,
|
||||
input_dependencies=[new_inputs[branch_idx]],
|
||||
)
|
||||
new_inputs[branch_idx] = cls._try_push_down_predicate(branch_filter)
|
||||
|
||||
# Return operator with updated input
|
||||
return cls._clone_op_with_new_inputs(conditional_op, new_inputs)
|
||||
|
||||
@classmethod
|
||||
def _clone_op_with_new_inputs(
|
||||
cls, op: LogicalOperator, new_inputs: List[LogicalOperator]
|
||||
) -> LogicalOperator:
|
||||
"""Clone an operator with new inputs.
|
||||
|
||||
Args:
|
||||
op: The operator to clone
|
||||
new_inputs: List of new input operators (can be single element list)
|
||||
|
||||
Returns:
|
||||
A shallow copy of the operator with updated input dependencies
|
||||
"""
|
||||
if isinstance(op, Limit):
|
||||
assert len(new_inputs) == 1, len(new_inputs)
|
||||
return Limit(op.limit, input_dependencies=[new_inputs[0]])
|
||||
if isinstance(op, AbstractMap) and is_dataclass(op):
|
||||
assert len(new_inputs) == 1, len(new_inputs)
|
||||
return replace(op, input_dependencies=[new_inputs[0]])
|
||||
if isinstance(op, AbstractAllToAll) and is_dataclass(op):
|
||||
assert len(new_inputs) == 1, len(new_inputs)
|
||||
kwargs = {"input_dependencies": [new_inputs[0]]}
|
||||
if isinstance(op, Repartition):
|
||||
kwargs["num_outputs"] = op.num_outputs
|
||||
if isinstance(op, RandomShuffle):
|
||||
kwargs["name"] = op.name
|
||||
return replace(op, **kwargs)
|
||||
if isinstance(op, Join) and is_dataclass(op):
|
||||
assert len(new_inputs) == 2, len(new_inputs)
|
||||
return Join(
|
||||
new_inputs[0],
|
||||
new_inputs[1],
|
||||
op.join_type,
|
||||
op.left_key_columns,
|
||||
op.right_key_columns,
|
||||
num_partitions=op.num_outputs,
|
||||
left_columns_suffix=op.left_columns_suffix,
|
||||
right_columns_suffix=op.right_columns_suffix,
|
||||
partition_size_hint=op.partition_size_hint,
|
||||
aggregator_ray_remote_args=op.aggregator_ray_remote_args,
|
||||
)
|
||||
if isinstance(op, Union) and is_dataclass(op):
|
||||
return Union(new_inputs)
|
||||
new_op = copy.copy(op)
|
||||
new_op.input_dependencies = new_inputs
|
||||
return new_op
|
||||
Reference in New Issue
Block a user