wehub-resource-sync
447 lines
18 KiB
Python
447 lines
18 KiB
Python
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
|