From aba5ea095e76f23e193190caa856896048068b64 Mon Sep 17 00:00:00 2001 From: Harry Shih Date: Thu, 5 Nov 2026 22:04:00 +0800 Subject: [PATCH 1/3] fix: let leaf expressions reach the scan past a RepartitionExec `ProjectionExec::benefits_from_input_partitioning` is true as soon as any expression in the projection has `KeepInPlace` placement. That single flag does two opposing things: `EnsureRequirements` inserts a `RepartitionExec` beneath the projection, and `RepartitionExec::try_swapping_with_projection` then refuses to let the same projection descend. A `MoveTowardsLeafNodes` expression sharing that projection can therefore never reach the source. `file_row_index()` and `input_file_name()` are then evaluated directly and fail with "is source dependent and cannot be evaluated directly", while `get_field` silently reads the whole struct instead of pruning to the field. The leaf call does not have to be wrapped in anything: an unrelated computed column in the same projection is enough, because the flag is a property of the whole projection. `ProjectionPushdown` now splits such a projection. The leaf expressions move into a lower projection that the source absorbs, and the computation stays above the repartition where it is still parallelised, so the intent of both guards is kept rather than one traded for the other. The split is committed only when the child really absorbs the leaf projection, and only `RoundRobinBatch` repartitioning is handled, because hash and range partitioning reference columns by index and would need remapping. Only the columns the outer expressions still reference outside the extracted leaves are passed through, so `s['a'] + 1` no longer reads the whole struct alongside the extracted field. Regression tests cover both symptom classes. Two traps are documented in the test files because either one makes the tests pass against broken code: the table must exceed `batch_size` for a repartition to be inserted at all, and a bare `LIMIT n` is pushed into the scan and hides the bug again. Closes #27131 Co-Authored-By: Claude Opus 5 --- .../src/projection_pushdown.rs | 204 +++++++++++++++++- .../test_files/file_row_index.slt | 139 ++++++++++++ .../test_files/projection_pushdown.slt | 115 ++++++++++ 3 files changed, 457 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-optimizer/src/projection_pushdown.rs b/datafusion/physical-optimizer/src/projection_pushdown.rs index 7faa72c1d947e..a45975b027012 100644 --- a/datafusion/physical-optimizer/src/projection_pushdown.rs +++ b/datafusion/physical-optimizer/src/projection_pushdown.rs @@ -18,14 +18,19 @@ use datafusion_common::tree_node::{ Transformed, TransformedResult, TreeNode, TreeNodeRecursion, }; use datafusion_common::{JoinSide, JoinType, Result, internal_datafusion_err}; +use datafusion_expr::ExpressionPlacement; use datafusion_physical_expr::expressions::{Column, LambdaExpr, LambdaVariable}; +use datafusion_physical_expr::projection::ProjectionExpr; use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, is_volatile}; -use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::joins::NestedLoopJoinExec; use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter}; use datafusion_physical_plan::projection::{ ProjectionExec, remove_unnecessary_projections, }; +use datafusion_physical_plan::repartition::RepartitionExec; +use datafusion_physical_plan::{ + ChildrenPropertiesMode, ExecutionPlan, Partitioning, ReplaceChildrenOptions, +}; /// This rule inspects `ProjectionExec`'s in the given physical plan and tries to /// remove or swap with its child. @@ -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() } @@ -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, +) -> Result>> { + let Some(projection) = plan.downcast_ref::() else { + return Ok(None); + }; + let Some(repartition) = projection.input().downcast_ref::() 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::>(); + 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::>>()?; + + // 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>> { + let mut leaves: Vec> = 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], +) -> Result> { + let mut indices: Vec = 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::() { + 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, + leaves: &[Arc], + passthrough_indices: &[usize], +) -> Result> { + // 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, + true, + TreeNodeRecursion::Jump, + )); + } + if let Some(column) = node.downcast_ref::() { + 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)); + } + } + Ok(Transformed::no(node)) + }) + .data() +} + #[cfg(test)] mod test { use super::*; diff --git a/datafusion/sqllogictest/test_files/file_row_index.slt b/datafusion/sqllogictest/test_files/file_row_index.slt index 935c4d6c72c01..7e135731e2c0b 100644 --- a/datafusion/sqllogictest/test_files/file_row_index.slt +++ b/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. +statement ok +SET datafusion.execution.batch_size = 8192; + +query II +SELECT file_row_index() + 1, column1 FROM parquet_table ORDER BY column1 +---- +1 10 +2 20 +3 30 +4 40 +5 50 + +# `target_partitions = 1` also removes the repartition, but is not covered here: +# `RESET` restores the global default rather than the harness value, which the +# harness rejects as leftover configuration. + statement ok DROP TABLE parquet_table; diff --git a/datafusion/sqllogictest/test_files/projection_pushdown.slt b/datafusion/sqllogictest/test_files/projection_pushdown.slt index b5b001026df76..a2e7cf5fb1248 100644 --- a/datafusion/sqllogictest/test_files/projection_pushdown.slt +++ b/datafusion/sqllogictest/test_files/projection_pushdown.slt @@ -2675,5 +2675,120 @@ A 2 statement ok DROP TABLE duplicate_alias; +# A leaf expression must still reach the scan when its projection also holds a +# `KeepInPlace` expression. See `try_split_leaf_projection` in +# `datafusion/physical-optimizer/src/projection_pushdown.rs` for why that used +# to be impossible. +# +# These are plan assertions on purpose. `get_field` stranded above the +# repartition still returns the right answer, it just reads the whole struct, +# so a result-only test passes whether or not the bug is present. +# +# Needs both a repartition (`target_partitions` > 1) and a row estimate above +# `batch_size`, hence the two SETs. Restore them with `SET` rather than `RESET`: +# `RESET` puts back the global default instead of the harness value, and the +# harness rejects a file that leaves configuration modified. + +statement ok +COPY ( + SELECT column1 AS id, column2 AS s + FROM VALUES + (1, {value: 100, label: 'alpha'}), + (2, {value: 200, label: 'beta'}), + (3, {value: 150, label: 'gamma'}), + (4, {value: 300, label: 'delta'}), + (5, {value: 250, label: 'epsilon'}) +) TO 'test_files/scratch/projection_pushdown/leaf_split.parquet' +STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE leaf_split STORED AS PARQUET +LOCATION 'test_files/scratch/projection_pushdown/leaf_split.parquet'; + +statement ok +SET datafusion.execution.target_partitions = 4; + +statement ok +SET datafusion.execution.batch_size = 2; + +# The scan must read only the `value` leaf. A plan holding a bare `s` in the +# `DataSourceExec` projection means the whole struct is being read. +query TT +EXPLAIN SELECT s['value'] + 1 AS v FROM leaf_split; +---- +logical_plan +01)Projection: get_field(leaf_split.s, Utf8("value")) + Int64(1) AS v +02)--TableScan: leaf_split projection=[s] +physical_plan +01)ProjectionExec: expr=[__datafusion_leaf_0@0 + 1 as v] +02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/leaf_split.parquet]]}, projection=[get_field(s@1, value) as __datafusion_leaf_0], file_type=parquet + +# The leaf is bare here: the unrelated computed column is what triggers it. +query TT +EXPLAIN SELECT s['value'] AS v, id + 1 AS w FROM leaf_split; +---- +logical_plan +01)Projection: get_field(leaf_split.s, Utf8("value")) AS v, leaf_split.id + Int64(1) AS w +02)--TableScan: leaf_split projection=[id, s] +physical_plan +01)ProjectionExec: expr=[__datafusion_leaf_0@1 as v, id@0 + 1 as w] +02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/leaf_split.parquet]]}, projection=[id, get_field(s@1, value) as __datafusion_leaf_0], file_type=parquet + +# `s` is genuinely needed by the outer projection here, so it must stay. +query TT +EXPLAIN SELECT s['value'] + 1 AS v, s AS whole FROM leaf_split; +---- +logical_plan +01)Projection: get_field(leaf_split.s, Utf8("value")) + Int64(1) AS v, leaf_split.s AS whole +02)--TableScan: leaf_split projection=[s] +physical_plan +01)ProjectionExec: expr=[__datafusion_leaf_0@1 + 1 as v, s@0 as whole] +02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/leaf_split.parquet]]}, projection=[s, get_field(s@1, value) as __datafusion_leaf_0], file_type=parquet + +# Two distinct leaves plus a computed sibling: both leaves must reach the scan +# and the struct must not be read whole. +query TT +EXPLAIN SELECT s['value'] AS v, s['label'] AS l, id + 1 AS w FROM leaf_split; +---- +logical_plan +01)Projection: get_field(leaf_split.s, Utf8("value")) AS v, get_field(leaf_split.s, Utf8("label")) AS l, leaf_split.id + Int64(1) AS w +02)--TableScan: leaf_split projection=[id, s] +physical_plan +01)ProjectionExec: expr=[__datafusion_leaf_0@1 as v, __datafusion_leaf_1@2 as l, id@0 + 1 as w] +02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/leaf_split.parquet]]}, projection=[id, get_field(s@1, value) as __datafusion_leaf_0, get_field(s@1, label) as __datafusion_leaf_1], file_type=parquet + +query ITI +SELECT s['value'] AS v, s['label'] AS l, id + 1 AS w FROM leaf_split ORDER BY v; +---- +100 alpha 2 +150 gamma 4 +200 beta 3 +250 epsilon 6 +300 delta 5 + +query I +SELECT s['value'] + 1 AS v FROM leaf_split ORDER BY v; +---- +101 +151 +201 +251 +301 + +statement ok +SET datafusion.execution.batch_size = 8192; + +# Back to 4, which is what this file leaves configured and what the harness +# started from. +statement ok +SET datafusion.execution.target_partitions = 4; + +statement ok +DROP TABLE leaf_split; + statement ok RESET datafusion.execution.parquet.pushdown_filters; From 322129abdcb33fbc306946f248d6b216b3d35537 Mon Sep 17 00:00:00 2001 From: Harry Shih Date: Fri, 6 Nov 2026 23:29:05 +0800 Subject: [PATCH 2/3] fix: rebase lambda parameters and keep metadata when splitting a projection Two regressions in the projection split, both reported in review. Each query below works on `main` and failed with the split. Lambda parameters are indexed after the input columns, so widening the input leaves them pointing at the wrong column: SELECT array_transform(arr, x -> x + 1), s['value'], s['label'] FROM t; `remap_filter_expr` in this file already rebases parameters for the join filter pushdown, so the node-level shift moves into `rebase_lambda_variable` and both callers use it rather than carrying two copies. The symptom depends on where the stale index lands. Out of range it is dropped and the query happens to be correct; in range it fails `LambdaVariable`'s field check; in range with a matching name and type it reads the wrong column and returns wrong values with no error. The last case is covered, since the field check is a backstop rather than the guarantee. `ProjectionExec::try_new` derives the output schema from the expressions alone and drops the field metadata the planner attached through `try_new_with_schema_metadata`, which the rule's own `schema_check` then rejects: PhysicalOptimizer rule 'ProjectionPushdown' failed. Schema mismatch. The split now rebuilds the projection with that constructor, taking the metadata from the projection it replaces. Tests. The value queries live in `projection_pushdown.slt`; the fixture gains `arr` and `tags`, appended after `s` so the existing `get_field(s@1, ...)` assertions do not move. A nested lambda needs two leaf expressions: with one the split is width-neutral, the parameter does not move, and the query passes against unfixed code. The metadata test is a Rust test because no SQL syntax attaches field metadata; it asserts on the physical plan's schema, not the DataFrame's, and uses Parquet because only a source that absorbs the leaf projection makes the split fire. Part of #27131 Co-Authored-By: Claude Opus 5 --- .../physical_optimizer/projection_pushdown.rs | 52 ++++++++++++++++ .../src/projection_pushdown.rs | 56 ++++++++++++----- .../test_files/projection_pushdown.slt | 61 +++++++++++++++++-- 3 files changed, 149 insertions(+), 20 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs index 0da2f759e3265..935497ddee2e2 100644 --- a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements +use std::collections::HashMap; use std::sync::Arc; use arrow::compute::SortOptions; @@ -10,7 +11,9 @@ use datafusion::datasource::listing::PartitionedFile; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::physical_plan::CsvSource; use datafusion::datasource::source::DataSourceExec; +use datafusion::prelude::{ParquetReadOptions, SessionContext}; use datafusion_common::config::{ConfigOptions, CsvOptions}; +use datafusion_common::metadata::FieldMetadata; use datafusion_common::{JoinSide, JoinType, NullEquality, Result, ScalarValue}; use datafusion_datasource::TableSchemaBuilder; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; @@ -20,6 +23,7 @@ use datafusion_expr::{ Operator, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, }; use datafusion_expr_common::columnar_value::ColumnarValue; +use datafusion_functions::core::expr_fn::get_field; use datafusion_physical_expr::expressions::{ BinaryExpr, CaseExpr, CastExpr, Column, Literal, NegativeExpr, binary, cast, col, }; @@ -1905,3 +1909,51 @@ fn test_sort_merge_join_interleaved_projection_embeds() -> Result<()> { Ok(()) } + +/// Splitting a projection rebuilds it, so it has to carry over the metadata the +/// planner attached; otherwise the rule's schema check fails the whole query. +/// Mirrors `keeps_metadata_of_constant_input` in `eliminate_constant_join`, but +/// over Parquet, because only a source that absorbs the leaf projection makes +/// the split fire at all. +#[tokio::test] +async fn split_keeps_field_metadata() -> Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("leaf_split_metadata.parquet"); + let path = path.to_str().expect("temp path is not utf-8"); + + let ctx = SessionContext::new(); + // Below the five rows, so `EnsureRequirements` inserts the repartition that + // the split needs. + ctx.sql("SET datafusion.execution.batch_size = 2").await?; + ctx.sql("SET datafusion.execution.target_partitions = 4") + .await?; + ctx.sql(&format!( + "COPY (SELECT column1 AS id, column2 AS s FROM VALUES \ + (1, {{value: 100}}), (2, {{value: 200}}), (3, {{value: 150}}), \ + (4, {{value: 300}}), (5, {{value: 250}})) \ + TO '{path}' STORED AS PARQUET" + )) + .await? + .collect() + .await?; + ctx.register_parquet("t", path, ParquetReadOptions::default()) + .await?; + + let metadata = + FieldMetadata::from(HashMap::from([("k".to_string(), "v".to_string())])); + let df = ctx.table("t").await?.select(vec![ + get_field(datafusion_expr::col("s"), "value") + .alias_with_metadata("v", Some(metadata)), + (datafusion_expr::col("id") + datafusion_expr::lit(1i64)).alias("w"), + ])?; + + // The physical plan is what the rule rewrites, so the assertion has to come + // from it rather than from the DataFrame's logical schema. + let plan = df.create_physical_plan().await?; + assert_eq!( + plan.schema().field(0).metadata().get("k"), + Some(&"v".to_string()) + ); + + Ok(()) +} diff --git a/datafusion/physical-optimizer/src/projection_pushdown.rs b/datafusion/physical-optimizer/src/projection_pushdown.rs index a45975b027012..ad3f0bbaaf68b 100644 --- a/datafusion/physical-optimizer/src/projection_pushdown.rs +++ b/datafusion/physical-optimizer/src/projection_pushdown.rs @@ -451,21 +451,32 @@ fn remap_filter_expr( } else if let Some(variable) = expr.downcast_ref::() && 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> { + 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"; @@ -547,6 +558,7 @@ fn try_split_leaf_projection( Arc::clone(&proj_expr.expr), &leaves, &passthrough_indices, + child_schema.fields().len(), )?, alias: proj_expr.alias.clone(), }) @@ -559,10 +571,15 @@ fn try_split_leaf_projection( vec![new_child], ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), )?; - Ok(Some(Arc::new(ProjectionExec::try_new( - new_exprs, - new_repartition, - )?))) + // `try_new` would derive the output schema from the expressions alone and + // drop the metadata the planner set, which `schema_check` then rejects. + Ok(Some(Arc::new( + ProjectionExec::try_new_with_schema_metadata( + new_exprs, + new_repartition, + projection.schema().as_ref(), + )?, + ))) } /// Collects the distinct outermost [`ExpressionPlacement::MoveTowardsLeafNodes`] @@ -614,14 +631,17 @@ fn collect_outer_columns( /// 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. +/// for it, and every surviving column or lambda parameter is remapped to its +/// new index. `old_input_len` is the child's width before the split. fn rewrite_outer_expr( expr: Arc, leaves: &[Arc], passthrough_indices: &[usize], + old_input_len: usize, ) -> Result> { // The leaf columns were appended after the passthrough ones. let first_leaf = passthrough_indices.len(); + let new_input_len = first_leaf + leaves.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); @@ -646,6 +666,14 @@ fn rewrite_outer_expr( return Ok(Transformed::yes(Arc::new(remapped) as Arc)); } } + // The split widens the input, which moves the lambda parameters indexed + // after it; see `remap_filter_expr`. + if let Some(variable) = node.downcast_ref::() + && old_input_len != new_input_len + { + let rebased = rebase_lambda_variable(variable, old_input_len, new_input_len)?; + return Ok(Transformed::yes(rebased)); + } Ok(Transformed::no(node)) }) .data() diff --git a/datafusion/sqllogictest/test_files/projection_pushdown.slt b/datafusion/sqllogictest/test_files/projection_pushdown.slt index a2e7cf5fb1248..65e78a35d873b 100644 --- a/datafusion/sqllogictest/test_files/projection_pushdown.slt +++ b/datafusion/sqllogictest/test_files/projection_pushdown.slt @@ -2689,15 +2689,17 @@ DROP TABLE duplicate_alias; # `RESET` puts back the global default instead of the harness value, and the # harness rejects a file that leaves configuration modified. +# `arr` and `tags` are only read by the lambda queries further down. They are +# appended after `s` so the `get_field(s@1, ...)` indices below do not move. statement ok COPY ( - SELECT column1 AS id, column2 AS s + SELECT column1 AS id, column2 AS s, column3 AS arr, column4 AS tags FROM VALUES - (1, {value: 100, label: 'alpha'}), - (2, {value: 200, label: 'beta'}), - (3, {value: 150, label: 'gamma'}), - (4, {value: 300, label: 'delta'}), - (5, {value: 250, label: 'epsilon'}) + (1, {value: 100, label: 'alpha'}, [1, 2, 3], ['a', 'b']), + (2, {value: 200, label: 'beta'}, [4, 5], ['c']), + (3, {value: 150, label: 'gamma'}, [6], ['d', 'e', 'f']), + (4, {value: 300, label: 'delta'}, [7, 8, 9], ['g']), + (5, {value: 250, label: 'epsilon'}, [10], ['h', 'i']) ) TO 'test_files/scratch/projection_pushdown/leaf_split.parquet' STORED AS PARQUET; @@ -2779,6 +2781,53 @@ SELECT s['value'] + 1 AS v FROM leaf_split ORDER BY v; 251 301 +# A lambda parameter is indexed after the input columns, so the split moves it; +# see `rewrite_outer_expr`. These are value queries because an index left behind +# is only sometimes an error, see the third one. + +statement ok +SET datafusion.sql_parser.dialect = databricks; + +query ?IT +SELECT array_transform(arr, x -> x + 1) AS a, s['value'] AS v, s['label'] AS l +FROM leaf_split ORDER BY v; +---- +[2, 3, 4] 100 alpha +[7] 150 gamma +[5, 6] 200 beta +[11] 250 epsilon +[8, 9, 10] 300 delta + +# Nested lambdas need two leaves: with one the split is width-neutral, the +# parameter does not move, and the query passes against broken code. +query ?IT +SELECT array_transform(arr, x -> x + array_length(array_transform(arr, y -> y * 2))) AS a, + s['value'] AS v, s['label'] AS l +FROM leaf_split ORDER BY v; +---- +[4, 5, 6] 100 alpha +[7] 150 gamma +[6, 7] 200 beta +[11] 250 epsilon +[10, 11, 12] 300 delta + +# Naming the parameter after the column a stale index would land on makes +# `LambdaVariable`'s field check match, so the wrong column is read with no +# error at all. The name has to track `leaf_column_name` to stay adversarial. +query ?ITI +SELECT array_transform(tags, __datafusion_leaf_1 -> concat(__datafusion_leaf_1, '!')) AS a, + s['value'] AS v, s['label'] AS l, id +FROM leaf_split ORDER BY v; +---- +[a!, b!] 100 alpha 1 +[d!, e!, f!] 150 gamma 3 +[c!] 200 beta 2 +[h!, i!] 250 epsilon 5 +[g!] 300 delta 4 + +statement ok +SET datafusion.sql_parser.dialect = generic; + statement ok SET datafusion.execution.batch_size = 8192; From 4ee8c648324502b0d0c23437104802564a9a5f76 Mon Sep 17 00:00:00 2001 From: Harry Shih Date: Sat, 7 Nov 2026 22:48:06 +0800 Subject: [PATCH 3/3] test: pin the leaf projection split's refusal path and the shapes it fixes Review follow-up. Nothing here changes what the split computes. The split is only committed when the child really absorbs the leaf projection, but no test held that: removing the check left all 545 SLT files, 50 physical-optimizer unit tests and 669 `core_integration` ones green. A memory table refuses the lower projection, because `all_alias_free_columns` rejects an aliased `get_field`, and that is also what any node that does not override `try_swapping_with_projection` does. An `EXPLAIN` on such a table now pins the plan; with the check removed it fails on the extra `ProjectionExec`. Collecting each distinct leaf once was likewise unpinned, and the values cannot show it: two calls to the same marker agree with each other whether they share a column or not. The duplicated-marker case gains an `EXPLAIN`, which fails with a second `__datafusion_leaf_1` in the scan projection if the de-duplication goes. Three shapes the split fixes had no test either. Each reaches a different path: a TopK sorts on a column the split produced and runs above the repartition with `preserve_partitioning=[true]`; a marker inside a lambda body has to be found through `LambdaExpr`; and CSV is a second scan implementation, while every other test here uses Parquet. All three error with `is source dependent and cannot be evaluated directly` without the split. The cleanup walk this PR had added after the split is removed rather than made conditional, because it cannot do anything. `remove_unnecessary_projections` has three paths for a projection: it is not removable, since the split renames and keeps the original output schema; `overrides_metadata` changes nothing; and `RepartitionExec` declines it either for not narrowing, when every expression stayed a plain column, or for holding a computed one, when an expression consumed more than one leaf and the projection did narrow. Nothing else in the tree moved, as the split preserves the projection's schema. Instrumenting the walk confirmed it changed no plan in any of the 19 splits the suite fires. Handling only `RoundRobinBatch` was unpinned the same way, and no SQL reaches that shape: the guard is hit 1326 times across the suite, every one of them round robin. The rule is a plain plan-to-plan function though, so the plan can be built by hand, which `split_declines_hash_repartition` does, after `ensure_requirements.rs` for the hash repartition and `test_utils.rs` for the scan. The scan has to be Parquet rather than `EmptyExec`: only a source that absorbs the leaf projection lets the split commit once the check is gone. Also drops a comment reference to a trap number that only exists in the pull request description, and rewrites the one on `collect_leaf_exprs` that used `s['a']['b']` as its example: the logical optimizer folds that into a single call, so the branch it describes never sees it. It now says what stopping at a leaf avoids instead. Part of #27131 Co-Authored-By: Claude Opus 5 --- .../physical_optimizer/projection_pushdown.rs | 80 ++++++++++++++++++- .../src/projection_pushdown.rs | 27 ++++--- .../test_files/file_row_index.slt | 26 +++++- .../test_files/input_file_name.slt | 21 +++++ .../test_files/projection_pushdown.slt | 36 +++++++++ 5 files changed, 176 insertions(+), 14 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs index 935497ddee2e2..c0936eb929b37 100644 --- a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs @@ -9,7 +9,7 @@ use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::datasource::listing::PartitionedFile; use datafusion::datasource::memory::MemorySourceConfig; -use datafusion::datasource::physical_plan::CsvSource; +use datafusion::datasource::physical_plan::{CsvSource, ParquetSource}; use datafusion::datasource::source::DataSourceExec; use datafusion::prelude::{ParquetReadOptions, SessionContext}; use datafusion_common::config::{ConfigOptions, CsvOptions}; @@ -1957,3 +1957,81 @@ async fn split_keeps_field_metadata() -> Result<()> { Ok(()) } + +/// The split only handles round robin partitioning, because hash and range keys +/// reference columns by index and would have to be remapped against the +/// rewritten child schema. No SQL reaches this shape, so the test builds it: +/// above a join, the leaf is already extracted into its own projection below the +/// join, and nothing else puts a projection on a hash repartition. +#[tokio::test] +async fn split_declines_hash_repartition() -> Result<()> { + let schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, true), + Field::new( + "s", + DataType::Struct(vec![Field::new("value", DataType::Int64, true)].into()), + true, + ), + ])); + let config = FileScanConfigBuilder::new( + ObjectStoreUrl::parse("test:///").unwrap(), + Arc::new(ParquetSource::new(Arc::clone(&schema))), + ) + .with_file(PartitionedFile::new("x".to_string(), 100)) + .build(); + // Parquet absorbs a leaf projection, so the split would commit here if the + // partitioning check did not stop it first. + let hash = Arc::new(RepartitionExec::try_new( + DataSourceExec::from_data_source(config), + Partitioning::Hash(vec![col("id", &schema)?], 4), + )?); + + let leaf: Arc = Arc::new(ScalarFunctionExpr::new( + "get_field", + datafusion_functions::core::get_field(), + vec![ + col("s", &schema)?, + Arc::new(Literal::new(ScalarValue::from("value"))), + ], + Arc::new(Field::new("v", DataType::Int64, true)), + Arc::new(ConfigOptions::default()), + )); + let projection: Arc = Arc::new(ProjectionExec::try_new( + vec![ + ProjectionExpr { + expr: leaf, + alias: "v".to_string(), + }, + ProjectionExpr { + expr: binary( + col("id", &schema)?, + Operator::Plus, + Arc::new(Literal::new(ScalarValue::Int64(Some(1)))), + &schema, + )?, + alias: "w".to_string(), + }, + ], + hash, + )?); + + let optimized = ProjectionPushdown::new() + .optimize(Arc::clone(&projection), &ConfigOptions::default())?; + let actual = displayable(optimized.as_ref()) + .indent(true) + .to_string() + .trim() + .to_string(); + + // Unchanged. Without the check the scan absorbs the leaf and the hash keys + // are left indexing the old schema. + assert_snapshot!( + actual, + @r" + ProjectionExec: expr=[get_field(s@1, value) as v, id@0 + 1 as w] + RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[id, s], file_type=parquet + " + ); + Ok(()) +} diff --git a/datafusion/physical-optimizer/src/projection_pushdown.rs b/datafusion/physical-optimizer/src/projection_pushdown.rs index ad3f0bbaaf68b..364ec8d75aab5 100644 --- a/datafusion/physical-optimizer/src/projection_pushdown.rs +++ b/datafusion/physical-optimizer/src/projection_pushdown.rs @@ -67,18 +67,19 @@ 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. + // stuck above a `RepartitionExec`. The split needs no cleanup pass of + // its own: the outer projection it leaves either keeps every expression + // a plain column, and is then no narrower than the widened child, or it + // narrows and holds a computed one. `RepartitionExec` declines both, it + // is not removable because it renames, and nothing else moved. 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), - }) + 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() + }) + .data() } fn name(&self) -> &str { @@ -585,8 +586,10 @@ fn try_split_leaf_projection( /// 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. +/// Recursion stops at a leaf, so one nested inside another is not collected too: +/// that would add a column no expression reads, and read more of the struct than +/// the outer leaf needs. Plain SQL does not reach it, because the logical +/// optimizer folds nested field access into a single call. fn collect_leaf_exprs(exprs: &[ProjectionExpr]) -> Result>> { let mut leaves: Vec> = vec![]; for proj_expr in exprs { diff --git a/datafusion/sqllogictest/test_files/file_row_index.slt b/datafusion/sqllogictest/test_files/file_row_index.slt index 7e135731e2c0b..45ccb38a8fbde 100644 --- a/datafusion/sqllogictest/test_files/file_row_index.slt +++ b/datafusion/sqllogictest/test_files/file_row_index.slt @@ -240,6 +240,29 @@ FROM parquet_table ORDER BY column1 3 4 40 4 5 50 +# The single column is what the values above cannot show: both outputs would +# also agree if the marker were read twice. +query TT +EXPLAIN SELECT file_row_index(), file_row_index() + 1, column1 FROM parquet_table +---- +logical_plan +01)Projection: file_row_index(), file_row_index() + Int64(1), parquet_table.column1 +02)--TableScan: parquet_table projection=[column1] +physical_plan +01)ProjectionExec: expr=[__datafusion_leaf_0@1 as file_row_index(), __datafusion_leaf_0@1 + 1 as file_row_index() + Int64(1), column1@0 as column1] +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, CAST(__datafusion_file_row_index@1 AS Int64;N) as __datafusion_leaf_0], file_type=parquet + +# `ORDER BY ... LIMIT` is a TopK that sorts on a column the split produced, +# with `preserve_partitioning=[true]`, so it runs above the repartition on the +# rewritten schema. No dynamic filter reaches the scan here: that needs the +# sort key to be the computed sibling rather than the rewritten marker. +query II +SELECT file_row_index() AS r, column1 + 1 FROM parquet_table ORDER BY r DESC LIMIT 2 +---- +4 51 +3 41 + # 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 @@ -273,7 +296,8 @@ SELECT file_row_index(), column1 FROM parquet_table ORDER BY column1 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. +# keep working, and it is also the control showing that `batch_size` is what +# triggers the repartition. statement ok SET datafusion.execution.batch_size = 8192; diff --git a/datafusion/sqllogictest/test_files/input_file_name.slt b/datafusion/sqllogictest/test_files/input_file_name.slt index 3ab598fb1e451..53639c17faf12 100644 --- a/datafusion/sqllogictest/test_files/input_file_name.slt +++ b/datafusion/sqllogictest/test_files/input_file_name.slt @@ -64,6 +64,27 @@ physical_plan query error Execution error: input_file_name\(\) is source dependent and cannot be evaluated directly SELECT input_file_name() FROM (VALUES (1)) v(x); +# CSV has its own scan implementation, so the leaf split has to reach it too. +# Needs a `KeepInPlace` sibling and a repartition, hence `column1 + 1` and the +# `batch_size` below the row count. +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; + statement ok DROP TABLE csv_table; diff --git a/datafusion/sqllogictest/test_files/projection_pushdown.slt b/datafusion/sqllogictest/test_files/projection_pushdown.slt index 65e78a35d873b..eb698f6ef9945 100644 --- a/datafusion/sqllogictest/test_files/projection_pushdown.slt +++ b/datafusion/sqllogictest/test_files/projection_pushdown.slt @@ -2825,9 +2825,45 @@ FROM leaf_split ORDER BY v; [h!, i!] 250 epsilon 5 [g!] 300 delta 4 +# A leaf expression inside the lambda body, rather than beside it: the split has +# to find it through `LambdaExpr` before the parameter rebase even applies. +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 + statement ok SET datafusion.sql_parser.dialect = generic; +# A source that does not absorb the projection keeps the plan unchanged: the +# split must not add a projection nobody can use. A memory table refuses +# because `all_alias_free_columns` rejects the aliased `get_field`, and that is +# the default for any node that does not override `try_swapping_with_projection`. +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; + statement ok SET datafusion.execution.batch_size = 8192;