diff --git a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs index 0da2f759e3265..c0936eb929b37 100644 --- a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs @@ -2,15 +2,18 @@ // 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; 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}; +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,129 @@ 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(()) +} + +/// 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 7faa72c1d947e..364ec8d75aab5 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,7 +66,20 @@ impl PhysicalOptimizerRule for ProjectionPushdown { }) .map(|t| t.data)?; - plan.transform_down(remove_unnecessary_projections).data() + // Push what can already descend, then split the projections that are + // 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()?; + plan.transform_down(|plan| { + Ok(match try_split_leaf_projection(&plan)? { + Some(new_plan) => Transformed::yes(new_plan), + None => Transformed::no(plan), + }) + }) + .data() } fn name(&self) -> &str { @@ -434,21 +452,236 @@ 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"; + +/// 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, + child_schema.fields().len(), + )?, + 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), + )?; + // `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`] +/// sub-expressions of `exprs`. +/// +/// 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 { + 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 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); + 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)); + } + } + // 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() +} + #[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..45ccb38a8fbde 100644 --- a/datafusion/sqllogictest/test_files/file_row_index.slt +++ b/datafusion/sqllogictest/test_files/file_row_index.slt @@ -151,6 +151,169 @@ 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 + +# 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 +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 showing that `batch_size` is what +# triggers the repartition. +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/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 b5b001026df76..eb698f6ef9945 100644 --- a/datafusion/sqllogictest/test_files/projection_pushdown.slt +++ b/datafusion/sqllogictest/test_files/projection_pushdown.slt @@ -2675,5 +2675,205 @@ 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. + +# `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, column3 AS arr, column4 AS tags + FROM VALUES + (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; + +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 + +# 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 + +# 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; + +# 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;