jayzhan211 commented on code in PR #27152:
URL: https://github.com/apache/datafusion/pull/27152#discussion_r4492591735
##########
datafusion/physical-optimizer/src/projection_pushdown.rs:
##########
@@ -449,6 +466,191 @@ fn remap_filter_expr(
.data()
}
+/// Prefix for the columns this rule adds when it splits a projection.
+const LEAF_COLUMN_PREFIX: &str = "__datafusion_leaf";
+
+/// Name of the column the split adds for the leaf expression at `index`.
+///
+/// Single source of truth: the lower projection's alias and the reference the
+/// outer projection is rewritten to must agree, or the rewrite silently points
+/// at nothing.
+fn leaf_column_name(index: usize) -> String {
+ format!("{LEAF_COLUMN_PREFIX}_{index}")
+}
+
+/// Splits a `ProjectionExec` that mixes leaf expressions with ones that must
+/// stay in place, moving the leaves into a projection the source absorbs.
+///
+/// `benefits_from_input_partitioning` is true as soon as any expression is
+/// `KeepInPlace`, which both inserts a `RepartitionExec` below the projection
+/// and stops that projection descending past it, so a `MoveTowardsLeafNodes`
+/// expression sharing it never reaches the source. Splitting leaves the
+/// computation above the repartition, still parallelised.
+fn try_split_leaf_projection(
+ plan: &Arc<dyn ExecutionPlan>,
+) -> Result<Option<Arc<dyn ExecutionPlan>>> {
+ let Some(projection) = plan.downcast_ref::<ProjectionExec>() else {
+ return Ok(None);
+ };
+ let Some(repartition) =
projection.input().downcast_ref::<RepartitionExec>() else {
+ return Ok(None);
+ };
+ // Hash and Range partitioning reference columns by index, which would have
+ // to be remapped against the rewritten child schema. Round robin carries
no
+ // column references, and is what `EnsureRequirements` inserts here.
+ if !matches!(repartition.partitioning(), Partitioning::RoundRobinBatch(_))
{
+ return Ok(None);
+ }
+
+ let leaves = collect_leaf_exprs(projection.expr())?;
+ if leaves.is_empty() {
+ return Ok(None);
+ }
+
+ // Pass through only the child columns the outer expressions still need. A
+ // column used *only* inside an extracted leaf must be dropped, otherwise
the
+ // source keeps reading it: for `s['a'] + 1` that would read the whole
struct
+ // alongside the extracted field, which is worse than not splitting at all.
+ let child = repartition.input();
+ let child_schema = child.schema();
+ let passthrough_indices = collect_outer_columns(projection.expr(),
&leaves)?;
+ let mut lower_exprs = passthrough_indices
+ .iter()
+ .map(|&idx| {
+ let field = child_schema.field(idx);
+ ProjectionExpr {
+ expr: Arc::new(Column::new(field.name(), idx)),
+ alias: field.name().to_string(),
+ }
+ })
+ .collect::<Vec<_>>();
+ for (index, leaf) in leaves.iter().enumerate() {
+ lower_exprs.push(ProjectionExpr {
+ expr: Arc::clone(leaf),
+ alias: leaf_column_name(index),
+ });
+ }
+ let lower_projection = ProjectionExec::try_new(lower_exprs,
Arc::clone(child))?;
+
+ // Only commit the split if the child really absorbs the leaf projection.
+ // Otherwise this would just add a projection that helps nobody.
+ let Some(new_child) =
child.try_swapping_with_projection(&lower_projection)? else {
+ return Ok(None);
+ };
+
+ let new_exprs = projection
+ .expr()
+ .iter()
+ .map(|proj_expr| {
+ Ok(ProjectionExpr {
+ expr: rewrite_outer_expr(
+ Arc::clone(&proj_expr.expr),
+ &leaves,
+ &passthrough_indices,
+ )?,
+ alias: proj_expr.alias.clone(),
+ })
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ // The child gained the leaf columns, so the repartition's properties have
to
+ // be recomputed rather than carried over.
+ let new_repartition = Arc::clone(projection.input()).replace_children(
+ vec![new_child],
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )?;
+ Ok(Some(Arc::new(ProjectionExec::try_new(
+ new_exprs,
+ new_repartition,
+ )?)))
+}
+
+/// Collects the distinct outermost
[`ExpressionPlacement::MoveTowardsLeafNodes`]
+/// sub-expressions of `exprs`.
+///
+/// Nested leaf expressions are not collected separately: `s['a']['b']` is one
+/// unit to hand to the source, not two.
+fn collect_leaf_exprs(exprs: &[ProjectionExpr]) -> Result<Vec<Arc<dyn
PhysicalExpr>>> {
+ let mut leaves: Vec<Arc<dyn PhysicalExpr>> = vec![];
+ for proj_expr in exprs {
+ proj_expr.expr.apply(|node| {
+ if node.placement() == ExpressionPlacement::MoveTowardsLeafNodes {
+ if !leaves.iter().any(|leaf| leaf == node) {
+ leaves.push(Arc::clone(node));
+ }
+ return Ok(TreeNodeRecursion::Jump);
+ }
+ Ok(TreeNodeRecursion::Continue)
+ })?;
+ }
+ Ok(leaves)
+}
+
+/// Collects the child column indices that `exprs` still reference *outside*
the
+/// extracted leaf sub-expressions, sorted and deduplicated.
+///
+/// Recursion stops at a leaf: columns used only inside one are supplied by the
+/// lower projection itself and must not be passed through.
+fn collect_outer_columns(
+ exprs: &[ProjectionExpr],
+ leaves: &[Arc<dyn PhysicalExpr>],
+) -> Result<Vec<usize>> {
+ let mut indices: Vec<usize> = vec![];
+ for proj_expr in exprs {
+ proj_expr.expr.apply(|node| {
+ if leaves.iter().any(|leaf| leaf == node) {
+ return Ok(TreeNodeRecursion::Jump);
+ }
+ if let Some(column) = node.downcast_ref::<Column>() {
+ indices.push(column.index());
+ }
+ Ok(TreeNodeRecursion::Continue)
+ })?;
+ }
+ indices.sort_unstable();
+ indices.dedup();
+ Ok(indices)
+}
+
+/// Rewrites an outer projection expression against the split child schema:
+/// every leaf sub-expression becomes a reference to the column the split added
+/// for it, and every surviving column is remapped to its new index.
+fn rewrite_outer_expr(
+ expr: Arc<dyn PhysicalExpr>,
+ leaves: &[Arc<dyn PhysicalExpr>],
+ passthrough_indices: &[usize],
+) -> Result<Arc<dyn PhysicalExpr>> {
+ // The leaf columns were appended after the passthrough ones.
+ let first_leaf = passthrough_indices.len();
+ expr.transform_down(|node| {
+ if let Some(pos) = leaves.iter().position(|leaf| leaf == &node) {
+ let column = Column::new(&leaf_column_name(pos), first_leaf + pos);
+ return Ok(Transformed::new(
+ Arc::new(column) as Arc<dyn PhysicalExpr>,
+ true,
+ TreeNodeRecursion::Jump,
+ ));
+ }
+ if let Some(column) = node.downcast_ref::<Column>() {
+ let new_index = passthrough_indices
+ .iter()
+ .position(|&idx| idx == column.index())
+ .ok_or_else(|| {
+ internal_datafusion_err!(
+ "projection split dropped column {} that is still
referenced",
+ column.index()
+ )
+ })?;
+ if new_index != column.index() {
+ let remapped = Column::new(column.name(), new_index);
+ return Ok(Transformed::yes(Arc::new(remapped) as Arc<dyn
PhysicalExpr>));
+ }
+ }
+ Ok(Transformed::no(node))
Review Comment:
Lambda parameters are indexed after the input columns, so they go stale when
the split widens the input. This query works on `main` but now fails with
`Field of physical LambdaVariable with index 0 doesn't match batch field`:
```sql
SET datafusion.sql_parser.dialect = databricks;
SELECT array_transform(arr, x -> x + 1) AS a, s['value'] AS v, s['label'] AS
l FROM t ORDER BY v;
```
This fixes it, with `child_schema.fields().len()` passed in as a new
`old_input_len` parameter and `let new_input_len = first_leaf + leaves.len();`
added next to `first_leaf`. Could you add the query above to the SLT, plus a
nested lambda case?
```suggestion
// Lambda parameters are appended after the input columns, so they
move
// when the split changes the input width.
if let Some(variable) = node.downcast_ref::<LambdaVariable>()
&& old_input_len != new_input_len
{
let rebased = LambdaVariable::new(
variable.index() - old_input_len + new_input_len,
Arc::clone(variable.field()),
);
return Ok(Transformed::yes(Arc::new(rebased) as Arc<dyn
PhysicalExpr>));
}
Ok(Transformed::no(node))
```
##########
datafusion/physical-optimizer/src/projection_pushdown.rs:
##########
@@ -449,6 +466,191 @@ fn remap_filter_expr(
.data()
}
+/// Prefix for the columns this rule adds when it splits a projection.
+const LEAF_COLUMN_PREFIX: &str = "__datafusion_leaf";
+
+/// Name of the column the split adds for the leaf expression at `index`.
+///
+/// Single source of truth: the lower projection's alias and the reference the
+/// outer projection is rewritten to must agree, or the rewrite silently points
+/// at nothing.
+fn leaf_column_name(index: usize) -> String {
+ format!("{LEAF_COLUMN_PREFIX}_{index}")
+}
+
+/// Splits a `ProjectionExec` that mixes leaf expressions with ones that must
+/// stay in place, moving the leaves into a projection the source absorbs.
+///
+/// `benefits_from_input_partitioning` is true as soon as any expression is
+/// `KeepInPlace`, which both inserts a `RepartitionExec` below the projection
+/// and stops that projection descending past it, so a `MoveTowardsLeafNodes`
+/// expression sharing it never reaches the source. Splitting leaves the
+/// computation above the repartition, still parallelised.
+fn try_split_leaf_projection(
+ plan: &Arc<dyn ExecutionPlan>,
+) -> Result<Option<Arc<dyn ExecutionPlan>>> {
+ let Some(projection) = plan.downcast_ref::<ProjectionExec>() else {
+ return Ok(None);
+ };
+ let Some(repartition) =
projection.input().downcast_ref::<RepartitionExec>() else {
+ return Ok(None);
+ };
+ // Hash and Range partitioning reference columns by index, which would have
+ // to be remapped against the rewritten child schema. Round robin carries
no
+ // column references, and is what `EnsureRequirements` inserts here.
+ if !matches!(repartition.partitioning(), Partitioning::RoundRobinBatch(_))
{
+ return Ok(None);
+ }
+
+ let leaves = collect_leaf_exprs(projection.expr())?;
+ if leaves.is_empty() {
+ return Ok(None);
+ }
+
+ // Pass through only the child columns the outer expressions still need. A
+ // column used *only* inside an extracted leaf must be dropped, otherwise
the
+ // source keeps reading it: for `s['a'] + 1` that would read the whole
struct
+ // alongside the extracted field, which is worse than not splitting at all.
+ let child = repartition.input();
+ let child_schema = child.schema();
+ let passthrough_indices = collect_outer_columns(projection.expr(),
&leaves)?;
+ let mut lower_exprs = passthrough_indices
+ .iter()
+ .map(|&idx| {
+ let field = child_schema.field(idx);
+ ProjectionExpr {
+ expr: Arc::new(Column::new(field.name(), idx)),
+ alias: field.name().to_string(),
+ }
+ })
+ .collect::<Vec<_>>();
+ for (index, leaf) in leaves.iter().enumerate() {
+ lower_exprs.push(ProjectionExpr {
+ expr: Arc::clone(leaf),
+ alias: leaf_column_name(index),
+ });
+ }
+ let lower_projection = ProjectionExec::try_new(lower_exprs,
Arc::clone(child))?;
+
+ // Only commit the split if the child really absorbs the leaf projection.
+ // Otherwise this would just add a projection that helps nobody.
+ let Some(new_child) =
child.try_swapping_with_projection(&lower_projection)? else {
+ return Ok(None);
+ };
+
+ let new_exprs = projection
+ .expr()
+ .iter()
+ .map(|proj_expr| {
+ Ok(ProjectionExpr {
+ expr: rewrite_outer_expr(
+ Arc::clone(&proj_expr.expr),
+ &leaves,
+ &passthrough_indices,
+ )?,
+ alias: proj_expr.alias.clone(),
+ })
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ // The child gained the leaf columns, so the repartition's properties have
to
+ // be recomputed rather than carried over.
+ let new_repartition = Arc::clone(projection.input()).replace_children(
+ vec![new_child],
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )?;
+ Ok(Some(Arc::new(ProjectionExec::try_new(
+ new_exprs,
+ new_repartition,
+ )?)))
Review Comment:
`try_new` drops metadata that the planner set through
`try_new_with_schema_metadata`, so this select now fails with
`PhysicalOptimizer rule 'ProjectionPushdown' failed. Schema mismatch`, while it
works on `main`:
```rs
ctx.table("t").await?.select(vec![
get_field(col("s"), "value").alias_with_metadata("v",
Some(metadata.into())),
(col("id") + lit(1i64)).alias("w"),
])?;
```
Could you apply this and add that select as a test, with `batch_size` below
the row count so the repartition appears, asserting the output field keeps its
metadata?
```suggestion
Ok(Some(Arc::new(
ProjectionExec::try_new_with_schema_metadata(
new_exprs,
new_repartition,
projection.schema().as_ref(),
)?,
)))
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
