cool9850311 commented on code in PR #27152:
URL: https://github.com/apache/datafusion/pull/27152#discussion_r4504820774
##########
datafusion/physical-optimizer/src/projection_pushdown.rs:
##########
@@ -434,21 +451,234 @@ fn remap_filter_expr(
} else if let Some(variable) = expr.downcast_ref::<LambdaVariable>()
&& old_input_len != new_input_len
{
- let offset =
variable.index().checked_sub(old_input_len).ok_or_else(|| {
- internal_datafusion_err!(
- "Lambda parameter {variable} precedes the {old_input_len}
input columns"
- )
- })?;
- return Ok(Transformed::yes(Arc::new(LambdaVariable::new(
- new_input_len + offset,
- Arc::clone(variable.field()),
- )) as _));
+ let rebased = rebase_lambda_variable(variable, old_input_len,
new_input_len)?;
+ return Ok(Transformed::yes(rebased));
}
Ok(Transformed::no(expr))
})
.data()
}
+/// Shifts a lambda parameter by the change in input width, for the reason
+/// [`remap_filter_expr`] documents.
+fn rebase_lambda_variable(
+ variable: &LambdaVariable,
+ old_input_len: usize,
+ new_input_len: usize,
+) -> Result<Arc<dyn PhysicalExpr>> {
+ let offset = variable.index().checked_sub(old_input_len).ok_or_else(|| {
+ internal_datafusion_err!(
+ "Lambda parameter {variable} precedes the {old_input_len} input
columns"
+ )
+ })?;
+ Ok(Arc::new(LambdaVariable::new(
+ new_input_len + offset,
+ Arc::clone(variable.field()),
+ )))
+}
+
+/// 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 {
Review Comment:
Done — your SLT block went in verbatim at the position you named, and the
expected plan matched exactly, including `partition_sizes=[3]`.
I measured how far "no test fails" reached, since it decides what the green
suite was worth. With the check removed:
| suite | result |
|---|---|
| sqllogictest, 546 files | `RC=0` |
| physical-optimizer unit, 50 | `RC=0` |
| `core_integration` `physical_optimizer`, 668 | `RC=0` |
Nothing in 1264 tests noticed, including the four this PR had added — all of
mine were positive tests, none checked that the split *doesn't* happen.
Running the same check over the rest of the function found a second unpinned
promise, so de-duplication has an `EXPLAIN` now too. Values can't show it
(two
calls to one marker agree either way) but the scan projection can:
```
projection=[column1, CAST(__datafusion_file_row_index@1 AS Int64;N) as
__datafusion_leaf_0]
```
Without de-duplication that line gains a second `CAST(...) as
__datafusion_leaf_1`,
which is what the new test fails on.
##########
datafusion/physical-optimizer/src/projection_pushdown.rs:
##########
@@ -61,6 +66,18 @@ impl PhysicalOptimizerRule for ProjectionPushdown {
})
.map(|t| t.data)?;
+ // Push what can already descend, then split the projections that are
+ // stuck above a `RepartitionExec`, then clean up once more: the split
+ // can leave an outer projection that is now only column references.
+ let plan = plan.transform_down(remove_unnecessary_projections).data()?;
+ let plan = plan
+ .transform_down(|plan| {
+ Ok(match try_split_leaf_projection(&plan)? {
+ Some(new_plan) => Transformed::yes(new_plan),
+ None => Transformed::no(plan),
+ })
+ })
+ .data()?;
plan.transform_down(remove_unnecessary_projections).data()
Review Comment:
I went a step further than this and removed the walk entirely — please tell
me
if you'd rather have your version. What pushed me that way is that it seems
not
to do anything even when a split does fire.
Instrumented to diff the plan string before and after it:
| suite | splits fired | walk changed a plan |
|---|---|---|
| sqllogictest, 545 files | 19 | 0 |
| unit | 0 | 0 |
| `core_integration` | 0 | 0 |
Both directions of the mutation are green too — never running it, and
running it
unconditionally — so no test observes it either way.
<details><summary>Why it looks like it can't fire</summary>
`remove_unnecessary_projections` has three paths for a projection
(`projection.rs:1047`), and as far as I can tell the outer projection takes
none:
1. `is_projection_removable` wants alias == column name and
`projection.schema() == input.schema()`. The split renames (`v` over
`__datafusion_leaf_0`) and keeps the original output schema, for the
metadata
reason you raised on the other thread. Instrumented: both false in 20/20
splits.
2. `overrides_metadata` returns `Transformed::no`.
3. `RepartitionExec::try_swapping_with_projection` declines either way. If
every
outer expression stayed a plain column, each original expression was
exactly
one leaf or exactly one passthrough, so `N >= P + L` and the narrowing
guard
refuses. If it does narrow, some expression consumed more than one leaf
and so
isn't a plain `Column`, and `all_columns` refuses.
Measured per split (`N` = outer expressions, `P` = passthrough, `L` =
leaves):
| N | P | L | count | refused by |
|---|---|---|---|---|
| 1 | 0 | 1 | 2 | narrowing guard |
| 2 | 1 | 1 | 8 | narrowing guard |
| 2 | 1 | 2 | 1 | `all_columns` |
| 3 | 1 | 1 | 3 | narrowing guard |
| 3 | 1 | 2 | 4 | narrowing guard |
| 3 | 1 | 3 | 1 | narrowing guard |
| 4 | 2 | 2 | 1 | narrowing guard |
18 of 20 on the first guard, 2 on the second — the narrowing case is a nested
lambda whose single output expression consumes two leaves.
Nothing else in the tree moves either: the split preserves the projection's
schema, so ancestors see no change, and the lower projection is absorbed by
the
source.
</details>
##########
datafusion/sqllogictest/test_files/file_row_index.slt:
##########
@@ -151,6 +151,145 @@ ORDER BY column1
statement ok
DROP TABLE parquet_two_files;
+# A leaf expression must still reach the scan when its projection also holds a
+# `KeepInPlace` one; see `try_split_leaf_projection` in
+# `datafusion/physical-optimizer/src/projection_pushdown.rs`. These need a
+# repartition, so `batch_size` is below the row count: do not raise it and do
+# not add a bare `LIMIT`, or the queries pass against broken code.
+
+statement ok
+SET datafusion.execution.batch_size = 2;
+
+# The marker wrapped in arithmetic.
+query II
+SELECT file_row_index() + 1, column1 FROM parquet_table ORDER BY column1
+----
+1 10
+2 20
+3 30
+4 40
+5 50
+
+# The marker is BARE here. An unrelated computed column in the same projection
+# is enough, because the flag is a property of the whole projection.
+query II
+SELECT file_row_index(), column1 + 1 FROM parquet_table ORDER BY column1
+----
+0 11
+1 21
+2 31
+3 41
+4 51
+
+# Not specific to arithmetic: any enclosing expression has the same effect.
+query TI
+SELECT arrow_typeof(file_row_index()), column1 FROM parquet_table ORDER BY
column1
+----
+Int64 10
+Int64 20
+Int64 30
+Int64 40
+Int64 50
+
+# The sibling marker behaves the same way. Compared against a predicate rather
+# than printed, so the expectation does not depend on the absolute file path.
+query BI
+SELECT length(input_file_name()) > 0, column1 FROM parquet_table ORDER BY
column1
+----
+true 10
+true 20
+true 30
+true 40
+true 50
+
+# Both markers in one projection, each inside an enclosing expression.
+query BI
+SELECT input_file_name() IS NOT NULL AND file_row_index() >= 0, column1
+FROM parquet_table ORDER BY column1
+----
+true 10
+true 20
+true 30
+true 40
+true 50
+
+# A bare `input_file_name()` beside a computed column: the same trigger as the
+# `file_row_index()` case above. Asserted through the plan rather than the
+# values, because the value is an absolute path.
+query TT
+EXPLAIN SELECT input_file_name(), column1 + 1 FROM parquet_table
+----
+logical_plan
+01)Projection: input_file_name(), CAST(parquet_table.column1 AS Int64) +
Int64(1)
+02)--TableScan: parquet_table projection=[column1]
+physical_plan
+01)ProjectionExec: expr=[__datafusion_leaf_0@1 as input_file_name(),
CAST(column1@0 AS Int64) + 1 as parquet_table.column1 + Int64(1)]
+02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+03)----DataSourceExec: file_groups={1 group:
[[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/file_row_index/parquet_table/data.parquet]]},
projection=[column1, input_file_name() as __datafusion_leaf_0],
file_type=parquet
+
+# The same marker twice. Both calls are one expression, so the split gives them
Review Comment:
All three are in. They match your expected values and all three error against
`main` with `is source dependent and cannot be evaluated directly` — TopK
gives
`4 51` / `3 41`, the lambda-body marker `[1, 2, 3] 100 alpha` …, CSV `true
11`
… `true 61`.
One thing I couldn't reproduce: the dynamic filter reaching the scan in the
TopK
query. With `ORDER BY r DESC LIMIT 2`, where `r` is the rewritten marker:
```
SortExec: TopK(fetch=2), expr=[r@0 DESC], preserve_partitioning=[true]
ProjectionExec: expr=[__datafusion_leaf_0@1 as r, column1@0 + 1 as ...]
RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
DataSourceExec: ..., projection=[column1, CAST(...) as
__datafusion_leaf_0], file_type=parquet
```
No predicate on the scan. Sorting on the computed sibling instead (`ORDER BY
w`)
does give `predicate=DynamicFilter` and `dynamic_rg_pruning=eligible` — my
guess
is `__datafusion_leaf_0` is a synthetic row-index column with no statistics
to
prune row groups with, but I haven't chased it.
The test still seems worth keeping for the path it does cover: the sort key
is a
column the split produced, and the TopK runs above the repartition with
`preserve_partitioning=[true]`. So I kept it and rewrote the comment to say
that
instead. I'd put your wording in without checking it first, which was my slip
rather than anything wrong with the test you suggested.
##########
datafusion/sqllogictest/test_files/file_row_index.slt:
##########
@@ -151,6 +151,145 @@ ORDER BY column1
statement ok
DROP TABLE parquet_two_files;
+# A leaf expression must still reach the scan when its projection also holds a
+# `KeepInPlace` one; see `try_split_leaf_projection` in
+# `datafusion/physical-optimizer/src/projection_pushdown.rs`. These need a
+# repartition, so `batch_size` is below the row count: do not raise it and do
+# not add a bare `LIMIT`, or the queries pass against broken code.
+
+statement ok
+SET datafusion.execution.batch_size = 2;
+
+# The marker wrapped in arithmetic.
+query II
+SELECT file_row_index() + 1, column1 FROM parquet_table ORDER BY column1
+----
+1 10
+2 20
+3 30
+4 40
+5 50
+
+# The marker is BARE here. An unrelated computed column in the same projection
+# is enough, because the flag is a property of the whole projection.
+query II
+SELECT file_row_index(), column1 + 1 FROM parquet_table ORDER BY column1
+----
+0 11
+1 21
+2 31
+3 41
+4 51
+
+# Not specific to arithmetic: any enclosing expression has the same effect.
+query TI
+SELECT arrow_typeof(file_row_index()), column1 FROM parquet_table ORDER BY
column1
+----
+Int64 10
+Int64 20
+Int64 30
+Int64 40
+Int64 50
+
+# The sibling marker behaves the same way. Compared against a predicate rather
+# than printed, so the expectation does not depend on the absolute file path.
+query BI
+SELECT length(input_file_name()) > 0, column1 FROM parquet_table ORDER BY
column1
+----
+true 10
+true 20
+true 30
+true 40
+true 50
+
+# Both markers in one projection, each inside an enclosing expression.
+query BI
+SELECT input_file_name() IS NOT NULL AND file_row_index() >= 0, column1
+FROM parquet_table ORDER BY column1
+----
+true 10
+true 20
+true 30
+true 40
+true 50
+
+# A bare `input_file_name()` beside a computed column: the same trigger as the
+# `file_row_index()` case above. Asserted through the plan rather than the
+# values, because the value is an absolute path.
+query TT
+EXPLAIN SELECT input_file_name(), column1 + 1 FROM parquet_table
+----
+logical_plan
+01)Projection: input_file_name(), CAST(parquet_table.column1 AS Int64) +
Int64(1)
+02)--TableScan: parquet_table projection=[column1]
+physical_plan
+01)ProjectionExec: expr=[__datafusion_leaf_0@1 as input_file_name(),
CAST(column1@0 AS Int64) + 1 as parquet_table.column1 + Int64(1)]
+02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+03)----DataSourceExec: file_groups={1 group:
[[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/file_row_index/parquet_table/data.parquet]]},
projection=[column1, input_file_name() as __datafusion_leaf_0],
file_type=parquet
+
+# The same marker twice. Both calls are one expression, so the split gives them
+# a single column and the marker is evaluated once per row, not twice. Both
+# outputs must still agree with each other.
+query III
+SELECT file_row_index(), file_row_index() + 1, column1
+FROM parquet_table ORDER BY column1
+----
+0 1 10
+1 2 20
+2 3 30
+3 4 40
+4 5 50
+
+# Shapes that already worked must keep working: the leaf rules extract from
+# `Filter`, `Sort` and `Aggregate`, so those never depended on this fix.
+query I
+SELECT column1 FROM parquet_table WHERE file_row_index() + 1 = 3
+----
+30
+
+query I
+SELECT column1 FROM parquet_table ORDER BY file_row_index() + 1
+----
+10
+20
+30
+40
+50
+
+query I
+SELECT sum(file_row_index() + 1) FROM parquet_table
+----
+15
+
+# A bare marker on its own is unaffected: no `KeepInPlace` expression shares
+# the projection, so nothing is inserted beneath it.
+query II
+SELECT file_row_index(), column1 FROM parquet_table ORDER BY column1
+----
+0 10
+1 20
+2 30
+3 40
+4 50
+
+# Raising `batch_size` above the row count removes the repartition. This must
+# keep working, and it is also the control that proves trap 1 above is real.
Review Comment:
Applied verbatim — thanks, that one was mine to catch. "trap 1" was a number
I
made up in the PR description and never defined in the file.
--
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]
