jayzhan211 commented on code in PR #27152:
URL: https://github.com/apache/datafusion/pull/27152#discussion_r4501784996
##########
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:
No test fails if this check is removed. `SELECT s['value'] AS v, id + 1 AS
w` on a memory table then gets an extra `ProjectionExec` that the source never
absorbs, and every test in this PR still passes. The PR description promises
this limit, so could you pin it here with an `EXPLAIN` in
`projection_pushdown.slt`?
<details>
<summary>SLT test, placed before the <code>batch_size = 8192</code> restore
(passes on this PR, fails with the check removed)</summary>
```sql
# A source that does not absorb the projection keeps the plan unchanged: the
# split must not add a projection nobody can use.
statement ok
CREATE TABLE leaf_split_mem AS
SELECT column1 AS id, column2 AS s FROM VALUES
(1, {value: 100}), (2, {value: 200}), (3, {value: 150}),
(4, {value: 300}), (5, {value: 250});
query TT
EXPLAIN SELECT s['value'] AS v, id + 1 AS w FROM leaf_split_mem;
----
logical_plan
01)Projection: get_field(leaf_split_mem.s, Utf8("value")) AS v,
leaf_split_mem.id + Int64(1) AS w
02)--TableScan: leaf_split_mem projection=[id, s]
physical_plan
01)ProjectionExec: expr=[get_field(s@1, value) as v, id@0 + 1 as w]
02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
03)----DataSourceExec: partitions=1, partition_sizes=[3]
statement ok
DROP TABLE leaf_split_mem;
```
</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:
Optional: three shapes fail on `main` and work with this PR, but have no
test yet: a TopK whose dynamic filter reaches the scan through the rewritten
marker column, `input_file_name()` on a CSV source, and a marker inside a
lambda body. The queries below all pass on this PR.
<details>
<summary>The three SLT tests and which file each goes in</summary>
In `file_row_index.slt`, inside the `batch_size = 2` block:
```sql
query II
SELECT file_row_index() AS r, column1 + 1 FROM parquet_table ORDER BY r DESC
LIMIT 2
----
4 51
3 41
```
In `projection_pushdown.slt`, inside the `databricks` dialect block:
```sql
query ?IT
SELECT array_transform(arr, x -> x + file_row_index()) AS a, s['value'] AS
v, s['label'] AS l
FROM leaf_split ORDER BY v;
----
[1, 2, 3] 100 alpha
[8] 150 gamma
[5, 6] 200 beta
[14] 250 epsilon
[10, 11, 12] 300 delta
```
In `input_file_name.slt`, while `csv_table` exists:
```sql
statement ok
SET datafusion.execution.batch_size = 2;
query BI
SELECT length(input_file_name()) > 0, column1 + 1
FROM csv_table
ORDER BY 2
----
true 11
true 21
true 31
true 41
true 51
true 61
statement ok
SET datafusion.execution.batch_size = 8192;
```
</details>
##########
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:
Optional cleanup: the last walk only matters when a split fired. Skipping it
otherwise means a plan with no split gets the same single cleanup walk as on
`main`. All 545 SLT files pass with this change.
```suggestion
let split = plan.transform_down(|plan| {
Ok(match try_split_leaf_projection(&plan)? {
Some(new_plan) => Transformed::yes(new_plan),
None => Transformed::no(plan),
})
})?;
if split.transformed {
split
.data
.transform_down(remove_unnecessary_projections)
.data()
} else {
Ok(split.data)
}
```
##########
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:
Optional: "trap 1" is only defined in the PR description, so this comment
won't make sense to someone reading the file.
```suggestion
# Raising `batch_size` above the row count removes the repartition. This must
# keep working, and it is also the control showing that `batch_size` is what
# triggers the repartition.
```
--
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]
