diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index 222c0ec688b78..09e1614ad4f61 100644 --- a/datafusion/core/Cargo.toml +++ b/datafusion/core/Cargo.toml @@ -285,6 +285,11 @@ harness = false name = "parquet_struct_projection" required-features = ["parquet"] +[[bench]] +harness = false +name = "parquet_wide_scan" +required-features = ["parquet"] + [[bench]] harness = false name = "parquet_struct_shared_prefix_pushdown" diff --git a/datafusion/core/benches/parquet_wide_scan.rs b/datafusion/core/benches/parquet_wide_scan.rs new file mode 100644 index 0000000000000..cc6bd8f6d4998 --- /dev/null +++ b/datafusion/core/benches/parquet_wide_scan.rs @@ -0,0 +1,196 @@ +// 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. + +//! Benchmarks for scanning a parquet table without a projection. +//! +//! `scan_construction` builds the physical scan directly and is where not +//! materializing an identity projection shows up: the work it removes is +//! proportional to the table's width. +//! +//! `planning` and `execution` are the end-to-end context for that, and are +//! controls rather than targets. Physical planning of `SELECT *` is dominated +//! by expanding the wildcard and running the optimizer over one expression per +//! column, and a full scan is dominated by decoding; neither moves measurably. +//! The narrow `SELECT c0` variants push a genuine projection and should not +//! move either. + +use std::hint::black_box; +use std::path::Path; +use std::sync::Arc; +use std::time::Duration; + +use arrow::array::{ArrayRef, Int32Array}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion::prelude::SessionContext; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::WriterProperties; +use tempfile::{Builder, NamedTempFile}; +use tokio::runtime::Runtime; + +/// Rows per written batch, and per row group, for the execution benchmarks +const EXEC_BATCH_ROWS: usize = 8192; +/// Number of batches written for the execution benchmarks +const EXEC_BATCHES: usize = 64; +/// Columns in the table scanned by the execution benchmarks +const EXEC_COLUMNS: usize = 128; + +fn schema(num_columns: usize) -> SchemaRef { + Arc::new(Schema::new( + (0..num_columns) + .map(|i| Field::new(format!("c{i}"), DataType::Int32, true)) + .collect::>(), + )) +} + +fn batch(schema: &SchemaRef, num_rows: usize) -> RecordBatch { + let columns: Vec = (0..schema.fields().len()) + .map(|i| { + let values = (0..num_rows).map(|row| (row * i) as i32); + Arc::new(Int32Array::from_iter_values(values)) as ArrayRef + }) + .collect(); + RecordBatch::try_new(Arc::clone(schema), columns).unwrap() +} + +/// Write `num_batches` batches of `num_rows` rows over a `num_columns` wide +/// Int32 schema, and register the result as `t` in a fresh context. +fn context( + rt: &Runtime, + num_columns: usize, + num_rows: usize, + num_batches: usize, +) -> (SessionContext, NamedTempFile) { + let schema = schema(num_columns); + let mut file = Builder::new().suffix(".parquet").tempfile().unwrap(); + let properties = WriterProperties::builder() + .set_max_row_group_row_count(Some(num_rows)) + .build(); + let mut writer = + ArrowWriter::try_new(&mut file, Arc::clone(&schema), Some(properties)).unwrap(); + let batch = batch(&schema, num_rows); + for _ in 0..num_batches { + writer.write(&batch).unwrap(); + } + writer.close().unwrap(); + + let path = file.path().display().to_string(); + assert!(Path::new(&path).exists(), "path not found"); + + let ctx = SessionContext::new(); + rt.block_on(ctx.register_parquet("t", &path, Default::default())) + .unwrap(); + (ctx, file) +} + +fn physical_plan(ctx: &SessionContext, rt: &Runtime, sql: &str) { + black_box(rt.block_on(async { + ctx.sql(sql) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap() + })); +} + +fn collect(ctx: &SessionContext, rt: &Runtime, sql: &str) { + black_box( + rt.block_on(async { ctx.sql(sql).await.unwrap().collect().await.unwrap() }), + ); +} + +fn planning_benchmarks(c: &mut Criterion) { + let rt = Runtime::new().unwrap(); + let mut group = c.benchmark_group("planning"); + group.sample_size(10); + group.warm_up_time(Duration::from_secs(1)); + group.measurement_time(Duration::from_secs(5)); + + for num_columns in [100, 1000] { + let (ctx, _file) = context(&rt, num_columns, 1, 1); + + group.bench_function(format!("select_all_{num_columns}_columns"), |b| { + b.iter(|| physical_plan(&ctx, &rt, "SELECT * FROM t")) + }); + group.bench_function(format!("select_one_of_{num_columns}_columns"), |b| { + b.iter(|| physical_plan(&ctx, &rt, "SELECT c0 FROM t")) + }); + } + + group.finish(); +} + +fn execution_benchmarks(c: &mut Criterion) { + let rt = Runtime::new().unwrap(); + let (ctx, _file) = context(&rt, EXEC_COLUMNS, EXEC_BATCH_ROWS, EXEC_BATCHES); + + let mut group = c.benchmark_group("execution"); + group.sample_size(10); + group.warm_up_time(Duration::from_secs(1)); + group.measurement_time(Duration::from_secs(5)); + + group.bench_function("select_all", |b| { + b.iter(|| collect(&ctx, &rt, "SELECT * FROM t")) + }); + group.bench_function("select_one", |b| { + b.iter(|| collect(&ctx, &rt, "SELECT c0 FROM t")) + }); + + group.finish(); +} + +/// Build the physical scan directly, without the SQL front end: this is the +/// work that is proportional to the table's width, and at these widths it is +/// swamped end-to-end by `SELECT *` expansion in the logical planner. +fn scan_construction_benchmarks(c: &mut Criterion) { + use datafusion::datasource::physical_plan::ParquetSource; + use datafusion_datasource::file_scan_config::FileScanConfigBuilder; + use datafusion_datasource::source::DataSourceExec; + use datafusion_execution::object_store::ObjectStoreUrl; + + let mut group = c.benchmark_group("scan_construction"); + group.sample_size(10); + group.warm_up_time(Duration::from_secs(1)); + group.measurement_time(Duration::from_secs(5)); + + for num_columns in [1_000, 10_000, 100_000] { + let schema = schema(num_columns); + group.bench_function(format!("unprojected_{num_columns}_columns"), |b| { + b.iter(|| { + let source = ParquetSource::new(Arc::clone(&schema)); + let config = FileScanConfigBuilder::new( + ObjectStoreUrl::parse("test:///").unwrap(), + Arc::new(source), + ) + .build(); + black_box(DataSourceExec::from_data_source(config)) + }) + }); + } + + group.finish(); +} + +criterion_group!( + benches, + scan_construction_benchmarks, + planning_benchmarks, + execution_benchmarks +); +criterion_main!(benches); diff --git a/datafusion/datasource-arrow/src/file_format.rs b/datafusion/datasource-arrow/src/file_format.rs index 2bee57ef17581..9ae6eb14f8b89 100644 --- a/datafusion/datasource-arrow/src/file_format.rs +++ b/datafusion/datasource-arrow/src/file_format.rs @@ -38,7 +38,7 @@ use datafusion_common::{ }; use datafusion_common_runtime::{JoinSet, SpawnedTask}; use datafusion_datasource::display::FileGroupDisplay; -use datafusion_datasource::file::FileSource; +use datafusion_datasource::file::{FileSource, projection_is_no_op}; use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; use datafusion_datasource::sink::{DataSink, DataSinkExec}; use datafusion_datasource::write::{ @@ -209,8 +209,10 @@ impl FileFormat for ArrowFormat { Err(e) => Err(e)?, }; - // Preserve projection from the original file source + // Preserve projection from the original file source, skipping one + // that is a no-op over the replacement's output. if let Some(projection) = conf.file_source.projection() + && !projection_is_no_op(source.as_ref(), projection) && let Some(new_source) = source.try_pushdown_projection(projection)? { source = new_source; diff --git a/datafusion/datasource-parquet/src/decoder_projection.rs b/datafusion/datasource-parquet/src/decoder_projection.rs index 89fdc01af4eda..dbf41d9fe7ae0 100644 --- a/datafusion/datasource-parquet/src/decoder_projection.rs +++ b/datafusion/datasource-parquet/src/decoder_projection.rs @@ -57,11 +57,18 @@ use crate::projection_read_plan::build_projection_read_plan; /// boundary) and calls [`Self::map`] on every decoded batch. pub(crate) struct DecoderProjection { projection_mask: ProjectionMask, + /// `None` when the decoder already yields `output_schema` and + /// [`map`](Self::map) is a pass-through. + transform: Option, +} + +/// The per-batch half of a [`DecoderProjection`]. +struct DecoderTransform { projector: Projector, output_schema: SchemaRef, /// `true` when the projector's output schema differs from `output_schema` - /// in metadata / nullability and [`map`](Self::map) must rebuild the batch - /// with `output_schema`. + /// in metadata / nullability and [`DecoderProjection::map`] must rebuild + /// the batch with `output_schema`. replace_schema: bool, } @@ -73,18 +80,52 @@ impl DecoderProjection { /// corresponding parquet [`SchemaDescriptor`]. `output_schema` is what /// consumers of the scan stream expect. /// + /// `projection` is `None` when the scan reads the table unprojected. When + /// the decoder's own output then already matches `output_schema` this + /// installs an all-columns mask and no per-batch transform, avoiding a + /// column-per-field projector on very wide unprojected scans. + /// /// `virtual_state`, when present, describes virtual columns the reader /// will append to each decoded batch (e.g. parquet `row_number`). Virtual /// columns are stripped from the projection fed into /// `build_projection_read_plan` (which only understands file columns) and /// appended to the stream schema so the projector can resolve them. pub(crate) fn try_new( - projection: &ProjectionExprs, + projection: Option<&ProjectionExprs>, physical_file_schema: &SchemaRef, parquet_schema: &SchemaDescriptor, output_schema: &SchemaRef, virtual_state: Option<&VirtualColumnsState>, ) -> Result { + if projection.is_none() { + // Unprojected scan: if the decoder's output (the file schema plus + // any appended virtual columns) already is the output schema, + // there is no transform to apply. + let stream_schema = match virtual_state { + Some(state) => { + append_fields(physical_file_schema, state.virtual_columns()) + } + None => Arc::clone(physical_file_schema), + }; + if stream_schema == *output_schema { + return Ok(Self { + projection_mask: ProjectionMask::all(), + transform: None, + }); + } + } + + // Anything else needs a concrete projection, including an unprojected + // scan whose decoder output does not line up with the output schema. + let materialized_identity; + let projection = match projection { + Some(projection) => projection, + None => { + materialized_identity = ProjectionExprs::identity(output_schema); + &materialized_identity + } + }; + // Virtual columns are produced by the reader separately from the // projection mask, so strip them from the expressions we feed into // `build_projection_read_plan`. We substitute each virtual column @@ -127,9 +168,11 @@ impl DecoderProjection { Ok(Self { projection_mask: read_plan.projection_mask, - projector, - output_schema: Arc::clone(output_schema), - replace_schema, + transform: Some(DecoderTransform { + projector, + output_schema: Arc::clone(output_schema), + replace_schema, + }), }) } @@ -145,17 +188,126 @@ impl DecoderProjection { /// batch with `output_schema` (some writers emit OPTIONAL fields even when /// the data has no nulls; some logical schemas carry field-level metadata /// the file schema does not). - pub(crate) fn map(&self, batch: &RecordBatch) -> Result { - let projected = self.projector.project_batch(batch)?; - if !self.replace_schema { + /// + /// When the decoder already yields the output schema the batch is returned + /// untouched. + pub(crate) fn map(&self, batch: RecordBatch) -> Result { + let Some(transform) = self.transform.as_ref() else { + return Ok(batch); + }; + let projected = transform.projector.project_batch(&batch)?; + if !transform.replace_schema { return Ok(projected); } let (_stream_schema, arrays, num_rows) = projected.into_parts(); let options = RecordBatchOptions::new().with_row_count(Some(num_rows)); Ok(RecordBatch::try_new_with_options( - Arc::clone(&self.output_schema), + Arc::clone(&transform.output_schema), arrays, &options, )?) } } + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::array::{Int32Array, StringArray}; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_physical_expr::expressions::Column; + use datafusion_physical_expr::projection::ProjectionExpr; + use parquet::arrow::ArrowSchemaConverter; + + fn test_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + ])) + } + + fn test_batch(schema: &SchemaRef) -> RecordBatch { + RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec!["x", "y", "z"])), + ], + ) + .unwrap() + } + + #[test] + fn absent_projection_reads_all_columns_without_transform() { + let schema = test_schema(); + let parquet_schema = ArrowSchemaConverter::new().convert(&schema).unwrap(); + + let decoder_projection = + DecoderProjection::try_new(None, &schema, &parquet_schema, &schema, None) + .unwrap(); + + assert!(decoder_projection.transform.is_none()); + assert_eq!(decoder_projection.projection_mask(), &ProjectionMask::all()); + + let batch = test_batch(&schema); + let mapped = decoder_projection.map(batch.clone()).unwrap(); + assert_eq!(mapped, batch); + } + + #[test] + fn absent_projection_falls_back_when_the_output_schema_differs() { + // A file whose column carries metadata the table schema does not: the + // decoder's own output is not the output schema, so the fallback + // materializes the identity rather than passing batches through. + let output_schema = test_schema(); + let physical_file_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true).with_metadata( + std::iter::once(("k".to_string(), "v".to_string())).collect(), + ), + Field::new("b", DataType::Utf8, true), + ])); + let parquet_schema = ArrowSchemaConverter::new() + .convert(&physical_file_schema) + .unwrap(); + + let decoder_projection = DecoderProjection::try_new( + None, + &physical_file_schema, + &parquet_schema, + &output_schema, + None, + ) + .unwrap(); + + assert!(decoder_projection.transform.is_some()); + let mapped = decoder_projection + .map(test_batch(&physical_file_schema)) + .unwrap(); + assert_eq!(mapped.schema(), output_schema); + assert_eq!(mapped, test_batch(&output_schema)); + } + + #[test] + fn narrowing_projection_masks_and_transforms() { + let schema = test_schema(); + let parquet_schema = ArrowSchemaConverter::new().convert(&schema).unwrap(); + let output_schema = + Arc::new(Schema::new(vec![Field::new("b", DataType::Utf8, true)])); + + let projection = ProjectionExprs::new([ProjectionExpr::new( + Arc::new(Column::new("b", 1)), + "b", + )]); + let decoder_projection = DecoderProjection::try_new( + Some(&projection), + &schema, + &parquet_schema, + &output_schema, + None, + ) + .unwrap(); + + assert!(decoder_projection.transform.is_some()); + assert_ne!(decoder_projection.projection_mask(), &ProjectionMask::all()); + } +} diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index 25c3bc9a77851..88a2e9a7ec421 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -238,8 +238,11 @@ fn validate_predicate_does_not_reference_virtual_columns( pub(super) struct ParquetMorselizer { /// Execution partition index pub(crate) partition_index: usize, - /// Projection to apply on top of the table schema (i.e. can reference partition columns). - pub projection: ProjectionExprs, + /// Projection to apply on top of the table schema (i.e. can reference + /// partition columns), or `None` when the output is the table schema + /// itself. `None` lets the scan read the file with an all-columns mask and + /// skip the per-batch transform entirely. + pub projection: Option, /// Target number of rows in each output RecordBatch pub batch_size: usize, /// Optional limit on the number of rows to read @@ -447,7 +450,9 @@ struct PreparedParquetOpen { logical_file_schema: SchemaRef, physical_file_schema: SchemaRef, output_schema: SchemaRef, - projection: ProjectionExprs, + /// `None` carries the same meaning as [`ParquetMorselizer::projection`]: + /// the decoder's own output is already `output_schema`. + projection: Option, predicate: Option>, /// Per-scan virtual-column state, Arc-cloned from [`ParquetMorselizer`] so /// each file shares validated fields, precomputed null replacements, and @@ -760,10 +765,12 @@ impl ParquetMorselizer { // Calculate the output schema from the original projection (before literal replacement) // so we get correct field names from column references let logical_file_schema = Arc::clone(self.table_schema.file_schema()); - let output_schema = Arc::new( - self.projection - .project_schema(self.table_schema.table_schema())?, - ); + let output_schema = match &self.projection { + Some(projection) => { + Arc::new(projection.project_schema(self.table_schema.table_schema())?) + } + None => Arc::clone(self.table_schema.table_schema()), + }; // Build a combined map for replacing column references with literal values. // This includes: @@ -802,16 +809,27 @@ impl ParquetMorselizer { let mut projection = self.projection.clone(); let mut predicate = self.predicate.clone(); if !literal_columns.is_empty() { - projection = projection.try_map_exprs(|expr| { + // Partition and constant columns are not read from the file, so + // the projection must name them as literals; an absent projection + // has to be materialized to hold them. + let concrete = + projection.unwrap_or_else(|| ProjectionExprs::identity(&output_schema)); + projection = Some(concrete.try_map_exprs(|expr| { replace_columns_with_literals(Arc::clone(&expr), &literal_columns) - })?; + })?); predicate = predicate .map(|p| replace_columns_with_literals(p, &literal_columns)) .transpose()?; } - // Replace any `input_file_name()` UDFs in the projection with a literal for this file. - projection = rewrite_input_file_name_in_projection(projection, &file_name)?; + // Replace any `input_file_name()` UDFs in the projection with a literal + // for this file. An absent projection is all plain column references, + // so there is nothing to rewrite. + projection = projection + .map(|projection| { + rewrite_input_file_name_in_projection(projection, &file_name) + }) + .transpose()?; let predicate_creation_errors = MetricBuilder::new(&self.metrics) .with_category(MetricCategory::Rows) @@ -1034,8 +1052,8 @@ impl MetadataLoadedParquetOpen { // columns are appended after file columns in the table schema), // types are the same, and there are no missing columns. Skip the // tree walk entirely in that case. - let needs_rewrite = prepared.predicate.is_some() - || prepared.logical_file_schema != physical_file_schema; + let schemas_differ = prepared.logical_file_schema != physical_file_schema; + let needs_rewrite = prepared.predicate.is_some() || schemas_differ; if needs_rewrite { // When virtual columns are requested, augment the logical and // physical schemas passed to the rewriter/simplifier with those @@ -1067,9 +1085,19 @@ impl MetadataLoadedParquetOpen { .predicate .map(|p| simplifier.simplify(rewriter.rewrite(p)?)) .transpose()?; - prepared.projection = prepared - .projection - .try_map_exprs(|p| simplifier.simplify(rewriter.rewrite(p)?))?; + let projection = match prepared.projection.take() { + // The rewriter leaves an absent projection alone when the + // logical and physical file schemas agree: indices already + // line up, there is nothing to cast, and no column is missing. + None if !schemas_differ => None, + // Otherwise materialize the identity so the casts and null + // fills land somewhere. + None => Some(ProjectionExprs::identity(&prepared.output_schema)), + projection => projection, + }; + prepared.projection = projection + .map(|p| p.try_map_exprs(|e| simplifier.simplify(rewriter.rewrite(e)?))) + .transpose()?; } prepared.physical_file_schema = Arc::clone(&physical_file_schema); @@ -1453,7 +1481,7 @@ impl RowGroupsPrunedParquetOpen { // opener's orchestration body focused on filter / decoder / stream // wiring. let decoder_projection = DecoderProjection::try_new( - &prepared.projection, + prepared.projection.as_ref(), &prepared.physical_file_schema, reader_metadata.parquet_schema(), &prepared.output_schema, @@ -2276,15 +2304,11 @@ mod test { ); let file_schema = Arc::clone(table_schema.file_schema()); - let projection = if let Some(projection) = self.projection { - projection - } else if let Some(indices) = self.projection_indices { - ProjectionExprs::from_indices(&indices, &file_schema) - } else { - // Default: project all columns - let all_indices: Vec = (0..file_schema.fields().len()).collect(); - ProjectionExprs::from_indices(&all_indices, &file_schema) - }; + // Default: no projection, i.e. the whole table. + let projection = self.projection.or_else(|| { + self.projection_indices + .map(|indices| ProjectionExprs::from_indices(&indices, &file_schema)) + }); let virtual_state = build_virtual_columns_state( table_schema.virtual_columns(), @@ -3099,6 +3123,69 @@ mod test { assert_eq!(num_rows, 0); } + #[tokio::test] + async fn test_unprojected_scan_fills_partition_columns() { + let store = Arc::new(InMemory::new()) as Arc; + + let batch = record_batch!( + ("a", Int32, vec![Some(1), Some(2), Some(3)]), + ("b", Float64, vec![Some(1.0), Some(2.0), None]) + ) + .unwrap(); + let data_size = + write_parquet(Arc::clone(&store), "part=7/file.parquet", batch.clone()).await; + let file_schema = batch.schema(); + let mut file = PartitionedFile::new( + "part=7/file.parquet".to_string(), + u64::try_from(data_size).unwrap(), + ); + file.partition_values = vec![ScalarValue::Int32(Some(7))]; + + let table_schema = TableSchemaBuilder::from(&file_schema) + .with_table_partition_cols(vec![Arc::new(Field::new( + "part", + DataType::Int32, + false, + ))]) + .build(); + + // No projection: the output is the whole table schema, [a, b, part], + // with the partition column materialized per file. + let opener = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_table_schema(table_schema) + .build(); + + let mut stream = open_file(&opener, file).await.unwrap(); + let mut batches = vec![]; + while let Some(batch) = stream.next().await { + batches.push(batch.unwrap()); + } + assert_eq!(batches.len(), 1); + let batch = &batches[0]; + assert_eq!( + batch + .schema() + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect::>(), + vec!["a", "b", "part"] + ); + let a = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(a, &arrow::array::Int32Array::from(vec![1, 2, 3])); + let part = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(part, &arrow::array::Int32Array::from(vec![7, 7, 7])); + } + /// Test that if the filter is not a dynamic filter and we have no stats we don't do extra pruning work at the file level. #[tokio::test] async fn test_opener_pruning_skipped_on_static_filters() { @@ -4005,6 +4092,38 @@ mod test { assert_eq!(row_numbers, vec![0, 1, 2, 3]); } + #[tokio::test] + async fn test_row_index_unprojected_scan() { + let store = Arc::new(InMemory::new()) as Arc; + let (file_schema, data_size) = + write_grouped_file(&store, "unprojected.parquet", 1, 5).await; + + let rn_field = row_number_field("row_number", false); + let table_schema = TableSchemaBuilder::new(Arc::clone(&file_schema)) + .with_virtual_columns(vec![Arc::clone(&rn_field)]) + .build(); + + // No projection: the output is the table schema itself, + // [value, row_number], with the reader appending the virtual + // column to each decoded batch. + let morselizer = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_table_schema(table_schema) + .build(); + + let file = PartitionedFile::new( + "unprojected.parquet".to_string(), + u64::try_from(data_size).unwrap(), + ); + let stream = open_file(&morselizer, file.clone()).await.unwrap(); + let values = collect_int64_values(stream, 0).await; + assert_eq!(values, vec![0, 1, 2, 3, 4]); + + let stream = open_file(&morselizer, file).await.unwrap(); + let row_numbers = collect_int64_values(stream, 1).await; + assert_eq!(row_numbers, vec![0, 1, 2, 3, 4]); + } + #[tokio::test] async fn test_input_file_name_projection() { let store = Arc::new(InMemory::new()) as Arc; diff --git a/datafusion/datasource-parquet/src/push_decoder.rs b/datafusion/datasource-parquet/src/push_decoder.rs index 94a0e5049fef6..b2147bb8461e1 100644 --- a/datafusion/datasource-parquet/src/push_decoder.rs +++ b/datafusion/datasource-parquet/src/push_decoder.rs @@ -442,7 +442,7 @@ impl PushDecoderStreamState { Some(Ok(batch)) => { let mut timer = self.baseline_metrics.elapsed_compute().timer(); self.copy_arrow_reader_metrics(); - let result = self.project_batch(&batch); + let result = self.project_batch(batch); timer.stop(); drop(timer); return Some((result, self)); @@ -695,7 +695,7 @@ impl PushDecoderStreamState { } } - fn project_batch(&self, batch: &RecordBatch) -> Result { + fn project_batch(&self, batch: RecordBatch) -> Result { self.decoder_projection.map(batch) } } diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index 087da503654d7..0214e8d9d7f1c 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -303,8 +303,9 @@ pub struct ParquetSource { pub(crate) batch_size: Option, /// Optional hint for the size of the parquet metadata pub(crate) metadata_size_hint: Option, - /// Projection to apply to the output. - pub(crate) projection: ProjectionExprs, + /// Projection to apply to the output, or `None` when the output is the + /// table schema itself (file columns + partition columns). + pub(crate) projection: Option, #[cfg(feature = "parquet_encryption")] pub(crate) encryption_factory: Option>, /// If true, the opener flips row-group iteration order. Within- @@ -323,13 +324,9 @@ impl ParquetSource { /// Uses default `TableParquetOptions`. /// To set custom options, use [ParquetSource::with_table_parquet_options`]. pub fn new(table_schema: impl Into) -> Self { - let table_schema = table_schema.into(); - // Projection over the full table schema (file columns + partition columns) - let full_schema = table_schema.table_schema(); - let indices: Vec = (0..full_schema.fields().len()).collect(); Self { - projection: ProjectionExprs::from_indices(&indices, full_schema), - table_schema, + projection: None, + table_schema: table_schema.into(), table_parquet_options: TableParquetOptions::default(), metrics: ExecutionPlanMetricsSet::new(), predicate: None, @@ -699,7 +696,11 @@ impl FileSource for ParquetSource { if !projection.iter().any(|projection_expr| { expr_references_scalar_udf::(&projection_expr.expr) }) { - source.projection = self.projection.try_merge(projection)?; + let table_schema = self.table_schema.table_schema(); + source.projection = Some(match &self.projection { + Some(existing) => existing.try_merge(projection)?, + None => projection.try_merge_onto_identity(table_schema)?, + }); return Ok(Some(Arc::new(source))); } @@ -708,18 +709,29 @@ impl FileSource for ParquetSource { let (table_schema, row_index_col) = table_schema_with_row_index_col(self.table_schema()); - source.table_schema = table_schema; - source.projection = rewrite_file_row_index_projection( - &self.projection, + // `rewrite_file_row_index_projection` appends the row-index column to + // the base projection, so it needs a concrete one. + let materialized_identity; + let base = match &self.projection { + Some(existing) => existing, + None => { + materialized_identity = + ProjectionExprs::identity(self.table_schema.table_schema()); + &materialized_identity + } + }; + source.projection = Some(rewrite_file_row_index_projection( + base, projection, &row_index_col, - )?; + )?); + source.table_schema = table_schema; Ok(Some(Arc::new(source))) } fn projection(&self) -> Option<&ProjectionExprs> { - Some(&self.projection) + self.projection.as_ref() } fn metrics(&self) -> &ExecutionPlanMetricsSet { @@ -1065,9 +1077,12 @@ impl FileSource for ParquetSource { ) -> datafusion_common::Result, ) -> datafusion_common::Result { datafusion_physical_plan::apply_expression_roots( - self.predicate - .iter() - .chain(self.projection.iter().map(|proj_expr| &proj_expr.expr)), + self.predicate.iter().chain( + self.projection + .iter() + .flat_map(|projection| projection.iter()) + .map(|proj_expr| &proj_expr.expr), + ), f, ) } @@ -1279,6 +1294,91 @@ mod tests { assert_eq!(parquet_source.predicate(), parquet_source.filter().as_ref()); } + #[test] + fn test_new_source_has_no_projection() { + use arrow::datatypes::{DataType, Field}; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let source = ParquetSource::new(Arc::clone(&schema)); + assert!(source.projection().is_none()); + + // Callers filter no-op projections (see `projection_is_no_op`), so + // only projections that change something reach the source. + let narrowed = ProjectionExprs::from_indices(&[1], &schema); + let pushed = source.try_pushdown_projection(&narrowed).unwrap().unwrap(); + assert_eq!(pushed.projection(), Some(&narrowed)); + + let reordered = ProjectionExprs::from_indices(&[1, 0], &schema); + let pushed = source.try_pushdown_projection(&reordered).unwrap().unwrap(); + assert_eq!(pushed.projection(), Some(&reordered)); + } + + #[test] + fn test_row_index_pushdown_onto_an_existing_projection() { + // `file_row_index()` pushed onto a source that already carries a + // projection must append the row-index column to *that* projection, + // not to the table's identity: the incoming expressions are indexed + // against the source's current output, not against the table. + use arrow::datatypes::{DataType, Field}; + use datafusion_expr::col; + use datafusion_functions::core::expr_fn::file_row_index; + use datafusion_physical_expr::planner::logical2physical; + use datafusion_physical_expr::projection::ProjectionExpr; + + let table_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Utf8, true), + ])); + let source = ParquetSource::new(Arc::clone(&table_schema)); + + // Narrow to `b`, so the source's output is a single column. + let narrowed = ProjectionExprs::from_indices(&[1], &table_schema); + let projected = source.try_pushdown_projection(&narrowed).unwrap().unwrap(); + assert_eq!(projected.projection(), Some(&narrowed)); + + // `SELECT b, file_row_index()` over that output: `b` is at index 0 here. + let output_schema = Schema::new(vec![Field::new("b", DataType::Utf8, true)]); + let pushed = ProjectionExprs::new([ + ProjectionExpr::new(logical2physical(&col("b"), &output_schema), "b"), + ProjectionExpr::new( + logical2physical(&file_row_index(), &output_schema), + "ri", + ), + ]); + let with_row_index = projected.try_pushdown_projection(&pushed).unwrap().unwrap(); + + let projection = with_row_index.projection().expect("projection is stored"); + assert_eq!( + projection + .iter() + .map(|expr| expr.alias.as_str()) + .collect::>(), + vec!["b", "ri"], + ); + + // Built on the existing projection: `b` still resolves to the table's + // index 1. Merging onto the table identity instead would have left `a` + // in the base and resolved `b` elsewhere. + let b = projection.as_ref()[0] + .expr + .downcast_ref::() + .expect("b stays a plain column reference"); + assert_eq!((b.name(), b.index()), ("b", 1)); + + // The row-index column is served by a virtual column on the rewritten + // table schema rather than read from the file. + assert!( + with_row_index + .table_schema() + .virtual_columns() + .iter() + .any(|f| f.name() == "__datafusion_file_row_index"), + ); + } + #[test] fn test_reverse_scan_default_value() { use arrow::datatypes::Schema; diff --git a/datafusion/datasource/src/file.rs b/datafusion/datasource/src/file.rs index f1a94f2e12363..61589aae0ff20 100644 --- a/datafusion/datasource/src/file.rs +++ b/datafusion/datasource/src/file.rs @@ -46,6 +46,28 @@ pub fn as_file_source(source: T) -> Arc Arc::new(source) } +/// Returns `true` when pushing `projection` into `source` would leave the +/// source's output exactly as it is: one column per output field, in order, +/// under the name that field already has. +/// +/// Callers of [`FileSource::try_pushdown_projection`] skip such projections: +/// storing one makes every consumer of [`FileSource::projection`] do work +/// proportional to the table's width, which dominates scan construction for +/// very wide tables. +pub fn projection_is_no_op( + source: &dyn FileSource, + projection: &ProjectionExprs, +) -> bool { + match source.projection() { + // A projected source's output fields are its projection's aliases. + Some(existing) => projection.is_identity_over_names( + existing.as_ref().iter().map(|expr| expr.alias.as_str()), + ), + // An unprojected source's output is its table schema. + None => projection.is_identity(source.table_schema().table_schema()), + } +} + /// File format specific behaviors for [`DataSource`] /// /// # Schema information @@ -117,6 +139,11 @@ pub trait FileSource: Any + Send + Sync { /// Return the projection that will be applied to the output stream on top /// of [`Self::table_schema`]. /// + /// `None` means the output *is* [`Self::table_schema`]: every field, in + /// order, under its own name. Callers of [`Self::try_pushdown_projection`] + /// use [`projection_is_no_op`] to avoid storing a projection that changes + /// nothing, so consumers can skip projection work when this is `None`. + /// /// Note you can use [`ProjectionExprs::project_schema`] on the table /// schema to get the effective output schema of this source. fn projection(&self) -> Option<&ProjectionExprs> { diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs index 4e72b5e83bd23..a3c5b31942f48 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -30,9 +30,9 @@ mod proto; use crate::file_groups::FileGroup; use crate::{ PartitionedFile, display::FileGroupsDisplay, file::FileSource, - file_compression_type::FileCompressionType, file_stream::FileStreamBuilder, - file_stream::work_source::SharedWorkSource, source::DataSource, - statistics::MinMaxStatistics, + file::projection_is_no_op, file_compression_type::FileCompressionType, + file_stream::FileStreamBuilder, file_stream::work_source::SharedWorkSource, + source::DataSource, statistics::MinMaxStatistics, }; use arrow::datatypes::Fields; use arrow::datatypes::{DataType, Schema, SchemaRef}; @@ -397,6 +397,9 @@ impl FileScanConfigBuilder { let Some(projection_exprs) = projection_exprs else { return Ok(self); }; + if projection_is_no_op(self.file_source.as_ref(), &projection_exprs) { + return Ok(self); + } let new_source = self .file_source .try_pushdown_projection(&projection_exprs) @@ -961,7 +964,12 @@ impl DataSource for FileScanConfig { projection.project_statistics(stat.clone(), &output_schema)?, )) } else { - Ok(Arc::new(stat.clone())) + // No projection: column statistics carry over as-is, but + // recompute the total byte size from the schema as + // `project_statistics` would. + let mut stat = stat.clone(); + stat.calculate_total_byte_size(&output_schema); + Ok(Arc::new(stat)) }; } // If no statistics available for this partition, return unknown @@ -970,14 +978,15 @@ impl DataSource for FileScanConfig { ))) } else { // Return aggregate statistics across all partitions - let statistics = self.statistics(); + let mut statistics = self.statistics(); let projection = self.file_source.projection(); let output_schema = self.projected_schema()?; if let Some(projection) = &projection { Ok(Arc::new( - projection.project_statistics(statistics.clone(), &output_schema)?, + projection.project_statistics(statistics, &output_schema)?, )) } else { + statistics.calculate_total_byte_size(&output_schema); Ok(Arc::new(statistics)) } } @@ -1013,6 +1022,13 @@ impl DataSource for FileScanConfig { { return Ok(None); } + // A no-op projection: report success with the scan unchanged so the + // caller drops the `ProjectionExec` without the source storing + // anything. `remove_unnecessary_projections` normally strips these + // before they reach a `DataSource`, so this is rarely hit. + if projection_is_no_op(self.file_source.as_ref(), projection) { + return Ok(Some(Arc::new(self.clone()) as Arc)); + } match self.file_source.try_pushdown_projection(projection)? { Some(new_source) => { let mut new_file_scan_config = self.clone(); @@ -2337,6 +2353,95 @@ mod tests { .build() } + #[test] + fn test_projection_is_no_op() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let table_schema = TableSchema::from(&schema); + + // `MockSource` reports the identity projection; the sort-pushdown + // source reports none. Either way the source's output is the table. + let projected: Arc = + Arc::new(MockSource::new(table_schema.clone())); + let unprojected: Arc = + Arc::new(InexactSortPushdownSource::new(table_schema.clone())); + assert!(projected.projection().is_some()); + assert!(unprojected.projection().is_none()); + + let identity = ProjectionExprs::identity(&schema); + let narrowed = ProjectionExprs::from_indices(&[1], &schema); + let reordered = ProjectionExprs::from_indices(&[1, 0], &schema); + for source in [&projected, &unprojected] { + assert!(projection_is_no_op(source.as_ref(), &identity)); + assert!(!projection_is_no_op(source.as_ref(), &narrowed)); + assert!(!projection_is_no_op(source.as_ref(), &reordered)); + } + + // A source that already reorders: only a projection reproducing *that* + // output changes nothing, and the table's own identity does not. + let swapped = projected + .try_pushdown_projection(&reordered) + .unwrap() + .unwrap(); + let over_swapped = ProjectionExprs::new([ + ProjectionExpr::new(Arc::new(Column::new("b", 0)), "b"), + ProjectionExpr::new(Arc::new(Column::new("a", 1)), "a"), + ]); + assert!(projection_is_no_op(swapped.as_ref(), &over_swapped)); + assert!(!projection_is_no_op(swapped.as_ref(), &identity)); + } + + #[test] + fn test_full_width_projection_indices_are_no_op() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let table_schema = TableSchema::from(&schema); + let source: Arc = + Arc::new(InexactSortPushdownSource::new(table_schema)); + + let config = FileScanConfigBuilder::new( + ObjectStoreUrl::parse("test:///").unwrap(), + Arc::clone(&source), + ) + .with_projection_indices(Some(vec![0, 1])) + .unwrap() + .build(); + + assert!(config.file_source().projection().is_none()); + assert_eq!(config.projected_schema().unwrap(), schema); + } + + #[test] + fn test_projection_swap_reports_success_for_no_op() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let table_schema = TableSchema::from(&schema); + let source: Arc = + Arc::new(InexactSortPushdownSource::new(table_schema)); + let config = FileScanConfigBuilder::new( + ObjectStoreUrl::parse("test:///").unwrap(), + source, + ) + .build(); + + // The swap reports success, so the caller drops its ProjectionExec, + // while the source stores nothing. + let identity = ProjectionExprs::identity(&schema); + let pushed = config + .try_swapping_with_projection(&identity) + .unwrap() + .expect("a no-op projection is reported as pushed"); + let pushed = pushed.downcast_ref::().unwrap(); + assert!(pushed.file_source().projection().is_none()); + assert_eq!(pushed.projected_schema().unwrap(), schema); + } + #[test] fn test_file_scan_config_builder() { let file_schema = aggr_test_schema(); @@ -2909,6 +3014,86 @@ mod tests { assert_eq!(partition_stats.total_byte_size, Precision::Exact(800)); } + #[test] + fn test_partition_statistics_no_projection_matches_identity() { + // A source returning `None` from `projection()` and one carrying the + // identity describe the same output, so their statistics must match, + // including `total_byte_size`, which the projected path recomputes + // from the output schema. + use crate::source::DataSourceExec; + use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; + + let schema = Arc::new(Schema::new(vec![ + Field::new("col0", DataType::Int32, false), + Field::new("col1", DataType::Int32, false), + ])); + let table_schema = TableSchema::from(&schema); + + let file_group_stats = Statistics { + num_rows: Precision::Exact(100), + total_byte_size: Precision::Exact(1024), + column_statistics: vec![ + ColumnStatistics { + null_count: Precision::Exact(0), + ..ColumnStatistics::new_unknown() + }, + ColumnStatistics { + null_count: Precision::Exact(5), + ..ColumnStatistics::new_unknown() + }, + ], + }; + let file_group = FileGroup::new(vec![PartitionedFile::new("test.parquet", 1024)]) + .with_statistics(Arc::new(file_group_stats.clone())); + + let statistics_of = |source: Arc, projected: bool| { + let mut builder = FileScanConfigBuilder::new( + ObjectStoreUrl::parse("test:///").unwrap(), + source, + ); + if projected { + builder = builder.with_projection_indices(Some(vec![0, 1])).unwrap(); + } + let config = builder + .with_file_groups(vec![file_group.clone()]) + .with_statistics(file_group_stats.clone()) + .build(); + let exec = DataSourceExec::from_data_source(config); + let context = StatisticsContext::new(); + ( + context + .compute(exec.as_ref(), &StatisticsArgs::new().with_partition(None)) + .unwrap(), + context + .compute( + exec.as_ref(), + &StatisticsArgs::new().with_partition(Some(0)), + ) + .unwrap(), + ) + }; + + // `InexactSortPushdownSource` does not override `projection()`, so it + // takes the `None` path; `MockSource` returns the identity projection. + let (unprojected_all, unprojected_partition) = statistics_of( + Arc::new(InexactSortPushdownSource::new(table_schema.clone())), + false, + ); + let (projected_all, projected_partition) = + statistics_of(Arc::new(MockSource::new(table_schema)), true); + + // 100 rows * 2 Int32 columns * 4 bytes + assert_eq!(unprojected_all.total_byte_size, Precision::Exact(800)); + assert_eq!( + unprojected_all.total_byte_size, + projected_all.total_byte_size + ); + assert_eq!( + unprojected_partition.total_byte_size, + projected_partition.total_byte_size + ); + } + #[test] fn test_statistics_with_filter() { assert_num_rows_with_filter(Precision::Absent, Precision::Absent); diff --git a/datafusion/datasource/src/file_scan_config/proto.rs b/datafusion/datasource/src/file_scan_config/proto.rs index 4071ea471b33f..a0018a2c8e0fa 100644 --- a/datafusion/datasource/src/file_scan_config/proto.rs +++ b/datafusion/datasource/src/file_scan_config/proto.rs @@ -51,7 +51,7 @@ use datafusion_physical_plan::proto::{ExecutionPlanDecodeCtx, ExecutionPlanEncod use datafusion_proto_models::datafusion_common::CompressionTypeVariant as ProtoCompressionTypeVariant; use datafusion_proto_models::protobuf; -use crate::file::FileSource; +use crate::file::{FileSource, projection_is_no_op}; use crate::file_compression_type::FileCompressionType; use crate::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; use crate::table_schema::TableSchema; @@ -243,9 +243,15 @@ impl FileScanConfig { let projection_exprs = ProjectionExprs::new(projection_exprs); - file_source - .try_pushdown_projection(&projection_exprs)? - .unwrap_or(file_source) + // Plans encoded before no-op projections were dropped may still + // carry one; skip it rather than store it on the source. + if projection_is_no_op(file_source.as_ref(), &projection_exprs) { + file_source + } else { + file_source + .try_pushdown_projection(&projection_exprs)? + .unwrap_or(file_source) + } } else { file_source }; diff --git a/datafusion/physical-expr/src/projection.rs b/datafusion/physical-expr/src/projection.rs index 0e8876f017379..0a0298ad92973 100644 --- a/datafusion/physical-expr/src/projection.rs +++ b/datafusion/physical-expr/src/projection.rs @@ -246,6 +246,50 @@ impl ProjectionExprs { Self::from_iter(projection_exprs) } + /// Creates the identity projection over `schema`: every field, in order, + /// under its own name. + pub fn identity(schema: &Schema) -> Self { + let projection_exprs = + schema + .fields() + .iter() + .enumerate() + .map(|(index, field)| ProjectionExpr { + expr: Arc::new(Column::new(field.name(), index)), + alias: field.name().clone(), + }); + + Self::from_iter(projection_exprs) + } + + /// Returns `true` if this projection is the identity over `schema`: every + /// field of `schema`, in order, under its own name. Applying it produces + /// the input unchanged. + pub fn is_identity(&self, schema: &Schema) -> bool { + self.is_identity_over_names( + schema.fields().iter().map(|field| field.name().as_str()), + ) + } + + /// [`Self::is_identity`] against an input described by its field names + /// alone, for callers whose input schema is not materialized as a + /// [`Schema`] (e.g. the output of another projection). + pub fn is_identity_over_names<'a>( + &self, + names: impl ExactSizeIterator, + ) -> bool { + self.exprs.len() == names.len() + && self.exprs.iter().zip(names).enumerate().all( + |(index, (proj_expr, name))| { + proj_expr.alias == name + && proj_expr + .expr + .downcast_ref::() + .is_some_and(|col| col.index() == index && col.name() == name) + }, + ) + } + /// Returns an iterator over the projection expressions pub fn iter(&self) -> impl Iterator { self.exprs.iter() @@ -389,6 +433,52 @@ impl ProjectionExprs { Ok(ProjectionExprs::new(new_exprs)) } + /// The result of [`ProjectionExprs::identity(schema).try_merge(self)`], + /// computed without materializing the identity. + /// + /// Merging onto the identity leaves every expression alone except for + /// column references whose name disagrees with the field `schema` carries + /// at that index, which are renamed. The cost is therefore proportional to + /// `self` rather than to the width of `schema`. + /// + /// [`ProjectionExprs::identity(schema).try_merge(self)`]: Self::try_merge + /// + /// # Errors + /// This function returns an error if any column reference is out of bounds + /// for `schema`, matching [`Self::try_merge`] against the identity. + pub fn try_merge_onto_identity(&self, schema: &Schema) -> Result { + let fields = schema.fields(); + let mut new_exprs = Vec::with_capacity(self.exprs.len()); + for proj_expr in self.exprs.iter() { + let expr = Arc::clone(&proj_expr.expr) + .transform_up(|expr| { + let Some(column) = expr.downcast_ref::() else { + return Ok(Transformed::no(expr)); + }; + let field = fields.get(column.index()).ok_or_else(|| { + internal_datafusion_err!( + "Column index {} out of bounds for projected expressions of length {}", + column.index(), + fields.len() + ) + })?; + if field.name() == column.name() { + return Ok(Transformed::no(expr)); + } + Ok(Transformed::yes(Arc::new(Column::new( + field.name(), + column.index(), + )) as _)) + }) + .data()?; + new_exprs.push(ProjectionExpr { + expr, + alias: proj_expr.alias.clone(), + }); + } + Ok(ProjectionExprs::new(new_exprs)) + } + /// Extract the column indices used in this projection. /// For example, for a projection `SELECT a AS x, b + 1 AS y`, where `a` is at index 0 and `b` is at index 1, /// this function would return `[0, 1]`. @@ -2422,6 +2512,109 @@ pub(crate) mod tests { // Tests for Projection struct + #[test] + fn test_identity_is_identity() { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, true), + ]); + assert!(ProjectionExprs::identity(&schema).is_identity(&schema)); + assert_eq!( + ProjectionExprs::identity(&schema), + ProjectionExprs::from_indices(&[0, 1], &schema) + ); + } + + #[test] + fn test_is_identity_rejects_non_identity_projections() { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, true), + ]); + + // Fewer columns + assert!(!ProjectionExprs::from_indices(&[0], &schema).is_identity(&schema)); + // Reordered + assert!(!ProjectionExprs::from_indices(&[1, 0], &schema).is_identity(&schema)); + // Duplicated + assert!(!ProjectionExprs::from_indices(&[0, 0], &schema).is_identity(&schema)); + // Renamed + assert!( + !ProjectionExprs::new([ + ProjectionExpr::new(Arc::new(Column::new("a", 0)), "renamed"), + ProjectionExpr::new(Arc::new(Column::new("b", 1)), "b"), + ]) + .is_identity(&schema) + ); + // Not a plain column reference + assert!( + !ProjectionExprs::new([ + ProjectionExpr::new( + Arc::new(Literal::new(ScalarValue::Int32(Some(1)))), + "a" + ), + ProjectionExpr::new(Arc::new(Column::new("b", 1)), "b"), + ]) + .is_identity(&schema) + ); + } + + #[test] + fn test_try_merge_onto_identity_matches_try_merge() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Utf8, true), + ]); + let identity = ProjectionExprs::identity(&schema); + + let cases = [ + // Plain column selection. + ProjectionExprs::from_indices(&[2, 0], &schema), + // Everything, in order. + ProjectionExprs::from_indices(&[0, 1, 2], &schema), + // A computed expression over two columns, plus a literal. + ProjectionExprs::new([ + ProjectionExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Plus, + Arc::new(Column::new("b", 1)), + )), + "sum", + ), + ProjectionExpr::new( + Arc::new(Literal::new(ScalarValue::Int32(Some(7)))), + "seven", + ), + ]), + // A column whose name disagrees with the schema at that index: + // merging onto the identity renames it. + ProjectionExprs::new([ProjectionExpr::new( + Arc::new(Column::new("stale", 1)), + "renamed", + )]), + ]; + + for projection in cases { + assert_eq!( + projection.try_merge_onto_identity(&schema)?, + identity.try_merge(&projection)?, + "mismatch for {projection}" + ); + } + + // Out of bounds fails the same way as merging onto the identity does. + let out_of_bounds = ProjectionExprs::new([ProjectionExpr::new( + Arc::new(Column::new("d", 3)), + "d", + )]); + assert!(out_of_bounds.try_merge_onto_identity(&schema).is_err()); + assert!(identity.try_merge(&out_of_bounds).is_err()); + + Ok(()) + } + #[test] fn test_projection_new() -> Result<()> { let exprs = vec![ diff --git a/datafusion/proto/tests/cases/plans/sources.rs b/datafusion/proto/tests/cases/plans/sources.rs index c17ff47e0f472..2da99531cce50 100644 --- a/datafusion/proto/tests/cases/plans/sources.rs +++ b/datafusion/proto/tests/cases/plans/sources.rs @@ -197,6 +197,71 @@ fn roundtrip_parquet_exec_attaches_cached_reader_factory_after_roundtrip() -> Re Ok(()) } +#[test] +fn roundtrip_parquet_exec_decodes_identity_projection_to_none() -> Result<()> { + use datafusion::physical_expr::projection::ProjectionExprs; + use datafusion_datasource::file::FileSource; + + let file_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, false), + ])); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + + let roundtrip_projection_is_none = + |file_source: Arc| -> Result { + let scan_config = FileScanConfigBuilder::new( + ObjectStoreUrl::local_filesystem(), + file_source, + ) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.parquet".to_string(), + 1024, + )])]) + .build(); + let roundtripped = roundtrip_test_and_return( + DataSourceExec::from_data_source(scan_config), + &ctx, + &codec, + &proto_converter, + )?; + let file_scan = roundtripped + .downcast_ref::() + .and_then(|exec| exec.data_source().downcast_ref::()) + .ok_or_else(|| { + internal_datafusion_err!("Expected FileScanConfig after roundtrip") + })?; + Ok(file_scan.file_source().projection().is_none()) + }; + + // A source that never stored a projection round-trips to none. + let plain: Arc = + Arc::new(ParquetSource::new(Arc::clone(&file_schema))); + assert!(plain.projection().is_none()); + + // Plans encoded before sources stopped storing no-op projections carry an + // explicit identity projection; build a source holding one directly (the + // no-op filtering lives at `try_pushdown_projection`'s call sites, not in + // the method itself). + let with_identity = plain + .try_pushdown_projection(&ProjectionExprs::identity(&file_schema))? + .expect("parquet source accepts projection pushdown"); + assert!(with_identity.projection().is_some()); + + assert!( + roundtrip_projection_is_none(plain)?, + "a projection-less source must decode back without a projection" + ); + assert!( + roundtrip_projection_is_none(with_identity)?, + "a legacy identity projection must be dropped on decode, not pushed back into the source" + ); + Ok(()) +} + /// Returns `FileSource::file_type` of a `DataSourceExec` file scan, e.g. /// "arrow" vs "arrow_stream". The two Arrow IPC formats print identically in /// plan debug output, so roundtrip tests must inspect the source directly.