diff --git a/datafusion/datasource/src/memory.rs b/datafusion/datasource/src/memory.rs index 49410ad695f77..61c9ccc6f23e5 100644 --- a/datafusion/datasource/src/memory.rs +++ b/datafusion/datasource/src/memory.rs @@ -1023,6 +1023,74 @@ mod tests { use datafusion_physical_plan::ExecutionPlan; + #[tokio::test] + #[expect(deprecated, reason = "Verify exact IPC dictionary properties")] + async fn projection_collapse_preserves_dictionary_stream_schema() -> Result<()> { + use arrow::array::StringDictionaryBuilder; + use arrow::datatypes::Int8Type; + use datafusion_physical_expr::expressions::Column; + use datafusion_physical_expr::projection::ProjectionExpr; + use datafusion_physical_plan::common::collect; + use datafusion_physical_plan::projection::{ + ProjectionExec, remove_unnecessary_projections, + }; + use std::collections::HashMap; + + let field = Field::new_dict( + "i", + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), + true, + 7, + true, + ); + let schema = Arc::new(Schema::new(vec![field.clone()])); + let mut values = StringDictionaryBuilder::::new(); + values.append("first")?; + values.append_null(); + values.append("second")?; + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(values.finish())])?; + let input: Arc = + MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None)?; + let expression = || ProjectionExpr { + expr: Arc::new(Column::new("i", 0)), + alias: "i".to_string(), + }; + let marked = + Schema::new(vec![field.clone().with_metadata(HashMap::from([( + "scope".to_string(), + "inner".to_string(), + )]))]); + let inner: Arc = Arc::new( + ProjectionExec::try_new_with_schema_metadata([expression()], input, &marked)?, + ); + let outer: Arc = + Arc::new(ProjectionExec::try_new_with_schema_metadata( + [expression()], + inner, + &Schema::new(vec![field]), + )?); + let expected_schema = outer.schema(); + assert_eq!(expected_schema.field(0).dict_id(), Some(0)); + assert_eq!(expected_schema.field(0).dict_is_ordered(), Some(false)); + let expected = + collect(outer.execute(0, Arc::new(TaskContext::default()))?).await?; + + let optimized = remove_unnecessary_projections(outer)?.data; + let stream = optimized.execute(0, Arc::new(TaskContext::default()))?; + for schema in [optimized.schema(), stream.schema()] { + assert_eq!(schema.field(0).dict_id(), Some(0)); + assert_eq!(schema.field(0).dict_is_ordered(), Some(false)); + } + let actual = collect(stream).await?; + assert_eq!(actual, expected); + for batch in actual { + assert_eq!(batch.schema().field(0).dict_id(), Some(0)); + assert_eq!(batch.schema().field(0).dict_is_ordered(), Some(false)); + } + Ok(()) + } + #[tokio::test] async fn exec_with_limit() -> Result<()> { let task_ctx = Arc::new(TaskContext::default()); diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 6bde9ca74d076..f8cef2e61640b 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -146,3 +146,7 @@ required-features = ["test_utils"] [[bench]] harness = false name = "bounded_window" + +[[bench]] +harness = false +name = "projection_metadata" diff --git a/datafusion/physical-plan/benches/projection_metadata.rs b/datafusion/physical-plan/benches/projection_metadata.rs new file mode 100644 index 0000000000000..b1bc604b4a8ba --- /dev/null +++ b/datafusion/physical-plan/benches/projection_metadata.rs @@ -0,0 +1,180 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Column-only projection chains, runnable unchanged before/after metadata collapse. + +use std::collections::HashMap; +use std::hint::black_box; +use std::sync::Arc; +use std::time::Duration; + +use arrow::array::{ArrayRef, Int64Array}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use criterion::async_executor::FuturesExecutor; +use criterion::{BatchSize, BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::tree_node::TreeNode; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::common::collect; +use datafusion_physical_plan::execution_plan::reset_plan_states; +use datafusion_physical_plan::projection::{ + ProjectionExec, ProjectionExpr, remove_unnecessary_projections, +}; +use datafusion_physical_plan::test::TestMemoryExec; +use futures::executor::block_on; + +fn schema(width: usize, layer: Option) -> SchemaRef { + let fields: Vec<_> = (0..width) + .map(|column| { + let field = Field::new(format!("c{column}"), DataType::Int64, true); + match layer { + Some(layer) => field.with_metadata(HashMap::from([( + "field-layer".to_string(), + format!("{layer}:{column}"), + )])), + None => field, + } + }) + .collect(); + let metadata = layer + .map(|layer| HashMap::from([("schema-layer".to_string(), layer.to_string())])) + .unwrap_or_default(); + Arc::new(Schema::new_with_metadata(fields, metadata)) +} + +fn fixture( + rows: usize, + width: usize, + depth: usize, + metadata: bool, +) -> (RecordBatch, Arc, SchemaRef) { + // Adjacent duplicate rows and periodic NULLs exercise multiplicity/validity. + let columns = (0..width) + .map(|column| { + Arc::new(Int64Array::from_iter((0..rows).map(|row| { + ((row / 2 + column) % 31 != 7) + .then_some((column * 100 + row / 2 % 17) as i64) + }))) as ArrayRef + }) + .collect(); + let input = RecordBatch::try_new(schema(width, None), columns).unwrap(); + let mut plan: Arc = + TestMemoryExec::try_new_exec(&[vec![input.clone()]], input.schema(), None) + .unwrap(); + for layer in 1..=depth { + let expressions = (0..width).map(|column| ProjectionExpr { + expr: Arc::new(Column::new(&format!("c{column}"), column)), + alias: format!("c{column}"), + }); + plan = Arc::new(if metadata { + ProjectionExec::try_new_with_schema_metadata( + expressions, + plan, + schema(width, Some(layer)).as_ref(), + ) + .unwrap() + } else { + ProjectionExec::try_new(expressions, plan).unwrap() + }); + } + // Expected output comes from the input and explicit outer metadata, not a plan. + (input, plan, schema(width, metadata.then_some(depth))) +} + +fn optimize(plan: Arc) -> Arc { + plan.transform_down(remove_unnecessary_projections) + .unwrap() + .data +} + +fn validate( + plan: &Arc, + input: &RecordBatch, + expected: &Schema, + context: &Arc, +) { + assert_eq!(plan.schema().as_ref(), expected); + let batches = + block_on(collect(plan.execute(0, Arc::clone(context)).unwrap())).unwrap(); + let mut offset = 0; + for batch in batches { + assert_eq!(batch.schema().as_ref(), expected); + for (actual, source) in batch.columns().iter().zip(input.columns()) { + let actual = actual.as_any().downcast_ref::().unwrap(); + let source = source.as_any().downcast_ref::().unwrap(); + assert_eq!(actual, &source.slice(offset, batch.num_rows())); + } + offset += batch.num_rows(); + } + assert_eq!(offset, input.num_rows()); +} + +fn projection_metadata(c: &mut Criterion) { + let mut group = c.benchmark_group("projection_metadata"); + group.sample_size(10); + group.warm_up_time(Duration::from_millis(200)); + group.measurement_time(Duration::from_secs(1)); + // One ~1 MiB fixture is live; all plans share its Arrow buffers. + for (rows, width, depth) in [(1, 1, 1), (8192, 1, 8), (8192, 16, 8)] { + for metadata in [false, true] { + let (input, original, expected) = fixture(rows, width, depth, metadata); + let optimized = optimize(Arc::clone(&original)); + let context = Arc::new(TaskContext::default()); + validate(&original, &input, expected.as_ref(), &context); + validate(&optimized, &input, expected.as_ref(), &context); + let mode = if metadata { "overrides" } else { "unchanged" }; + let name = format!("{mode}/rows{rows}_width{width}_depth{depth}"); + for (phase, plan, apply_rule) in [ + ("optimize_and_collect", &original, true), + ("execute_already_optimized", &optimized, false), + ] { + group.bench_function(BenchmarkId::new(phase, &name), |b| { + b.to_async(FuturesExecutor).iter_batched( + // Untimed: fresh metrics on the prebuilt topology; otherwise + // repeated execute() appends metrics without bound. + || reset_plan_states(Arc::clone(plan)).unwrap(), + |plan| { + let context = Arc::clone(&context); + async move { + let running = Arc::clone(&plan); + let running = if apply_rule { + optimize(running) + } else { + running + }; + let batches = + collect(running.execute(0, context).unwrap()) + .await + .unwrap(); + // Timed: Arc clones, optional rule, stream setup, + // collection/materialization, batch and plan drops. + drop(black_box(batches)); + } + }, + BatchSize::PerIteration, + ); + }); + } + } + } + group.finish(); +} + +criterion_group!(benches, projection_metadata); +criterion_main!(benches); diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index 81ce9a585fe78..58bf0843dacf1 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -1004,18 +1004,44 @@ pub fn remove_unnecessary_projections( // Swapping a projection with observable metadata can change query results // by changing the metadata visible to its child expressions. if projection.overrides_metadata() { - return Ok(Transformed::no(plan)); + if let Some(inner) = projection.input().downcast_ref::() + && is_column_only_projection(projection) + && is_column_only_projection(inner) + { + try_collapse_projection_chain(projection)? + } else { + return Ok(Transformed::no(plan)); + } + } else { + // Otherwise, check if we can push it under its child(ren): + projection + .input() + .try_swapping_with_projection(projection)? } - // Otherwise, check if we can push it under its child(ren): - projection - .input() - .try_swapping_with_projection(projection)? } else { return Ok(Transformed::no(plan)); }; + let maybe_modified = maybe_modified.filter(|replacement| { + root_dictionary_properties_match(&plan.schema(), &replacement.schema()) + }); Ok(maybe_modified.map_or_else(|| Transformed::no(plan), Transformed::yes)) } +/// Column projections rebuild root fields but clone nested storage types. +/// Arrow schema equality deliberately excludes these root IPC properties. +#[expect(deprecated, reason = "Preserve observable IPC dictionary properties")] +fn root_dictionary_properties_match(left: &Schema, right: &Schema) -> bool { + left.fields().len() == right.fields().len() + && left + .fields() + .iter() + .zip(right.fields()) + .all(|(left, right)| { + left.dict_id() == right.dict_id() + && left.dict_is_ordered() == right.dict_is_ordered() + }) +} + /// Compare the inputs and outputs of the projection. All expressions must be /// columns without alias, and projection does not change the order of fields. /// The input and output schemas must also match exactly to preserve metadata. @@ -1029,7 +1055,11 @@ fn is_projection_removable(projection: &ProjectionExec) -> bool { }; col.name() == proj_expr.alias && col.index() == idx }) && exprs.len() == projection.input().schema().fields().len() - && projection.schema() == projection.input().schema() + && { + let output = projection.schema(); + let input = projection.input().schema(); + output == input && root_dictionary_properties_match(&output, &input) + } } /// Given the expression set of a projection, checks if the projection causes @@ -1338,7 +1368,8 @@ pub fn update_join_filter( /// The projection-removal optimizer checks `outer.overrides_metadata()` before /// reaching this helper. The unified projection also keeps `outer`'s schema, so /// collapsing cannot lose its output metadata. Inner projections still need the -/// check below because outer expressions may observe their metadata. +/// check below because outer expressions may observe their metadata. Column-only +/// chains cannot inspect metadata and retain the outer schema when collapsed. fn try_collapse_projection_chain( outer: &ProjectionExec, ) -> Result>> { @@ -1348,7 +1379,12 @@ fn try_collapse_projection_chain( let mut collapsed_any = false; 'outer: while let Some(inner_proj) = current_input.downcast_ref::() { - if inner_proj.overrides_metadata() { + if inner_proj.overrides_metadata() + && !(current_exprs + .iter() + .all(|expr| expr.expr.downcast_ref::().is_some()) + && is_column_only_projection(inner_proj)) + { break; } @@ -1411,6 +1447,13 @@ fn try_collapse_projection_chain( remove_unnecessary_projections(unified).data().map(Some) } +fn is_column_only_projection(projection: &ProjectionExec) -> bool { + projection + .expr() + .iter() + .all(|expr| expr.expr.downcast_ref::().is_some()) +} + /// Collect all column indices from the given projection expressions. fn collect_column_indices(exprs: &[ProjectionExpr]) -> Vec { // Collect column indices in a deterministic order that preserves the @@ -1716,7 +1759,7 @@ mod tests { expr: Arc::new(arrow_metadata), alias: "metadata".to_string(), }], - inner, + Arc::clone(&inner), )?); let outer_projection = outer @@ -1724,6 +1767,36 @@ mod tests { .expect("test plan should be a ProjectionExec"); assert!(try_collapse_projection_chain(outer_projection)?.is_none()); + // Collapsing a column wrapper first must not move the resulting + // metadata-observing expression through the next inner projection. + let renamed: Arc = Arc::new(ProjectionExec::try_new( + [ProjectionExpr { + expr: Arc::new(Column::new("metadata", 0)), + alias: "renamed_metadata".to_string(), + }], + Arc::clone(&outer), + )?); + let rewritten = remove_unnecessary_projections(renamed)?; + assert!(rewritten.transformed); + let projection = rewritten + .data + .downcast_ref::() + .expect("metadata-observing projection remains"); + assert!(Arc::ptr_eq(projection.input(), &inner)); + assert_eq!(projection.schema().field(0).name(), "renamed_metadata"); + let batches = collect( + rewritten + .data + .execute(0, Arc::new(TaskContext::default()))?, + ) + .await?; + let values = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .expect("metadata expression should return Utf8"); + assert_eq!(values.value(0), "true"); + let optimized = remove_unnecessary_projections(outer)?.data; let batches = collect(optimized.execute(0, Arc::new(TaskContext::default()))?).await?; @@ -1892,6 +1965,131 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_column_projection_chain_collapses_preserving_metadata_and_rows() + -> Result<()> { + for (field_metadata, schema_metadata) in [ + ( + HashMap::from([("event_field".to_string(), "outer".to_string())]), + HashMap::new(), + ), + ( + HashMap::new(), + HashMap::from([("scope".to_string(), "outer".to_string())]), + ), + ] { + let input = test::scan_partitioned(1); + let inner = identity_projection_with_metadata( + Arc::clone(&input), + HashMap::from([("event_field".to_string(), "inner".to_string())]), + HashMap::from([("scope".to_string(), "inner".to_string())]), + )?; + let outer = identity_projection_with_metadata( + inner, + field_metadata.clone(), + schema_metadata.clone(), + )?; + let expected = + collect(outer.execute(0, Arc::new(TaskContext::default()))?).await?; + + let optimized = remove_unnecessary_projections(Arc::clone(&outer))?; + assert!(optimized.transformed, "column-only chain should collapse"); + let projection = optimized + .data + .downcast_ref::() + .expect("outer metadata still requires a projection"); + assert!(Arc::ptr_eq(projection.input(), &input)); + assert_eq!(projection.schema(), outer.schema()); + assert_eq!(projection.schema().field(0).metadata(), &field_metadata); + assert_eq!(projection.schema().metadata(), &schema_metadata); + let actual = collect( + optimized + .data + .execute(0, Arc::new(TaskContext::default()))?, + ) + .await?; + assert_eq!(actual, expected); + } + Ok(()) + } + + #[test] + #[expect(deprecated, reason = "Verify exact IPC dictionary properties")] + fn test_clearing_projection_metadata_keeps_dictionary_properties() -> Result<()> { + let source_field = Field::new_dict( + "i", + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), + true, + 7, + true, + ); + let input: Arc = + Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![source_field])))); + let inner = identity_projection_with_metadata( + input, + HashMap::from([("scope".to_string(), "inner".to_string())]), + HashMap::new(), + )?; + let outer = + identity_projection_with_metadata(inner, HashMap::new(), HashMap::new())?; + let expected = outer.schema(); + assert_eq!(expected.field(0).dict_id(), Some(0)); + assert_eq!(expected.field(0).dict_is_ordered(), Some(false)); + let optimized = remove_unnecessary_projections(outer)?; + let actual = optimized.data.schema(); + assert_eq!(actual.field(0).dict_id(), expected.field(0).dict_id()); + assert_eq!( + actual.field(0).dict_is_ordered(), + expected.field(0).dict_is_ordered() + ); + Ok(()) + } + + #[test] + #[expect(deprecated, reason = "Verify exact IPC dictionary properties")] + fn test_column_projection_chain_preserves_dictionary_properties() -> Result<()> { + let dictionary = Field::new_dict( + "member", + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), + true, + 7, + true, + ); + let fields = [ + dictionary.clone().with_name("i"), + Field::new("i", DataType::Struct(vec![dictionary.clone()].into()), true), + Field::new("i", DataType::List(Arc::new(dictionary)), true), + ]; + for field in fields { + let input: Arc = + Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![field])))); + let inner = identity_projection_with_metadata( + Arc::clone(&input), + HashMap::from([("scope".to_string(), "inner".to_string())]), + HashMap::new(), + )?; + let outer = identity_projection_with_metadata( + inner, + HashMap::from([("scope".to_string(), "outer".to_string())]), + HashMap::new(), + )?; + let optimized = remove_unnecessary_projections(Arc::clone(&outer))?; + assert!(optimized.transformed); + let actual = optimized.data.schema(); + let expected = outer.schema(); + assert_eq!(actual, expected); + for (actual, expected) in actual + .flattened_fields() + .into_iter() + .zip(expected.flattened_fields()) + { + assert_eq!(actual.dict_id(), expected.dict_id()); + assert_eq!(actual.dict_is_ordered(), expected.dict_is_ordered()); + } + } + Ok(()) + } + #[test] fn test_join_projection_pushdown_preserves_output_metadata() -> Result<()> { let source_metadata =