wehub-resource-sync
524 lines
20 KiB
Python
524 lines
20 KiB
Python
from typing import Dict, List, Optional, Set, Tuple
|
|
|
|
from ray.data._internal.logical.interfaces import (
|
|
LogicalOperator,
|
|
LogicalOperatorSupportsProjectionPushdown,
|
|
LogicalPlan,
|
|
Rule,
|
|
)
|
|
from ray.data._internal.logical.operators import Project
|
|
from ray.data._internal.planner.plan_expression.expression_visitors import (
|
|
_ColumnReferenceCollector,
|
|
_ColumnSubstitutionVisitor,
|
|
_is_col_expr,
|
|
)
|
|
from ray.data.expressions import (
|
|
AliasExpr,
|
|
Expr,
|
|
StarExpr,
|
|
is_rename_expr,
|
|
)
|
|
|
|
__all__ = [
|
|
"ProjectionPushdown",
|
|
]
|
|
|
|
|
|
def _collect_referenced_columns(exprs: List[Expr]) -> Optional[List[str]]:
|
|
"""
|
|
Extract all column names referenced by the given expressions.
|
|
|
|
Recursively traverses expression trees to find all ColumnExpr nodes
|
|
and collects their names.
|
|
|
|
Example: For expression "col1 + col2", returns {"col1", "col2"}
|
|
"""
|
|
# ``StarExpr`` is eagerly expanded to explicit ``col()`` refs in
|
|
# ``Project.__post_init__`` when the input schema is known. So this
|
|
# branch is hit only on the UDF-fallback path (Project on top of an
|
|
# opaque-schema input like ``MapBatches``), where we can't enumerate
|
|
# columns and have to fall back to "all columns" (``None``).
|
|
if any(isinstance(expr, StarExpr) for expr in exprs):
|
|
return None
|
|
|
|
collector = _ColumnReferenceCollector()
|
|
for expr in exprs or []:
|
|
collector.visit(expr)
|
|
|
|
return collector.get_column_refs()
|
|
|
|
|
|
def _analyze_upstream_project(
|
|
upstream_project: Project,
|
|
) -> Tuple[Set[str], dict[str, Expr], Set[str]]:
|
|
"""
|
|
Analyze what the upstream project produces and identifies removed columns.
|
|
|
|
Example: Upstream exprs [col("x").alias("y")] → removed_by_renames = {"x"} if "x" not in output
|
|
"""
|
|
output_column_names = {
|
|
expr.name for expr in upstream_project.exprs if not isinstance(expr, StarExpr)
|
|
}
|
|
|
|
# Compose column definitions in the form of a mapping of
|
|
# - Target column name
|
|
# - Target expression
|
|
output_column_defs = {
|
|
expr.name: expr for expr in _filter_out_star(upstream_project.exprs)
|
|
}
|
|
|
|
# Identify upstream input columns removed by renaming (ie not propagated into
|
|
# its output)
|
|
upstream_column_renaming_map = _extract_input_columns_renaming_mapping(
|
|
upstream_project.exprs
|
|
)
|
|
|
|
return (
|
|
output_column_names,
|
|
output_column_defs,
|
|
set(upstream_column_renaming_map.keys()),
|
|
)
|
|
|
|
|
|
def _validate_fusion(
|
|
downstream_project: Project,
|
|
upstream_has_all: bool,
|
|
upstream_output_columns: Set[str],
|
|
removed_by_renames: Set[str],
|
|
) -> Tuple[bool, Set[str]]:
|
|
"""
|
|
Validate if fusion is possible without rewriting expressions.
|
|
|
|
Args:
|
|
downstream_project: The downstream Project operator
|
|
upstream_has_all: True if the upstream Project has all columns, False otherwise
|
|
upstream_output_columns: Set of column names that are available in the upstream Project
|
|
removed_by_renames: Set of column names that are removed by renames in the upstream Project
|
|
|
|
Returns:
|
|
Tuple of (is_valid, missing_columns)
|
|
- is_valid: True if all expressions can be fused, False otherwise
|
|
- missing_columns: Set of column names that are referenced but not available
|
|
|
|
Example: Downstream refs "x" but upstream renamed "x" to "y" and dropped "x"
|
|
→ (False, {"x"})
|
|
"""
|
|
missing_columns = set()
|
|
|
|
for expr in downstream_project.exprs:
|
|
if isinstance(expr, StarExpr):
|
|
continue
|
|
|
|
column_refs = _collect_referenced_columns([expr])
|
|
column_refs_set = set(column_refs or [])
|
|
|
|
columns_from_original = column_refs_set - (
|
|
column_refs_set & upstream_output_columns
|
|
)
|
|
|
|
# Validate accessibility
|
|
if not upstream_has_all and columns_from_original:
|
|
# Example: Upstream selects ["a", "b"], Downstream refs "c" → can't fuse
|
|
missing_columns.update(columns_from_original)
|
|
|
|
if any(col in removed_by_renames for col in columns_from_original):
|
|
# Example: Upstream renames "x" to "y" (dropping "x"), Downstream refs "x" → can't fuse
|
|
removed_cols = {
|
|
col for col in columns_from_original if col in removed_by_renames
|
|
}
|
|
missing_columns.update(removed_cols)
|
|
|
|
is_valid = len(missing_columns) == 0
|
|
return is_valid, missing_columns
|
|
|
|
|
|
def _would_duplicate_nonidempotent_expr(
|
|
upstream_project: Project,
|
|
downstream_project: Project,
|
|
upstream_column_defs: Dict[str, Expr],
|
|
) -> bool:
|
|
"""Return ``True`` if fusing would materialize a non-idempotent column >1 time.
|
|
|
|
Fusion inlines each upstream output column definition into every downstream
|
|
reference. For a non-idempotent definition (random/uuid/monotonically_increasing_id)
|
|
this changes its evaluation count, so we block the fusion and let the upstream
|
|
evaluate it exactly once.
|
|
|
|
Materialization count for an upstream column post-fusion is its downstream reference
|
|
multiplicity, plus one if it survives as a passthrough (the composition/star case).
|
|
"""
|
|
nonidem_cols = {
|
|
name
|
|
for name, def_expr in upstream_column_defs.items()
|
|
if not def_expr.is_idempotent()
|
|
}
|
|
if not nonidem_cols:
|
|
return False
|
|
|
|
# Downstream reference multiplicity (counts ``x + x`` as 2).
|
|
counter = _ColumnReferenceCollector()
|
|
# Note: We ignore _common_sub_exprs field as CSE rule is applied post-optimization.
|
|
# If order is to be changed, we also need to invoke the _ColumnReferenceCollector
|
|
# on _common_sub_exprs.
|
|
for e in _filter_out_star(downstream_project.exprs):
|
|
counter.visit(e)
|
|
ref_counts = counter.get_counts()
|
|
|
|
# In the composition (downstream-star) case the upstream column also survives in
|
|
# the fused output unless downstream renames it away, adding one materialization.
|
|
passthrough: Set[str] = set()
|
|
if downstream_project.has_star_expr():
|
|
renamed_away = _extract_input_columns_renaming_mapping(downstream_project.exprs)
|
|
passthrough = {
|
|
e.name
|
|
for e in upstream_project.exprs
|
|
if not isinstance(e, StarExpr) and e.name not in renamed_away
|
|
}
|
|
|
|
for name in nonidem_cols:
|
|
total = ref_counts.get(name, 0) + (1 if name in passthrough else 0)
|
|
if total > 1:
|
|
return True
|
|
return False
|
|
|
|
|
|
def _try_fuse(upstream_project: Project, downstream_project: Project) -> Project:
|
|
"""
|
|
Attempt to merge two consecutive Project operations into one.
|
|
|
|
Example: Upstream: [star(), col("x").alias("y")], Downstream: [star(), (col("y") + 1).alias("z")] → Fused: [star(), (col("x") + 1).alias("z")]
|
|
"""
|
|
# Check resource compatibility before attempting fusion
|
|
# This ensures with_column respects resource boundaries like map_batches does
|
|
from ray.data._internal.logical.rules.operator_fusion import (
|
|
FuseOperators,
|
|
are_op_remote_args_compatible,
|
|
)
|
|
|
|
# Check if remote args (num_cpus, num_gpus, etc.) are compatible and that
|
|
# neither op specifies a `ray_remote_args_fn`.
|
|
if not are_op_remote_args_compatible(upstream_project, downstream_project):
|
|
# Resources don't match - cannot fuse
|
|
return downstream_project
|
|
|
|
# Check if compute strategies are compatible
|
|
fused_compute = FuseOperators._fuse_compute_strategy(
|
|
upstream_project.compute, downstream_project.compute
|
|
)
|
|
if fused_compute is None:
|
|
# Compute strategies incompatible - cannot fuse
|
|
return downstream_project
|
|
|
|
upstream_has_star: bool = upstream_project.has_star_expr()
|
|
|
|
# TODO add validations that
|
|
# - exprs only depend on input attrs (ie no dep on output of other exprs)
|
|
|
|
# Analyze upstream
|
|
(
|
|
upstream_output_cols,
|
|
upstream_column_defs,
|
|
upstream_input_cols_removed,
|
|
) = _analyze_upstream_project(upstream_project)
|
|
|
|
# Validate fusion possibility
|
|
is_valid, missing_columns = _validate_fusion(
|
|
downstream_project,
|
|
upstream_has_star,
|
|
upstream_output_cols,
|
|
upstream_input_cols_removed,
|
|
)
|
|
|
|
if not is_valid:
|
|
# Raise KeyError to match expected error type in tests
|
|
raise KeyError(
|
|
f"Column(s) {sorted(missing_columns)} not found. "
|
|
f"Available columns: {sorted(upstream_output_cols) if not upstream_has_star else 'all columns (has star)'}"
|
|
)
|
|
|
|
if _would_duplicate_nonidempotent_expr(
|
|
upstream_project, downstream_project, upstream_column_defs
|
|
):
|
|
# Fusing would inline a non-idempotent expression (random/uuid/
|
|
# monotonically_increasing_id) into multiple references, changing its
|
|
# evaluation count. Leave the two Projects unfused so the upstream evaluates
|
|
# it exactly once and downstream reads the materialized column.
|
|
return downstream_project
|
|
|
|
# Following invariants are upheld for each ``Project`` logical op:
|
|
#
|
|
# 1. ``Project``s list of expressions are bound to op's input columns **only**
|
|
# (ie there could be no inter-dependency b/w expressions themselves)
|
|
#
|
|
# 2. `Each of expressions on the `Project``s list constitutes an output
|
|
# column definition, where column's name is derived from ``expr.name`` and
|
|
# column itself is derived by executing that expression against the op's
|
|
# input block.
|
|
#
|
|
# Therefore to abide by and satisfy aforementioned invariants, when fusing
|
|
# 2 ``Project`` operators, following scenarios are considered:
|
|
#
|
|
# 1. Composition: downstream including (and potentially renaming) upstream
|
|
# output columns (this is the case when downstream holds ``StarExpr``).
|
|
#
|
|
# 2. Projection: downstream projecting upstream output columns (by for ex,
|
|
# only selecting & transforming some of the upstream output columns).
|
|
#
|
|
|
|
# Upstream output column refs inside downstream expressions need to be bound
|
|
# to upstream output column definitions to satisfy invariant #1 (common for both
|
|
# composition/projection cases)
|
|
v = _ColumnSubstitutionVisitor(upstream_column_defs)
|
|
|
|
rebound_downstream_exprs = [
|
|
v.visit(e) for e in _filter_out_star(downstream_project.exprs)
|
|
]
|
|
|
|
if not downstream_project.has_star_expr():
|
|
# Projection case: this is when downstream is a *selection* (ie, not including
|
|
# the upstream columns with ``StarExpr``). With eager expansion of
|
|
# ``StarExpr`` in ``Project.__post_init__`` this is the common case
|
|
# for typed chains (no ``StarExpr`` reaches the optimizer).
|
|
#
|
|
# Example:
|
|
# Upstream: Project([col("a").alias("b")])
|
|
# Downstream: Project([col("b").alias("c")])
|
|
#
|
|
# Result: Project([col("a").alias("c")])
|
|
new_exprs = rebound_downstream_exprs
|
|
else:
|
|
# Composition case: downstream has ``StarExpr`` (entailing that downstream
|
|
# output will be including all of the upstream output columns). This
|
|
# is the UDF-fallback path; for typed chains
|
|
# ``Project.__post_init__`` would have replaced the ``StarExpr``
|
|
# with explicit ``col()`` refs already.
|
|
#
|
|
# Example 1:
|
|
# Upstream: [star(), col("a").alias("b")],
|
|
# Downstream: [star(), col("b").alias("c")]
|
|
#
|
|
# Result: [star(), col("a").alias("b"), col("a").alias("c")]
|
|
#
|
|
# Example 2:
|
|
# Input (columns): ["a", "b"]
|
|
# Upstream: [star({"b": "z"}), col("a").alias("x")],
|
|
# Downstream: [star({"x": "y"}), col("z")]
|
|
#
|
|
# Result: [star(), col("a").alias("y"), col("b").alias("z")]
|
|
|
|
# Extract downstream's input column rename map (downstream inputs are
|
|
# upstream's outputs)
|
|
downstream_input_column_rename_map = _extract_input_columns_renaming_mapping(
|
|
downstream_project.exprs
|
|
)
|
|
# Collect upstream output column expression "projected" to become
|
|
# downstream expressions
|
|
projected_upstream_output_col_exprs = []
|
|
|
|
# When fusing 2 projections
|
|
for e in upstream_project.exprs:
|
|
# NOTE: We have to filter out upstream output columns that are
|
|
# being *renamed* by downstream expression
|
|
if e.name not in downstream_input_column_rename_map:
|
|
projected_upstream_output_col_exprs.append(e)
|
|
|
|
new_exprs = projected_upstream_output_col_exprs + rebound_downstream_exprs
|
|
|
|
return Project(
|
|
exprs=new_exprs,
|
|
input_dependencies=[upstream_project.input_dependencies[0]],
|
|
compute=fused_compute,
|
|
ray_remote_args=downstream_project.ray_remote_args,
|
|
)
|
|
|
|
|
|
def _filter_out_star(exprs: List[Expr]) -> List[Expr]:
|
|
return [e for e in exprs if not isinstance(e, StarExpr)]
|
|
|
|
|
|
class ProjectionPushdown(Rule):
|
|
"""
|
|
Optimization rule that pushes projections (column selections) down the query plan.
|
|
|
|
This rule performs two optimizations:
|
|
1. Fuses consecutive Project operations to eliminate redundant projections
|
|
2. Pushes projections into data sources (e.g., Read operations) to enable
|
|
column pruning at the storage layer
|
|
"""
|
|
|
|
def apply(self, plan: LogicalPlan) -> LogicalPlan:
|
|
"""Apply projection pushdown optimization to the entire plan."""
|
|
dag = plan.dag
|
|
# Insert a pruning projection below consuming ops (e.g. ``Aggregate``)
|
|
# first, so the fuse/push steps can carry the narrowed columns into the
|
|
# read.
|
|
new_dag = dag._apply_transform(self._prune_aggregate_input)
|
|
new_dag = new_dag._apply_transform(self._try_fuse_projects)
|
|
new_dag = new_dag._apply_transform(self._push_projection_into_read_op)
|
|
return LogicalPlan(new_dag, plan.context) if dag is not new_dag else plan
|
|
|
|
@classmethod
|
|
def _prune_aggregate_input(cls, op: LogicalOperator) -> LogicalOperator:
|
|
"""Insert a ``Project`` below an ``Aggregate`` that keeps only the
|
|
columns it consumes (group keys + each aggregation's target column).
|
|
|
|
The aggregation drops every other column anyway, so pruning them before
|
|
the shuffle avoids dragging unused (often wide) columns through it --
|
|
which otherwise inflates both the aggregator's memory reservation and
|
|
the bytes shuffled. The inserted projection is fused/pushed toward the
|
|
read by the steps in ``apply``.
|
|
"""
|
|
from dataclasses import replace
|
|
|
|
from ray.data._internal.logical.operators.all_to_all_operator import Aggregate
|
|
from ray.data.aggregate import AggregateFnV2
|
|
from ray.data.expressions import col
|
|
|
|
if not isinstance(op, Aggregate):
|
|
return op
|
|
|
|
keys = op.key if isinstance(op.key, list) else ([op.key] if op.key else [])
|
|
required: List[str] = list(keys)
|
|
for agg in op.aggs:
|
|
# A generic ``AggregateFn`` may read arbitrary columns, so only
|
|
# prune when every aggregation declares the column it reads.
|
|
if not isinstance(agg, AggregateFnV2):
|
|
return op
|
|
target = agg.get_target_column()
|
|
if target is not None:
|
|
required.append(target)
|
|
|
|
# Order-preserving dedup; empty means nothing safe to prune to (e.g. a
|
|
# global count reading no columns).
|
|
required = list(dict.fromkeys(required))
|
|
if not required:
|
|
return op
|
|
|
|
input_op = op.input_dependencies[0]
|
|
schema = input_op.infer_schema()
|
|
if schema is None or not hasattr(schema, "names"):
|
|
return op # unknown schema: can't prove pruning helps
|
|
|
|
# Insert only when ``required`` is a strict subset of the input columns:
|
|
# this guarantees there's something to drop and keeps the rule
|
|
# idempotent (once the input yields exactly ``required`` nothing more is
|
|
# inserted, so the fixed-point optimizer terminates).
|
|
if not set(required) < set(schema.names):
|
|
return op
|
|
|
|
prune = Project(exprs=[col(c) for c in required], input_dependencies=[input_op])
|
|
return replace(op, input_dependencies=[prune])
|
|
|
|
@classmethod
|
|
def _try_fuse_projects(cls, op: LogicalOperator) -> LogicalOperator:
|
|
"""
|
|
Optimize a single Project operator.
|
|
|
|
Steps:
|
|
1. Iteratively fuse with upstream Project operations
|
|
2. Push the resulting projection into the data source if possible
|
|
"""
|
|
if not isinstance(op, Project):
|
|
return op
|
|
|
|
# Step 1: Iteratively fuse with upstream Project operations
|
|
current_project: Project = op
|
|
|
|
upstream_op = current_project.input_dependencies[0]
|
|
if not isinstance(upstream_op, Project):
|
|
return op
|
|
|
|
fused = _try_fuse(upstream_op, current_project)
|
|
|
|
return fused
|
|
|
|
@classmethod
|
|
def _push_projection_into_read_op(cls, op: LogicalOperator) -> LogicalOperator:
|
|
|
|
if not isinstance(op, Project):
|
|
return op
|
|
|
|
current_project: Project = op
|
|
|
|
# Step 2: Push projection into the data source if supported
|
|
input_op = current_project.input_dependencies[0]
|
|
if (
|
|
isinstance(input_op, LogicalOperatorSupportsProjectionPushdown)
|
|
and input_op.supports_projection_pushdown()
|
|
):
|
|
# Collect the set of input columns this ``Project`` reads.
|
|
# ``None`` means "all columns" (a ``StarExpr`` is present, so
|
|
# we can't enumerate). Renames are NOT pushed into the read:
|
|
# column renaming always stays as an ``AliasExpr`` in a
|
|
# ``Project`` on top of the read. This keeps the read stage
|
|
# in a single column namespace (the scanner's on-disk names)
|
|
# and avoids the predicate-pushdown rebinding dance.
|
|
if current_project.has_star_expr():
|
|
# UDF-fallback path: if ``StarExpr`` survives to this rule
|
|
# the input schema was unknown at construction time, so
|
|
# we can't enumerate columns and have to push "all".
|
|
required_columns = None
|
|
else:
|
|
# Otherwise, collect required columns to push projection down
|
|
# into the reader. (``Project.__post_init__`` expands
|
|
# ``StarExpr`` to explicit ``col()`` refs when the input
|
|
# schema is known, so this is the common path.)
|
|
required_columns = _collect_referenced_columns(current_project.exprs)
|
|
|
|
# Build a pure-prune projection map (identity, no renames).
|
|
projection_map = (
|
|
None
|
|
if required_columns is None
|
|
else {name: name for name in required_columns}
|
|
)
|
|
projected_input_op = input_op.apply_projection(projection_map)
|
|
|
|
# If the ``Project`` is a pure-prune (only ``col()`` refs,
|
|
# no renames, no computed expressions), the projection
|
|
# pushdown into the read op fully subsumes it — discard it.
|
|
# Otherwise (renames or computed expressions present), keep
|
|
# the ``Project`` on top so it runs above the (pruned) read.
|
|
# Physical operator fusion later merges the kept ``Project``
|
|
# into the same ``MapOperator`` as the read, so the
|
|
# runtime cost is the same either way.
|
|
has_renames = any(isinstance(e, AliasExpr) for e in current_project.exprs)
|
|
all_col_refs = all(
|
|
_is_col_expr(e) for e in _filter_out_star(current_project.exprs)
|
|
)
|
|
is_pure_prune = not has_renames and all_col_refs
|
|
if is_pure_prune:
|
|
return projected_input_op
|
|
|
|
return Project(
|
|
exprs=current_project.exprs,
|
|
input_dependencies=[projected_input_op],
|
|
compute=current_project.compute,
|
|
ray_remote_args=current_project.ray_remote_args,
|
|
)
|
|
|
|
return current_project
|
|
|
|
|
|
def _extract_input_columns_renaming_mapping(
|
|
projection_exprs: List[Expr],
|
|
) -> Dict[str, str]:
|
|
"""Fetches renaming mapping of all input columns names being renamed (replaced).
|
|
Format is source column name -> new column name.
|
|
"""
|
|
|
|
return dict(
|
|
[
|
|
_get_renaming_mapping(expr)
|
|
for expr in _filter_out_star(projection_exprs)
|
|
if is_rename_expr(expr)
|
|
]
|
|
)
|
|
|
|
|
|
def _get_renaming_mapping(expr: Expr) -> Tuple[str, str]:
|
|
assert is_rename_expr(expr)
|
|
|
|
alias: AliasExpr = expr
|
|
|
|
return alias.expr.name, alias.name
|