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