diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 454af28c14b4a..30a9bdd49195c 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1181,6 +1181,14 @@ config_namespace! { /// (reading) Use any available bloom filters when reading parquet files pub bloom_filter_on_read: bool, default = true + /// (reading) If true, when a projected nested column (struct, or + /// struct nested in lists) is read through a cast to a narrower + /// nested type — as happens when the table schema declares fewer + /// struct fields than the parquet file contains — the reader only + /// fetches and decodes the leaf columns the cast retains, instead of + /// the entire column + pub nested_projection_pruning: bool, default = true + /// (reading) The maximum predicate cache size, in bytes. When /// `pushdown_filters` is enabled, sets the maximum memory used to cache /// the results of predicate evaluation between filter evaluation and diff --git a/datafusion/common/src/file_options/parquet_writer.rs b/datafusion/common/src/file_options/parquet_writer.rs index 320bfcf33e488..76e791d20a0b2 100644 --- a/datafusion/common/src/file_options/parquet_writer.rs +++ b/datafusion/common/src/file_options/parquet_writer.rs @@ -242,6 +242,7 @@ impl ParquetOptions { maximum_parallel_row_group_writers: _, maximum_buffered_record_batches_per_stream: _, bloom_filter_on_read: _, // reads not used for writer props + nested_projection_pruning: _, // reads not used for writer props schema_force_view_types: _, binary_as_string: _, // not used for writer props coerce_int96: _, // not used for writer props @@ -500,6 +501,7 @@ mod tests { maximum_buffered_record_batches_per_stream: defaults .maximum_buffered_record_batches_per_stream, bloom_filter_on_read: defaults.bloom_filter_on_read, + nested_projection_pruning: defaults.nested_projection_pruning, schema_force_view_types: defaults.schema_force_view_types, binary_as_string: defaults.binary_as_string, skip_arrow_metadata: defaults.skip_arrow_metadata, @@ -620,6 +622,8 @@ mod tests { maximum_buffered_record_batches_per_stream: global_options_defaults .maximum_buffered_record_batches_per_stream, bloom_filter_on_read: global_options_defaults.bloom_filter_on_read, + nested_projection_pruning: global_options_defaults + .nested_projection_pruning, max_predicate_cache_size: global_options_defaults .max_predicate_cache_size, schema_force_view_types: global_options_defaults.schema_force_view_types, diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index 60cff658a6a97..10435220a1f0d 100644 --- a/datafusion/core/Cargo.toml +++ b/datafusion/core/Cargo.toml @@ -247,6 +247,11 @@ harness = false name = "parquet_struct_query" required-features = ["parquet"] +[[bench]] +harness = false +name = "parquet_nested_schema_pruning" +required-features = ["parquet"] + [[bench]] harness = false name = "parquet_struct_projection" diff --git a/datafusion/core/benches/parquet_nested_schema_pruning.rs b/datafusion/core/benches/parquet_nested_schema_pruning.rs new file mode 100644 index 0000000000000..39640da078601 --- /dev/null +++ b/datafusion/core/benches/parquet_nested_schema_pruning.rs @@ -0,0 +1,442 @@ +// 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 schema-driven nested projection pruning in Parquet. +//! +//! A table's declared (logical) schema can be *narrower* than the physical +//! parquet type of a nested column — e.g. the table declares +//! `events: LIST>` while the file contains +//! `events: LIST>`. Engines like Spark +//! communicate nested projection pruning to the scan exactly this way +//! (a clipped read schema), so the reader should fetch and decode only the +//! leaves the declared schema names. +//! +//! Each dataset shape is measured three ways: +//! +//! 1. **narrow_schema**: wide file, narrow declared table schema — the +//! interesting case; ideally close to (3) +//! 2. **full_schema**: wide file, full table schema — the cost of reading +//! everything +//! 3. **physically_narrow**: a file that only contains the narrow columns — +//! the floor +//! +//! At setup the benchmark prints the `bytes_scanned` metric for (1) and (2) +//! so the IO pattern is visible in addition to wall time. + +use arrow::array::{ + ArrayRef, Int32Array, Int64Array, ListArray, StringArray, StructArray, +}; +use arrow::buffer::OffsetBuffer; +use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use arrow::util::pretty::pretty_format_batches; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion::datasource::listing::{ + ListingTable, ListingTableConfig, ListingTableConfigExt, +}; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_datasource::ListingTableUrl; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::{WriterProperties, WriterVersion}; +use std::hint::black_box; +use std::sync::Arc; +use std::time::Duration; +use tempfile::NamedTempFile; +use tokio::runtime::Runtime; + +const NUM_BATCHES: usize = 2; +const ROWS_PER_BATCH: usize = 256; +const ROW_GROUP_ROW_COUNT: usize = 256; +const ELEMS_PER_ROW: usize = 3; +const NUM_PAD_FIELDS: usize = 8; +const PAD_LEN: usize = 2048; + +/// The narrow item fields: the subset of the struct the table declares. +fn narrow_item_fields() -> Fields { + Fields::from(vec![ + Field::new("x", DataType::Int64, true), + Field::new("y", DataType::Utf8, true), + ]) +} + +/// The wide item fields as written to the file: the narrow fields plus +/// `NUM_PAD_FIELDS` fat string fields the table schema does not mention. +fn wide_item_fields() -> Fields { + let mut fields = vec![ + Field::new("x", DataType::Int64, false), + Field::new("y", DataType::Utf8, true), + ]; + for i in 0..NUM_PAD_FIELDS { + fields.push(Field::new(format!("pad_{i}"), DataType::Utf8, false)); + } + Fields::from(fields) +} + +fn list_schema(item_fields: Fields) -> SchemaRef { + let item = Arc::new(Field::new("item", DataType::Struct(item_fields), true)); + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("events", DataType::List(item), true), + ])) +} + +fn struct_schema(item_fields: Fields) -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("s", DataType::Struct(item_fields), true), + ])) +} + +/// Distinct pad values so dictionary encoding cannot collapse them. +fn pad_values(count: usize, seed: usize) -> ArrayRef { + let base = "x".repeat(PAD_LEN); + let values: Vec = (0..count) + .map(|i| format!("{:08}{base}", seed + i)) + .collect(); + Arc::new(StringArray::from(values)) +} + +/// Struct children for `count` elements, restricted to `fields`. +fn item_columns(fields: &Fields, count: usize, seed: usize) -> Vec { + fields + .iter() + .enumerate() + .map(|(i, field)| match field.name().as_str() { + "x" => Arc::new(Int64Array::from_iter_values( + (0..count).map(|j| (seed + j) as i64), + )) as ArrayRef, + "y" => Arc::new(StringArray::from_iter_values( + (0..count).map(|j| format!("y-{}", seed + j)), + )) as ArrayRef, + _ => pad_values(count, seed * (i + 1)), + }) + .collect() +} + +fn list_batch(fields: &Fields, batch_id: usize) -> RecordBatch { + let num_elems = ROWS_PER_BATCH * ELEMS_PER_ROW; + let seed = batch_id * num_elems; + let struct_array = + StructArray::new(fields.clone(), item_columns(fields, num_elems, seed), None); + let item = Arc::new(Field::new("item", DataType::Struct(fields.clone()), true)); + let events = ListArray::new( + item, + OffsetBuffer::from_lengths(std::iter::repeat_n(ELEMS_PER_ROW, ROWS_PER_BATCH)), + Arc::new(struct_array), + None, + ); + let ids = Int32Array::from_iter_values( + (0..ROWS_PER_BATCH).map(|i| (batch_id * ROWS_PER_BATCH + i) as i32), + ); + RecordBatch::try_new( + list_schema(fields.clone()), + vec![Arc::new(ids), Arc::new(events)], + ) + .unwrap() +} + +fn struct_batch(fields: &Fields, batch_id: usize) -> RecordBatch { + let seed = batch_id * ROWS_PER_BATCH; + let struct_array = StructArray::new( + fields.clone(), + item_columns(fields, ROWS_PER_BATCH, seed), + None, + ); + let ids = + Int32Array::from_iter_values((0..ROWS_PER_BATCH).map(|i| (seed + i) as i32)); + RecordBatch::try_new( + struct_schema(fields.clone()), + vec![Arc::new(ids), Arc::new(struct_array)], + ) + .unwrap() +} + +fn generate_file( + schema: SchemaRef, + batch_fn: impl Fn(usize) -> RecordBatch, + prefix: &str, +) -> NamedTempFile { + let mut named_file = tempfile::Builder::new() + .prefix(prefix) + .suffix(".parquet") + .tempfile() + .unwrap(); + + let properties = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_2_0) + .set_dictionary_enabled(false) + .set_max_row_group_row_count(Some(ROW_GROUP_ROW_COUNT)) + .build(); + + let mut writer = + ArrowWriter::try_new(&mut named_file, schema, Some(properties)).unwrap(); + for batch_id in 0..NUM_BATCHES { + writer.write(&batch_fn(batch_id)).unwrap(); + } + let metadata = writer.close().unwrap(); + println!( + "Generated {} ({} rows, {} row groups, {} bytes)", + named_file.path().display(), + metadata.file_metadata().num_rows(), + metadata.row_groups().len(), + std::fs::metadata(named_file.path()).unwrap().len(), + ); + named_file +} + +/// Register `path` as `table`, declaring `table_schema` (which may be narrower +/// than the file's physical schema). +fn register_table( + ctx: &SessionContext, + rt: &Runtime, + table: &str, + path: &str, + table_schema: SchemaRef, +) { + let url = ListingTableUrl::parse(path).unwrap(); + let config = rt + .block_on(ListingTableConfig::new(url).infer_options(&ctx.state())) + .unwrap() + .with_schema(table_schema); + let provider = ListingTable::try_new(config).unwrap(); + ctx.register_table(table, Arc::new(provider)).unwrap(); +} + +fn query(ctx: &SessionContext, rt: &Runtime, sql: &str) { + let df = rt.block_on(ctx.sql(sql)).unwrap(); + black_box(rt.block_on(df.collect()).unwrap()); +} + +/// Print the parquet scan's `bytes_scanned` for a query so the IO pattern is +/// visible alongside the wall-time measurements. +fn report_bytes_scanned(ctx: &SessionContext, rt: &Runtime, label: &str, sql: &str) { + let df = rt + .block_on(ctx.sql(&format!("EXPLAIN ANALYZE {sql}"))) + .unwrap(); + let batches = rt.block_on(df.collect()).unwrap(); + let text = pretty_format_batches(&batches).unwrap().to_string(); + let bytes_scanned = text + .split("bytes_scanned=") + .nth(1) + .map(|rest| { + rest.split([',', ']']) + .next() + .unwrap_or("") + .trim() + .to_string() + }) + .unwrap_or_else(|| "".to_string()); + println!("{label}: bytes_scanned={bytes_scanned} ({sql})"); +} + +struct Fixture { + ctx: SessionContext, + /// Same tables, `nested_projection_pruning` disabled: the pre-pruning + /// behavior, kept so the before/after stays visible. + ctx_pruning_disabled: SessionContext, + rt: Runtime, + _files: Vec, +} + +/// Tables: +/// `_narrow_schema`: wide file, narrow declared schema +/// `_full_schema`: wide file, full declared schema +/// `_physically_narrow`: narrow file, narrow declared schema +fn setup( + name: &str, + schema_fn: fn(Fields) -> SchemaRef, + batch_fn: fn(&Fields, usize) -> RecordBatch, +) -> Fixture { + let rt = Runtime::new().unwrap(); + let ctx = SessionContext::new(); + let mut cfg_off = SessionConfig::new(); + cfg_off + .options_mut() + .execution + .parquet + .nested_projection_pruning = false; + let ctx_pruning_disabled = SessionContext::new_with_config(cfg_off); + + let wide = wide_item_fields(); + let narrow = narrow_item_fields(); + + let wide_file = generate_file(schema_fn(wide.clone()), |i| batch_fn(&wide, i), name); + let narrow_file = generate_file( + schema_fn(narrow.clone()), + |i| batch_fn(&narrow, i), + &format!("{name}_narrow"), + ); + let wide_path = wide_file.path().display().to_string(); + let narrow_path = narrow_file.path().display().to_string(); + + register_table( + &ctx, + &rt, + &format!("{name}_narrow_schema"), + &wide_path, + schema_fn(narrow.clone()), + ); + register_table( + &ctx, + &rt, + &format!("{name}_full_schema"), + &wide_path, + schema_fn(wide.clone()), + ); + register_table( + &ctx, + &rt, + &format!("{name}_physically_narrow"), + &narrow_path, + schema_fn(narrow.clone()), + ); + register_table( + &ctx_pruning_disabled, + &rt, + &format!("{name}_narrow_schema"), + &wide_path, + schema_fn(narrow.clone()), + ); + + Fixture { + ctx, + ctx_pruning_disabled, + rt, + _files: vec![wide_file, narrow_file], + } +} + +fn list_struct_benchmarks(c: &mut Criterion) { + let f = setup("list_struct", list_schema, list_batch); + let (ctx, rt) = (&f.ctx, &f.rt); + + for table in [ + "list_struct_narrow_schema", + "list_struct_full_schema", + "list_struct_physically_narrow", + ] { + report_bytes_scanned(ctx, rt, table, &format!("SELECT events FROM {table}")); + } + report_bytes_scanned( + &f.ctx_pruning_disabled, + rt, + "list_struct_narrow_schema (pruning disabled)", + "SELECT events FROM list_struct_narrow_schema", + ); + + let mut group = c.benchmark_group("list_struct"); + group.sample_size(10); + group.warm_up_time(Duration::from_secs(1)); + group.measurement_time(Duration::from_secs(3)); + + // wide file, narrow declared schema: should only read the narrow leaves + group.bench_function("select_events_narrow_schema", |b| { + b.iter(|| query(ctx, rt, "SELECT events FROM list_struct_narrow_schema")) + }); + + // wide file, full schema: the cost of reading everything + group.bench_function("select_events_full_schema", |b| { + b.iter(|| query(ctx, rt, "SELECT events FROM list_struct_full_schema")) + }); + + // narrow file: the floor + group.bench_function("select_events_physically_narrow", |b| { + b.iter(|| query(ctx, rt, "SELECT events FROM list_struct_physically_narrow")) + }); + + // wide file, narrow schema, pruning disabled: the pre-pruning behavior + group.bench_function("select_events_narrow_schema_pruning_disabled", |b| { + b.iter(|| { + query( + &f.ctx_pruning_disabled, + rt, + "SELECT events FROM list_struct_narrow_schema", + ) + }) + }); + + // aggregation over one narrow leaf through unnest + group.bench_function("sum_x_narrow_schema", |b| { + b.iter(|| { + query( + ctx, + rt, + "SELECT SUM(e['x']) FROM (SELECT UNNEST(events) AS e FROM list_struct_narrow_schema)", + ) + }) + }); + + group.finish(); +} + +fn top_level_struct_benchmarks(c: &mut Criterion) { + let f = setup("struct", struct_schema, struct_batch); + let (ctx, rt) = (&f.ctx, &f.rt); + + for table in [ + "struct_narrow_schema", + "struct_full_schema", + "struct_physically_narrow", + ] { + report_bytes_scanned(ctx, rt, table, &format!("SELECT s FROM {table}")); + } + report_bytes_scanned( + &f.ctx_pruning_disabled, + rt, + "struct_narrow_schema (pruning disabled)", + "SELECT s FROM struct_narrow_schema", + ); + + let mut group = c.benchmark_group("top_level_struct"); + group.sample_size(10); + group.warm_up_time(Duration::from_secs(1)); + group.measurement_time(Duration::from_secs(3)); + + group.bench_function("select_struct_narrow_schema", |b| { + b.iter(|| query(ctx, rt, "SELECT s FROM struct_narrow_schema")) + }); + + group.bench_function("select_struct_full_schema", |b| { + b.iter(|| query(ctx, rt, "SELECT s FROM struct_full_schema")) + }); + + group.bench_function("select_struct_physically_narrow", |b| { + b.iter(|| query(ctx, rt, "SELECT s FROM struct_physically_narrow")) + }); + + // wide file, narrow schema, pruning disabled: the pre-pruning behavior + group.bench_function("select_struct_narrow_schema_pruning_disabled", |b| { + b.iter(|| { + query( + &f.ctx_pruning_disabled, + rt, + "SELECT s FROM struct_narrow_schema", + ) + }) + }); + + // get_field on a schema-narrowed struct column: the expression-level + // pruning path interacting with the schema-level narrowing + group.bench_function("sum_x_narrow_schema", |b| { + b.iter(|| query(ctx, rt, "SELECT SUM(s['x']) FROM struct_narrow_schema")) + }); + + group.finish(); +} + +criterion_group!(benches, list_struct_benchmarks, top_level_struct_benchmarks); +criterion_main!(benches); diff --git a/datafusion/core/tests/parquet/expr_adapter.rs b/datafusion/core/tests/parquet/expr_adapter.rs index fd70d74a9140c..024ee877b239d 100644 --- a/datafusion/core/tests/parquet/expr_adapter.rs +++ b/datafusion/core/tests/parquet/expr_adapter.rs @@ -1126,3 +1126,392 @@ async fn test_physical_expr_adapter_factory_reuse_across_tables() { ]; assert_batches_eq!(expected, &batches); } + +// --------------------------------------------------------------------------- +// Nested projection pruning: when the table schema declares a nested column +// narrower than the physical parquet file, the scan should only read the +// leaves the declared schema names (`datafusion.execution.parquet. +// nested_projection_pruning`), instead of reading the whole column and +// discarding the extra subfields in the adapter-inserted cast. +// --------------------------------------------------------------------------- + +mod nested_projection_pruning { + use super::*; + use arrow::array::{Array, Int64Array}; + use arrow::buffer::NullBuffer; + use datafusion::physical_plan::collect; + use datafusion_physical_plan::metrics::MetricsSet; + + use crate::parquet::utils::MetricsFinder; + + const NUM_ELEMENTS: usize = 64; + const PAD_LEN: usize = 2048; + + /// Physical item struct written to the file: the narrow fields plus fat + /// pads the table schema will not mention. `x` is Int32 in the file (the + /// table declares Int64 to also exercise leaf promotion). + fn wide_item_fields() -> Fields { + Fields::from(vec![ + Field::new("x", DataType::Int32, false), + Field::new("y", DataType::Utf8, true), + Field::new("pad_a", DataType::Utf8, false), + Field::new("pad_b", DataType::Utf8, false), + Field::new("pad_c", DataType::Utf8, false), + ]) + } + + /// The narrow item struct the table declares: a subset of the physical + /// fields in a different order, a promoted leaf type for `x`, plus `z` + /// which does not exist in the file (null-filled by the cast). + fn narrow_item_fields() -> Fields { + Fields::from(vec![ + Field::new("y", DataType::Utf8, true), + Field::new("x", DataType::Int64, true), + Field::new("z", DataType::Int64, true), + ]) + } + + fn wide_struct_values(validity: Option) -> StructArray { + let pad = |seed: usize| { + let base = "x".repeat(PAD_LEN); + Arc::new(StringArray::from_iter_values( + (0..NUM_ELEMENTS).map(|i| format!("{}{base}", seed + i)), + )) as ArrayRef + }; + StructArray::new( + wide_item_fields(), + vec![ + Arc::new(Int32Array::from_iter_values(0..NUM_ELEMENTS as i32)), + Arc::new(StringArray::from_iter_values( + (0..NUM_ELEMENTS).map(|i| format!("y-{i}")), + )), + pad(1000), + pad(2000), + pad(3000), + ], + validity, + ) + } + + /// File batch: `id Int32`, `events List` (one element per row). + fn wide_list_batch() -> RecordBatch { + let item = Arc::new(Field::new( + "item", + DataType::Struct(wide_item_fields()), + true, + )); + let events = ListArray::new( + Arc::clone(&item), + OffsetBuffer::from_lengths(std::iter::repeat_n(1, NUM_ELEMENTS)), + Arc::new(wide_struct_values(None)), + None, + ); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("events", DataType::List(item), true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from_iter_values(0..NUM_ELEMENTS as i32)), + Arc::new(events), + ], + ) + .unwrap() + } + + fn narrow_list_table_schema() -> SchemaRef { + let item = Arc::new(Field::new( + "item", + DataType::Struct(narrow_item_fields()), + true, + )); + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("events", DataType::List(item), true), + ])) + } + + /// File batch: `id Int32`, `s `, with per-row struct + /// validity so struct-level nullability can be asserted. + fn wide_struct_batch() -> RecordBatch { + // rows 0, 10, 20, ... have a NULL struct + let validity = + NullBuffer::from((0..NUM_ELEMENTS).map(|i| i % 10 != 0).collect::>()); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("s", DataType::Struct(wide_item_fields()), true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from_iter_values(0..NUM_ELEMENTS as i32)), + Arc::new(wide_struct_values(Some(validity))), + ], + ) + .unwrap() + } + + fn narrow_struct_table_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("s", DataType::Struct(narrow_item_fields()), true), + ])) + } + + async fn setup( + batches: Vec<(&str, RecordBatch)>, + table_schema: SchemaRef, + pruning_enabled: bool, + ) -> SessionContext { + let store = Arc::new(InMemory::new()) as Arc; + for (name, batch) in batches { + write_parquet(batch, Arc::clone(&store), &format!("data/{name}")).await; + } + let mut cfg = SessionConfig::new().with_collect_statistics(false); + cfg.options_mut() + .execution + .parquet + .nested_projection_pruning = pruning_enabled; + let ctx = SessionContext::new_with_config(cfg); + register_memory_listing_table(&ctx, store, "memory:///data/", table_schema).await; + ctx + } + + async fn run(ctx: &SessionContext, sql: &str) -> (Vec, MetricsSet) { + let df = ctx.sql(sql).await.unwrap(); + let (state, logical) = df.into_parts(); + let plan = state.create_physical_plan(&logical).await.unwrap(); + let batches = collect(Arc::clone(&plan), state.task_ctx()).await.unwrap(); + let metrics = MetricsFinder::find_metrics(plan.as_ref()).unwrap(); + (batches, metrics) + } + + fn bytes_scanned(metrics: &MetricsSet) -> usize { + metrics + .sum(|m| m.value().name() == "bytes_scanned") + .map(|v| v.as_usize()) + .expect("bytes_scanned metric") + } + + /// Run `sql` with pruning on and off against the same data; assert + /// identical results and that pruning read strictly less than half of + /// the unpruned bytes (the pads dominate the file). + async fn assert_prunes( + batches: Vec<(&str, RecordBatch)>, + table_schema: SchemaRef, + sql: &str, + ) -> Vec { + let ctx_on = setup(batches.clone(), Arc::clone(&table_schema), true).await; + let ctx_off = setup(batches, table_schema, false).await; + + let (result_on, metrics_on) = run(&ctx_on, sql).await; + let (result_off, metrics_off) = run(&ctx_off, sql).await; + + let on = concat_batches(&result_on[0].schema(), &result_on).unwrap(); + let off = concat_batches(&result_off[0].schema(), &result_off).unwrap(); + assert_eq!(on, off, "results must not change under pruning: {sql}"); + + let (on_bytes, off_bytes) = + (bytes_scanned(&metrics_on), bytes_scanned(&metrics_off)); + assert!( + on_bytes * 2 < off_bytes, + "expected pruned scan to read less than half of {off_bytes} bytes, \ + read {on_bytes}: {sql}" + ); + result_on + } + + #[tokio::test] + async fn prunes_list_of_struct() { + // Narrow schema over the wide file: subset of fields, reordered, + // promoted leaf (x: Int32 -> Int64), missing subfield z null-filled. + let batches = assert_prunes( + vec![("wide.parquet", wide_list_batch())], + narrow_list_table_schema(), + "SELECT events FROM t ORDER BY id", + ) + .await; + + let events = batches[0].column(0); + let list = events.as_any().downcast_ref::().unwrap(); + let items = list + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(items.fields().len(), 3); + let x = items + .column_by_name("x") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(x.value(5), 5); + let z = items.column_by_name("z").unwrap(); + assert_eq!(z.null_count(), z.len(), "z is not in the file"); + } + + #[tokio::test] + async fn prunes_top_level_struct() { + assert_prunes( + vec![("wide.parquet", wide_struct_batch())], + narrow_struct_table_schema(), + "SELECT s FROM t ORDER BY id", + ) + .await; + } + + /// Struct-level nullability must survive the clip: rows where the struct + /// itself is NULL stay NULL (not `{y: NULL, x: NULL, z: NULL}`). + #[tokio::test] + async fn preserves_struct_nullability() { + let batches = assert_prunes( + vec![("wide.parquet", wide_struct_batch())], + narrow_struct_table_schema(), + "SELECT id, s IS NULL AS s_null, s FROM t ORDER BY id", + ) + .await; + + let combined = concat_batches(&batches[0].schema(), &batches).unwrap(); + let s_null = combined + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..NUM_ELEMENTS { + assert_eq!(s_null.value(i), i % 10 == 0, "row {i}"); + } + } + + /// `get_field` on a schema-narrowed struct becomes + /// `get_field(CAST(s), 'x')`; the read clips to the cast target. + #[tokio::test] + async fn prunes_get_field_on_narrowed_struct() { + let batches = assert_prunes( + vec![("wide.parquet", wide_struct_batch())], + narrow_struct_table_schema(), + "SELECT s['x'] AS x FROM t ORDER BY id", + ) + .await; + let combined = concat_batches(&batches[0].schema(), &batches).unwrap(); + let x = combined + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(x.value(5), 5); + assert_eq!(x.value(NUM_ELEMENTS - 1), NUM_ELEMENTS as i64 - 1); + } + + /// Mixed access: the whole (narrowed) column and a subfield of it. + #[tokio::test] + async fn prunes_mixed_struct_and_subfield_access() { + assert_prunes( + vec![("wide.parquet", wide_struct_batch())], + narrow_struct_table_schema(), + "SELECT s, s['y'] AS y FROM t ORDER BY id", + ) + .await; + } + + /// Predicate on a primitive column with filter pushdown enabled while + /// the projected nested column is clipped. + #[tokio::test] + async fn prunes_with_filter_pushdown() { + let setup_with_pushdown = |enabled| async move { + let store = Arc::new(InMemory::new()) as Arc; + write_parquet(wide_list_batch(), Arc::clone(&store), "data/wide.parquet") + .await; + let mut cfg = SessionConfig::new().with_collect_statistics(false); + cfg.options_mut().execution.parquet.pushdown_filters = true; + cfg.options_mut() + .execution + .parquet + .nested_projection_pruning = enabled; + let ctx = SessionContext::new_with_config(cfg); + register_memory_listing_table( + &ctx, + store, + "memory:///data/", + narrow_list_table_schema(), + ) + .await; + ctx + }; + + let sql = "SELECT events FROM t WHERE id >= 32 ORDER BY id"; + let (result_on, metrics_on) = run(&setup_with_pushdown(true).await, sql).await; + let (result_off, metrics_off) = run(&setup_with_pushdown(false).await, sql).await; + + let on = concat_batches(&result_on[0].schema(), &result_on).unwrap(); + let off = concat_batches(&result_off[0].schema(), &result_off).unwrap(); + assert_eq!(on, off); + assert_eq!(on.num_rows(), NUM_ELEMENTS / 2); + assert!(bytes_scanned(&metrics_on) * 2 < bytes_scanned(&metrics_off)); + } + + /// A scan over two files where one matches the table schema exactly (no + /// cast is inserted) and one is wider (clipped): both must be read + /// correctly in the same scan. + #[tokio::test] + async fn mixed_files_narrow_and_wide() { + // The physically-narrow file has exactly the table's item struct. + let narrow_item = narrow_item_fields(); + let item = Arc::new(Field::new( + "item", + DataType::Struct(narrow_item.clone()), + true, + )); + let events = ListArray::new( + Arc::clone(&item), + OffsetBuffer::from_lengths([1]), + Arc::new(StructArray::new( + narrow_item, + vec![ + Arc::new(StringArray::from(vec![Some("y-narrow")])) as ArrayRef, + Arc::new(Int64Array::from(vec![Some(4242)])) as ArrayRef, + Arc::new(Int64Array::from(vec![Some(7)])) as ArrayRef, + ], + None, + )), + None, + ); + let narrow_batch = RecordBatch::try_new( + narrow_list_table_schema(), + vec![ + Arc::new(Int32Array::from(vec![NUM_ELEMENTS as i32])), + Arc::new(events), + ], + ) + .unwrap(); + + let ctx = setup( + vec![ + ("wide.parquet", wide_list_batch()), + ("narrow.parquet", narrow_batch), + ], + narrow_list_table_schema(), + true, + ) + .await; + + let (batches, _) = run( + &ctx, + "SELECT id, e['x'] AS x, e['z'] AS z \ + FROM (SELECT id, unnest(events) AS e FROM t) ORDER BY id", + ) + .await; + let combined = concat_batches(&batches[0].schema(), &batches).unwrap(); + assert_eq!(combined.num_rows(), NUM_ELEMENTS + 1); + let x = combined + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(x.value(NUM_ELEMENTS), 4242, "row from the narrow file"); + let z = combined.column(2); + // z is null-filled for the wide file, present in the narrow file + assert_eq!(z.null_count(), NUM_ELEMENTS); + } +} diff --git a/datafusion/datasource-parquet/src/decoder_projection.rs b/datafusion/datasource-parquet/src/decoder_projection.rs index 192600ce7a607..ee44c1194a504 100644 --- a/datafusion/datasource-parquet/src/decoder_projection.rs +++ b/datafusion/datasource-parquet/src/decoder_projection.rs @@ -45,7 +45,7 @@ use parquet::arrow::ProjectionMask; use parquet::schema::types::SchemaDescriptor; use crate::opener::{VirtualColumnsState, append_fields}; -use crate::row_filter::build_projection_read_plan; +use crate::projection_read_plan::build_projection_read_plan; /// Per-file decoder projection: the [`ProjectionMask`] installed on the /// parquet decoder, plus the per-batch transform that maps the decoder's @@ -84,6 +84,7 @@ impl DecoderProjection { parquet_schema: &SchemaDescriptor, output_schema: &SchemaRef, virtual_state: Option<&VirtualColumnsState>, + nested_projection_pruning: bool, ) -> Result { // Virtual columns are produced by the reader separately from the // projection mask, so strip them from the expressions we feed into @@ -101,6 +102,7 @@ impl DecoderProjection { projection_for_read_plan.expr_iter(), physical_file_schema, parquet_schema, + nested_projection_pruning, ); // The reader produces projected file columns followed by any virtual diff --git a/datafusion/datasource-parquet/src/mod.rs b/datafusion/datasource-parquet/src/mod.rs index e6e372cd788f1..35f831230b305 100644 --- a/datafusion/datasource-parquet/src/mod.rs +++ b/datafusion/datasource-parquet/src/mod.rs @@ -30,8 +30,10 @@ mod decoder_projection; pub mod file_format; pub mod metadata; mod metrics; +mod nested_schema_pruning; mod opener; mod page_filter; +mod projection_read_plan; mod push_decoder; mod reader; mod row_filter; diff --git a/datafusion/datasource-parquet/src/nested_schema_pruning.rs b/datafusion/datasource-parquet/src/nested_schema_pruning.rs new file mode 100644 index 0000000000000..aaef776e6d6c1 --- /dev/null +++ b/datafusion/datasource-parquet/src/nested_schema_pruning.rs @@ -0,0 +1,560 @@ +// 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. + +//! Schema-driven nested projection pruning. +//! +//! When a scan's projection consumes a nested column only through a cast to a +//! *narrower* nested type — e.g. the file contains +//! `events: List>` but the expression is +//! `CAST(events AS List>)` — the parquet reader does not need to +//! fetch or decode the leaves the cast target never names. This module +//! computes which parquet leaves survive such a cast by walking the physical +//! and target type trees in parallel, matching struct fields by name (the +//! equivalent of Spark's `ParquetReadSupport.clipParquetSchema`). +//! +//! This situation arises whenever a table's logical schema declares a nested +//! column narrower than the physical parquet file: the physical expression +//! adapter rewrites the projected column into exactly such a whole-column +//! cast (see `datafusion_physical_expr_adapter`). Engines like Spark +//! communicate nested projection pruning to the scan this way — as a clipped +//! read *schema* rather than as `get_field` expressions. +//! +//! # Safety of clipping +//! +//! The runtime cast for nested types +//! ([`datafusion_common::nested_struct::cast_column`]) consumes source struct +//! children exclusively by looking up the *target* field names, recursively +//! through list wrappers. Physical subtrees not named by the target are +//! provably dead: removing them from the read cannot change the cast's +//! output. Struct-level nullability is preserved because the parquet reader +//! reconstructs ancestor validity from the definition levels of any surviving +//! leaf, and the clip always keeps at least one leaf per struct level. +//! +//! The clip is *total*: any type shape it does not understand (maps, +//! dictionaries, wrapper-kind mismatches, ...) keeps all of its leaves, so +//! the worst case is today's behavior of reading the full column. Map values +//! are deliberately not clipped: the runtime cast routes maps through Arrow's +//! positional struct cast, which requires all children to be present. + +use std::collections::BTreeSet; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, FieldRef}; + +/// Clip `physical` against `cast_target`, returning which parquet leaves the +/// cast actually consumes, as offsets relative to the root column's first +/// leaf (sorted ascending and non-empty). The Arrow type the reader will +/// emit for these offsets is derived with [`prune_type_by_kept_offsets`]. +/// +/// Returns `None` when nothing can be pruned (every leaf is consumed, or the +/// shapes do not allow safe clipping), in which case the caller should read +/// the whole column as before. This function never fails: unknown shapes +/// degrade to keeping all leaves. +pub(crate) fn clip_for_cast( + physical: &DataType, + cast_target: &DataType, +) -> Option> { + let total = count_leaves(physical); + let mut kept = Vec::new(); + let mut next_leaf = 0; + clip_type(physical, cast_target, &mut next_leaf, &mut kept); + debug_assert_eq!(next_leaf, total, "leaf accounting must cover the type"); + if kept.is_empty() || kept.len() >= total { + return None; + } + Some(kept) +} + +/// Number of parquet leaf columns a (parquet-derived) Arrow type occupies. +pub(crate) fn count_leaves(dt: &DataType) -> usize { + match dt { + DataType::Struct(fields) => { + fields.iter().map(|f| count_leaves(f.data_type())).sum() + } + DataType::List(f) + | DataType::LargeList(f) + | DataType::ListView(f) + | DataType::LargeListView(f) + | DataType::FixedSizeList(f, _) + | DataType::Map(f, _) => count_leaves(f.data_type()), + DataType::Dictionary(_, value) => count_leaves(value), + DataType::RunEndEncoded(_, value) => count_leaves(value.data_type()), + _ => 1, + } +} + +/// Recursive walker: advances `next_leaf` across every leaf of `physical`, +/// pushing the offsets the cast target consumes into `kept`. +fn clip_type( + physical: &DataType, + target: &DataType, + next_leaf: &mut usize, + kept: &mut Vec, +) { + match (physical, target) { + (DataType::Struct(p_children), DataType::Struct(t_children)) => { + let matches: Vec> = p_children + .iter() + .map(|pc| t_children.iter().find(|tc| tc.name() == pc.name())) + .collect(); + + if matches.iter().all(Option::is_none) { + // No field name overlap at this level. The runtime cast will + // fail (or a custom cast may do something unusual), so keep + // the first non-empty child wholesale to preserve at least + // one leaf and the struct's definition levels; behavior is + // then identical to an unclipped read. + return keep_first_child(p_children, next_leaf, kept); + } + + for (pc, tc) in p_children.iter().zip(matches) { + match tc { + Some(tc) => { + clip_type(pc.data_type(), tc.data_type(), next_leaf, kept) + } + None => skip_leaves(pc.data_type(), next_leaf), + } + } + } + (DataType::List(p_item), DataType::List(t_item)) + | (DataType::LargeList(p_item), DataType::LargeList(t_item)) => { + clip_type(p_item.data_type(), t_item.data_type(), next_leaf, kept) + } + // Anything else — leaf pairs, wrapper-kind mismatches, maps, + // dictionaries, fixed-size lists, views — is kept wholesale. + _ => keep_all_leaves(physical, next_leaf, kept), + } +} + +/// Fallback for a struct level with zero name overlap: keep the first child +/// that owns at least one leaf, skip the rest. +fn keep_first_child( + children: &arrow::datatypes::Fields, + next_leaf: &mut usize, + kept: &mut Vec, +) { + let mut kept_one = false; + for child in children { + if !kept_one && count_leaves(child.data_type()) > 0 { + keep_all_leaves(child.data_type(), next_leaf, kept); + kept_one = true; + } else { + skip_leaves(child.data_type(), next_leaf); + } + } +} + +fn keep_all_leaves(dt: &DataType, next_leaf: &mut usize, kept: &mut Vec) { + let n = count_leaves(dt); + kept.extend(*next_leaf..*next_leaf + n); + *next_leaf += n; +} + +fn skip_leaves(dt: &DataType, next_leaf: &mut usize) { + *next_leaf += count_leaves(dt); +} + +/// Derive the Arrow type the reader emits for `physical` when only the leaf +/// offsets in `kept` (relative to this type's first leaf) are selected. +/// +/// Used when a root column's kept set is the *union* of several accesses +/// (e.g. a narrowing cast plus `get_field` accesses), where no single cast +/// target describes the result. Subtrees that own no kept leaf are removed; +/// struct levels that lose all children are removed entirely (`None`). +pub(crate) fn prune_type_by_kept_offsets( + physical: &DataType, + kept: &BTreeSet, +) -> Option { + let mut next_leaf = 0; + prune_type_walker(physical, kept, &mut next_leaf) +} + +fn prune_type_walker( + physical: &DataType, + kept: &BTreeSet, + next_leaf: &mut usize, +) -> Option { + match physical { + DataType::Struct(children) => { + let kept_children: Vec = children + .iter() + .filter_map(|c| { + prune_type_walker(c.data_type(), kept, next_leaf) + .map(|t| Arc::new(c.as_ref().clone().with_data_type(t)) as _) + }) + .collect(); + (!kept_children.is_empty()).then(|| DataType::Struct(kept_children.into())) + } + DataType::List(item) => prune_type_walker(item.data_type(), kept, next_leaf) + .map(|t| DataType::List(Arc::new(item.as_ref().clone().with_data_type(t)))), + DataType::LargeList(item) => prune_type_walker(item.data_type(), kept, next_leaf) + .map(|t| { + DataType::LargeList(Arc::new(item.as_ref().clone().with_data_type(t))) + }), + other => { + // Opaque subtree: all-or-nothing. Kept sets are always aligned to + // whole subtrees below struct levels, so partial overlap cannot + // occur; any kept leaf implies the whole subtree was kept. + let n = count_leaves(other); + let start = *next_leaf; + *next_leaf += n; + kept.range(start..start + n) + .next() + .is_some() + .then(|| other.clone()) + } + } +} + +/// A projected root column that is consumed through a cast to a narrower +/// nested type (`CAST(col AS target_type)`), recorded during projection +/// analysis. +#[derive(Debug, Clone)] +pub(crate) struct CastColumnAccess { + /// Arrow root column index of the column in the file schema. + pub(crate) root_index: usize, + /// The cast's target type. + pub(crate) target_type: DataType, +} + +/// Rebuild `field` with a new data type, preserving name, nullability and +/// metadata. +pub(crate) fn field_with_type(field: &Field, data_type: DataType) -> FieldRef { + Arc::new(field.clone().with_data_type(data_type)) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::datatypes::Fields; + + /// Emitted type for kept offsets, via the same derivation production + /// uses. + fn emitted(physical: &DataType, kept: &[usize]) -> DataType { + let set: BTreeSet = kept.iter().copied().collect(); + prune_type_by_kept_offsets(physical, &set).unwrap() + } + + fn utf8(name: &str) -> Field { + Field::new(name, DataType::Utf8, true) + } + + fn int64(name: &str) -> Field { + Field::new(name, DataType::Int64, true) + } + + fn struct_of(fields: Vec) -> DataType { + DataType::Struct(Fields::from(fields)) + } + + fn list_of(item: DataType) -> DataType { + DataType::List(Arc::new(Field::new("item", item, true))) + } + + #[test] + fn count_leaves_shapes() { + assert_eq!(count_leaves(&DataType::Int32), 1); + assert_eq!(count_leaves(&struct_of(vec![utf8("a"), int64("b")])), 2); + assert_eq!( + count_leaves(&list_of(struct_of(vec![ + utf8("a"), + struct_of(vec![int64("x"), int64("y")]).into_field("s") + ]))), + 3 + ); + let map = DataType::Map( + Arc::new(Field::new( + "entries", + struct_of(vec![utf8("key"), int64("value")]), + false, + )), + false, + ); + assert_eq!(count_leaves(&map), 2); + let dict = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + assert_eq!(count_leaves(&dict), 1); + } + + /// `{a, b, c} CAST TO {b}` keeps only b's leaf. + #[test] + fn clip_struct_subset() { + let physical = struct_of(vec![utf8("a"), int64("b"), utf8("c")]); + let target = struct_of(vec![int64("b")]); + let kept = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![1]); + assert_eq!(emitted(&physical, &kept), struct_of(vec![int64("b")])); + } + + /// Target field order does not matter: emitted type is in physical order. + #[test] + fn clip_struct_reordered_target() { + let physical = struct_of(vec![utf8("a"), int64("b"), utf8("c")]); + let target = struct_of(vec![utf8("c"), utf8("a")]); + let kept = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0, 2]); + assert_eq!( + emitted(&physical, &kept), + struct_of(vec![utf8("a"), utf8("c")]) + ); + } + + /// Target fields missing from the physical type are ignored (the runtime + /// cast null-fills them). + #[test] + fn clip_struct_target_field_missing_from_physical() { + let physical = struct_of(vec![utf8("a"), int64("b")]); + let target = struct_of(vec![utf8("a"), int64("z")]); + let kept = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0]); + assert_eq!(emitted(&physical, &kept), struct_of(vec![utf8("a")])); + } + + /// Leaf-level type mismatch (promotion) still clips: the emitted type + /// keeps the physical leaf type; the cast performs the promotion. + #[test] + fn clip_keeps_physical_leaf_types() { + let physical = + struct_of(vec![Field::new("x", DataType::Int32, true), utf8("pad")]); + let target = struct_of(vec![int64("x")]); + let kept = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0]); + assert_eq!( + emitted(&physical, &kept), + struct_of(vec![Field::new("x", DataType::Int32, true)]) + ); + } + + /// Zero name overlap keeps the first physical child (>=1 leaf invariant). + #[test] + fn clip_struct_no_overlap_keeps_first_child() { + let physical = struct_of(vec![utf8("a"), int64("b")]); + let target = struct_of(vec![utf8("z")]); + let kept = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0]); + assert_eq!(emitted(&physical, &kept), struct_of(vec![utf8("a")])); + } + + /// Nested struct-in-struct clips at both levels. + #[test] + fn clip_nested_struct() { + let inner_physical = struct_of(vec![int64("x"), utf8("pad_inner")]); + let physical = struct_of(vec![ + inner_physical.clone().into_field("inner"), + utf8("pad_outer"), + ]); + let target = struct_of(vec![struct_of(vec![int64("x")]).into_field("inner")]); + let kept = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0]); + assert_eq!( + emitted(&physical, &kept), + struct_of(vec![struct_of(vec![int64("x")]).into_field("inner")]) + ); + } + + /// List — the headline case. + #[test] + fn clip_list_of_struct() { + let physical = list_of(struct_of(vec![int64("x"), utf8("y"), utf8("pad")])); + let target = list_of(struct_of(vec![int64("x"), utf8("y")])); + let kept = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0, 1]); + assert_eq!( + emitted(&physical, &kept), + list_of(struct_of(vec![int64("x"), utf8("y")])) + ); + } + + #[test] + fn clip_large_list_of_struct() { + let item = |fields| Arc::new(Field::new("item", struct_of(fields), true)); + let physical = DataType::LargeList(item(vec![int64("x"), utf8("pad")])); + let target = DataType::LargeList(item(vec![int64("x")])); + let kept = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0]); + assert_eq!( + emitted(&physical, &kept), + DataType::LargeList(item(vec![int64("x")])) + ); + } + + /// Wrapper-kind mismatch cannot be clipped. + #[test] + fn no_clip_on_wrapper_mismatch() { + let physical = list_of(struct_of(vec![int64("x"), utf8("pad")])); + let target = DataType::LargeList(Arc::new(Field::new( + "item", + struct_of(vec![int64("x")]), + true, + ))); + assert!(clip_for_cast(&physical, &target).is_none()); + } + + /// Maps are opaque: never clipped. + #[test] + fn no_clip_on_map() { + let entries = |fields| Arc::new(Field::new("entries", struct_of(fields), false)); + let physical = + DataType::Map(entries(vec![utf8("key"), int64("a"), int64("b")]), false); + let target = DataType::Map(entries(vec![utf8("key"), int64("a")]), false); + assert!(clip_for_cast(&physical, &target).is_none()); + } + + /// Identical types: nothing to prune. + #[test] + fn no_clip_when_identical() { + let t = struct_of(vec![utf8("a"), int64("b")]); + assert!(clip_for_cast(&t, &t).is_none()); + } + + /// Non-nested types: nothing to prune. + #[test] + fn no_clip_on_primitives() { + assert!(clip_for_cast(&DataType::Int32, &DataType::Int64).is_none()); + } + + #[test] + fn prune_type_by_offsets_union() { + // struct{a, inner{x, y}, c} with kept leaves {0 (a), 2 (inner.y)} + let physical = struct_of(vec![ + utf8("a"), + struct_of(vec![int64("x"), int64("y")]).into_field("inner"), + utf8("c"), + ]); + let kept: BTreeSet = [0, 2].into_iter().collect(); + let pruned = prune_type_by_kept_offsets(&physical, &kept).unwrap(); + assert_eq!( + pruned, + struct_of(vec![ + utf8("a"), + struct_of(vec![int64("y")]).into_field("inner"), + ]) + ); + } + + #[test] + fn prune_type_by_offsets_through_list() { + let physical = list_of(struct_of(vec![int64("x"), utf8("y"), utf8("pad")])); + let kept: BTreeSet = [1].into_iter().collect(); + let pruned = prune_type_by_kept_offsets(&physical, &kept).unwrap(); + assert_eq!(pruned, list_of(struct_of(vec![utf8("y")]))); + } + + /// Pins the arrow-rs behavior this module relies on: selecting a subset + /// of leaves under a `List` column with `ProjectionMask::leaves` + /// makes the reader emit exactly the type predicted by [`clip_for_cast`], + /// and null list rows / null struct elements survive (their validity is + /// reconstructed from the surviving leaves' definition levels). + #[test] + fn arrow_reader_emits_clipped_type_for_masked_list_struct() { + use arrow::array::{ + Array, ArrayRef, Int64Array, ListArray, StringArray, StructArray, + }; + use arrow::buffer::{NullBuffer, OffsetBuffer}; + use arrow::record_batch::RecordBatch; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use parquet::arrow::{ArrowWriter, ProjectionMask}; + + let item_fields = Fields::from(vec![int64("x"), utf8("y"), utf8("pad")]); + let item_field = Arc::new(Field::new( + "item", + DataType::Struct(item_fields.clone()), + true, + )); + let schema = Arc::new(arrow::datatypes::Schema::new(vec![Field::new( + "events", + DataType::List(Arc::clone(&item_field)), + true, + )])); + + // 3 elements; element 1 is a NULL struct. Rows: [e0, e1], NULL, [e2]. + let columns: Vec = vec![ + Arc::new(Int64Array::from(vec![Some(1), None, Some(3)])), + Arc::new(StringArray::from(vec![Some("a"), None, Some("c")])), + Arc::new(StringArray::from(vec![Some("p0"), None, Some("p2")])), + ]; + let struct_validity = NullBuffer::from(vec![true, false, true]); + let values = StructArray::new(item_fields, columns, Some(struct_validity)); + let list_validity = NullBuffer::from(vec![true, false, true]); + let events = ListArray::new( + item_field, + OffsetBuffer::from_lengths([2, 0, 1]), + Arc::new(values), + Some(list_validity), + ); + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(events)]).unwrap(); + + let file = tempfile::NamedTempFile::new().unwrap(); + let mut writer = + ArrowWriter::try_new(file.reopen().unwrap(), schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + // Clip to the narrow target {x, y}. + let physical = batch.schema().field(0).data_type().clone(); + let target = list_of(struct_of(vec![int64("x"), utf8("y")])); + let kept = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0, 1]); + + let builder = + ParquetRecordBatchReaderBuilder::try_new(file.reopen().unwrap()).unwrap(); + let mask = ProjectionMask::leaves(builder.parquet_schema(), kept.iter().copied()); + let reader = builder.with_projection(mask).build().unwrap(); + let out: Vec = reader.map(|b| b.unwrap()).collect(); + assert_eq!(out.len(), 1); + let out = &out[0]; + + // Emitted type matches the prediction. + assert_eq!( + out.schema().field(0).data_type(), + &emitted(&physical, &kept) + ); + + // Null semantics survive the clip. + let events = out.column(0).as_any().downcast_ref::().unwrap(); + assert!(events.is_valid(0)); + assert!(events.is_null(1)); + assert!(events.is_valid(2)); + let structs = events + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(structs.len(), 3); + assert!(structs.is_valid(0)); + assert!(structs.is_null(1)); + assert!(structs.is_valid(2)); + let x = structs + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(x.value(0), 1); + assert_eq!(x.value(2), 3); + } + + trait IntoField { + fn into_field(self, name: &str) -> Field; + } + + impl IntoField for DataType { + fn into_field(self, name: &str) -> Field { + Field::new(name, self, true) + } + } +} diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index 87ec341f590da..fd35ae3294abc 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -271,6 +271,9 @@ pub(super) struct ParquetMorselizer { pub enable_bloom_filter: bool, /// Should row group pruning be applied pub enable_row_group_stats_pruning: bool, + /// Should projected nested columns read through a narrowing cast be + /// clipped to the leaves the cast retains + pub nested_projection_pruning: bool, /// Coerce INT96 timestamps to specific TimeUnit pub coerce_int96: Option, /// Optional timezone applied to INT96-coerced timestamps. When `Some`, the @@ -445,6 +448,7 @@ struct PreparedParquetOpen { enable_page_index: bool, enable_bloom_filter: bool, enable_row_group_stats_pruning: bool, + nested_projection_pruning: bool, limit: Option, coerce_int96: Option, coerce_int96_tz: Option>, @@ -844,6 +848,7 @@ impl ParquetMorselizer { enable_page_index: self.enable_page_index, enable_bloom_filter: self.enable_bloom_filter, enable_row_group_stats_pruning: self.enable_row_group_stats_pruning, + nested_projection_pruning: self.nested_projection_pruning, limit: self.limit, coerce_int96: self.coerce_int96, coerce_int96_tz: self.coerce_int96_tz.clone(), @@ -1384,6 +1389,7 @@ impl RowGroupsPrunedParquetOpen { reader_metadata.parquet_schema(), &prepared.output_schema, prepared.virtual_state.as_deref(), + prepared.nested_projection_pruning, )?; let (decoder, rg_plan) = { @@ -1750,6 +1756,7 @@ mod test { enable_page_index: bool, enable_bloom_filter: bool, enable_row_group_stats_pruning: bool, + nested_projection_pruning: bool, coerce_int96: Option, max_predicate_cache_size: Option, reverse_row_groups: bool, @@ -1858,6 +1865,7 @@ mod test { enable_page_index: false, enable_bloom_filter: false, enable_row_group_stats_pruning: false, + nested_projection_pruning: true, coerce_int96: None, max_predicate_cache_size: None, reverse_row_groups: false, @@ -2026,6 +2034,7 @@ mod test { enable_page_index: self.enable_page_index, enable_bloom_filter: self.enable_bloom_filter, enable_row_group_stats_pruning: self.enable_row_group_stats_pruning, + nested_projection_pruning: self.nested_projection_pruning, coerce_int96: self.coerce_int96, // End-to-end coercion behavior (including timezone) is // covered by parquet.slt. No opener-level test currently diff --git a/datafusion/datasource-parquet/src/projection_read_plan.rs b/datafusion/datasource-parquet/src/projection_read_plan.rs new file mode 100644 index 0000000000000..9d445a285491a --- /dev/null +++ b/datafusion/datasource-parquet/src/projection_read_plan.rs @@ -0,0 +1,657 @@ +// 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. + +//! Resolution of expressions against a Parquet file's schema into a +//! [`ParquetReadPlan`]: the leaf-level [`ProjectionMask`] to install on the +//! decoder plus the Arrow schema the decoder will emit under that mask. +//! +//! This is shared by the opener's projection handling (via +//! [`build_projection_read_plan`]) and row-filter construction (via +//! [`crate::row_filter`]), which both need to translate column and struct +//! field references into Parquet leaf indices. + +use std::collections::{BTreeMap, BTreeSet}; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use parquet::arrow::ProjectionMask; +use parquet::schema::types::SchemaDescriptor; + +use datafusion_common::tree_node::TreeNode; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr::utils::collect_columns; + +use crate::nested_schema_pruning::{ + CastColumnAccess, clip_for_cast, count_leaves, field_with_type, + prune_type_by_kept_offsets, +}; +use crate::row_filter::PushdownChecker; + +/// The result of resolving which Parquet leaf columns and Arrow schema fields +/// are needed to evaluate an expression against a Parquet file +/// +/// This is the shared output of the column resolution pipeline used by both +/// the row filter to build `ArrowPredicate`s and the opener to build `ProjectionMask`s +#[derive(Debug, Clone)] +pub(crate) struct ParquetReadPlan { + /// Projection mask built from leaf column indices in the Parquet schema. + /// Using a `ProjectionMask` directly (rather than raw indices) prevents + /// bugs from accidentally mixing up root vs leaf indices. + pub projection_mask: ProjectionMask, + /// The projected Arrow schema containing only the columns/fields required + /// Struct types are pruned to include only the accessed sub-fields + pub projected_schema: SchemaRef, +} + +/// Records a struct field access via `get_field(struct_col, 'field1', 'field2', ...)`. +/// +/// This allows the row filter to project only the specific Parquet leaf columns +/// needed by the filter, rather than all leaves of the struct. +#[derive(Debug, Clone)] +pub(crate) struct StructFieldAccess { + /// Arrow root column index of the struct in the file schema. + pub(crate) root_index: usize, + /// Field names forming the path into the struct. + /// e.g., `["value"]` for `s['value']`, `["outer", "inner"]` for `s['outer']['inner']`. + pub(crate) field_path: Vec, +} + +/// Builds a unified [`ParquetReadPlan`] for a set of projection expressions +/// +/// Unlike [`crate::row_filter::build_parquet_read_plan`] (which is used for +/// filter pushdown and returns `None` when an expression references +/// unsupported nested types or missing columns), this function always +/// succeeds. It collects every column that *can* be resolved in the file and +/// produces a leaf-level projection mask. Columns missing from the file are +/// silently skipped since the projection layer handles those by inserting +/// nulls. +pub(crate) fn build_projection_read_plan( + exprs: impl IntoIterator>, + file_schema: &Schema, + schema_descr: &SchemaDescriptor, + prune_nested_casts: bool, +) -> ParquetReadPlan { + // fast path: if every expression is a plain Column reference, skip all + // struct analysis and use root-level projection directly + let exprs = exprs.into_iter().collect::>(); + let all_plain_columns = exprs.iter().all(|e| e.downcast_ref::().is_some()); + + if all_plain_columns { + let mut root_indices: Vec = exprs + .iter() + .map(|e| e.downcast_ref::().unwrap().index()) + .collect(); + root_indices.sort_unstable(); + root_indices.dedup(); + + let projection_mask = + ProjectionMask::roots(schema_descr, root_indices.iter().copied()); + let projected_schema = Arc::new( + file_schema + .project(&root_indices) + .expect("valid column indices"), + ); + + return ParquetReadPlan { + projection_mask, + projected_schema, + }; + } + + // secondary fast path: if no column contains a struct at any nesting + // level, there are no leaves to prune and we can skip PushdownChecker + // traversal and use root-level projection + let has_struct_columns = file_schema + .fields() + .iter() + .any(|f| contains_struct(f.data_type())); + + if !has_struct_columns { + let mut root_indices = exprs + .into_iter() + .flat_map(|e| collect_columns(&e).into_iter().map(|col| col.index())) + .collect::>(); + + root_indices.sort_unstable(); + root_indices.dedup(); + + let projection_mask = + ProjectionMask::roots(schema_descr, root_indices.iter().copied()); + + let projected_schema = Arc::new( + file_schema + .project(&root_indices) + .expect("valid column indices"), + ); + + return ParquetReadPlan { + projection_mask, + projected_schema, + }; + } + + let mut all_root_indices = Vec::new(); + let mut all_struct_accesses = Vec::new(); + let mut all_cast_accesses = Vec::new(); + + for expr in exprs { + let mut checker = PushdownChecker::new(file_schema, true) + .with_cast_collection(prune_nested_casts); + let _ = expr.visit(&mut checker); + let columns = checker.into_sorted_columns(); + + all_root_indices.extend_from_slice(&columns.required_columns); + all_struct_accesses.extend(columns.struct_field_accesses); + all_cast_accesses.extend(columns.cast_accesses); + } + + all_root_indices.sort_unstable(); + all_root_indices.dedup(); + + // A whole-column reference reads every leaf of the root, so clipping the + // same root for a cast (or a struct field access) would be overridden + // anyway: drop those accesses up front. + let whole_roots: BTreeSet = all_root_indices.iter().copied().collect(); + all_cast_accesses.retain(|c| !whole_roots.contains(&c.root_index)); + + if !all_cast_accesses.is_empty() { + return build_read_plan_with_cast_clipping( + file_schema, + schema_descr, + &all_root_indices, + &all_struct_accesses, + &all_cast_accesses, + ); + } + + // when no struct field accesses were found, fall back to root-level projection + // to match the performance of the simple path + if all_struct_accesses.is_empty() { + let projection_mask = + ProjectionMask::roots(schema_descr, all_root_indices.iter().copied()); + let projected_schema = Arc::new( + file_schema + .project(&all_root_indices) + .expect("valid column indices"), + ); + + return ParquetReadPlan { + projection_mask, + projected_schema, + }; + } + + let leaf_indices = { + let mut out = + leaf_indices_for_roots(all_root_indices.iter().copied(), schema_descr); + let struct_leaf_indices = + resolve_struct_field_leaves(&all_struct_accesses, file_schema, schema_descr); + + out.extend_from_slice(&struct_leaf_indices); + out.sort_unstable(); + out.dedup(); + + out + }; + + let projection_mask = + ProjectionMask::leaves(schema_descr, leaf_indices.iter().copied()); + + let projected_schema = + build_filter_schema(file_schema, &all_root_indices, &all_struct_accesses); + + ParquetReadPlan { + projection_mask, + projected_schema, + } +} + +/// Does this type contain a struct at any nesting depth? +fn contains_struct(dt: &DataType) -> bool { + match dt { + DataType::Struct(_) => true, + DataType::List(f) + | DataType::LargeList(f) + | DataType::ListView(f) + | DataType::LargeListView(f) + | DataType::FixedSizeList(f, _) + | DataType::Map(f, _) => contains_struct(f.data_type()), + DataType::Dictionary(_, value) => contains_struct(value), + DataType::RunEndEncoded(_, value) => contains_struct(value.data_type()), + _ => false, + } +} + +/// Builds a [`ParquetReadPlan`] when at least one projected root column is +/// consumed through a cast to a narrower nested type. +/// +/// Per root, in ascending root-index order: +/// - roots referenced as whole columns keep every leaf and their full +/// physical field (whole-column reads take precedence; cast accesses on +/// such roots were already dropped by the caller); +/// - roots consumed only through casts and/or `get_field` accesses keep the +/// union of the leaves those accesses consume, and their projected field +/// type is derived from that union. +/// +/// Falls back to keeping all leaves of a root whenever its clipping is not +/// provably safe (see [`crate::nested_schema_pruning`]). +fn build_read_plan_with_cast_clipping( + file_schema: &Schema, + schema_descr: &SchemaDescriptor, + whole_root_indices: &[usize], + struct_accesses: &[StructFieldAccess], + cast_accesses: &[CastColumnAccess], +) -> ParquetReadPlan { + let whole_roots: BTreeSet = whole_root_indices.iter().copied().collect(); + + // Union of kept *absolute* leaf indices per clipped root. + let mut kept_by_root: BTreeMap> = BTreeMap::new(); + // Roots where clipping had to be abandoned (treated as whole reads). + let mut fallback_roots: BTreeSet = BTreeSet::new(); + + for access in cast_accesses { + let root = access.root_index; + if fallback_roots.contains(&root) { + continue; + } + let physical_type = file_schema.field(root).data_type(); + let root_leaves = leaf_indices_for_roots([root], schema_descr); + + // Defensive: the arrow type's leaf count must agree with the parquet + // schema (it can diverge if the file embeds a different arrow + // schema). If not, never risk a wrong mask — read the whole root. + if root_leaves.len() != count_leaves(physical_type) { + fallback_roots.insert(root); + kept_by_root.remove(&root); + continue; + } + + match clip_for_cast(physical_type, &access.target_type) { + Some(kept_offsets) => { + let start = root_leaves[0]; + kept_by_root + .entry(root) + .or_default() + .extend(kept_offsets.iter().map(|o| start + o)); + } + // Nothing prunable for this cast: every leaf is consumed. + None => { + fallback_roots.insert(root); + kept_by_root.remove(&root); + } + } + } + + // `get_field` accesses union into the same per-root sets. Accesses on + // whole-read or fallback roots are covered by the full read. + for access in struct_accesses { + let root = access.root_index; + if whole_roots.contains(&root) || fallback_roots.contains(&root) { + continue; + } + let leaves = resolve_struct_field_leaves( + std::slice::from_ref(access), + file_schema, + schema_descr, + ); + match kept_by_root.get_mut(&root) { + Some(kept) => kept.extend(leaves), + None => { + kept_by_root.entry(root).or_default().extend(leaves); + } + } + } + + // Assemble mask + schema over every referenced root in ascending order. + let all_roots: BTreeSet = whole_roots + .iter() + .copied() + .chain(fallback_roots.iter().copied()) + .chain(kept_by_root.keys().copied()) + .collect(); + + let mut leaf_indices: Vec = Vec::new(); + let mut fields: Vec> = Vec::with_capacity(all_roots.len()); + + for &root in &all_roots { + let field = file_schema.field(root); + match kept_by_root.get(&root) { + Some(kept) if !kept.is_empty() => { + let root_leaves = leaf_indices_for_roots([root], schema_descr); + let start = root_leaves[0]; + let relative: BTreeSet = + kept.iter().map(|abs| abs - start).collect(); + match prune_type_by_kept_offsets(field.data_type(), &relative) { + Some(pruned_type) => { + leaf_indices.extend(kept.iter().copied()); + fields.push(field_with_type(field, pruned_type)); + } + // Cannot express the kept set as a type: full read. + None => { + leaf_indices.extend(root_leaves); + fields.push(Arc::new(field.clone())); + } + } + } + // Whole-column read (whole reference, fallback, or empty set). + _ => { + leaf_indices.extend(leaf_indices_for_roots([root], schema_descr)); + fields.push(Arc::new(field.clone())); + } + } + } + + leaf_indices.sort_unstable(); + leaf_indices.dedup(); + + ParquetReadPlan { + projection_mask: ProjectionMask::leaves( + schema_descr, + leaf_indices.iter().copied(), + ), + projected_schema: Arc::new(Schema::new_with_metadata( + fields, + file_schema.metadata().clone(), + )), + } +} + +pub(crate) fn leaf_indices_for_roots( + root_indices: I, + schema_descr: &SchemaDescriptor, +) -> Vec +where + I: IntoIterator, +{ + // Always map root (Arrow) indices to Parquet leaf indices via the schema + // descriptor. Arrow root indices only equal Parquet leaf indices when the + // schema has no group columns (Struct, Map, etc.); when group columns + // exist, their children become separate leaves and shift all subsequent + // leaf indices. + // Struct columns are unsupported. + let root_set: BTreeSet<_> = root_indices.into_iter().collect(); + + (0..schema_descr.num_columns()) + .filter(|leaf_idx| { + root_set.contains(&schema_descr.get_column_root_idx(*leaf_idx)) + }) + .collect() +} + +/// Resolves struct field access to specific Parquet leaf column indices +/// +/// For every `StructFieldAccess`, finds the leaf columns in the Parquet schema +/// whose path matches the struct root name + field path. This avoids reading all +/// leaves of a struct when only specific fields are needed +pub(crate) fn resolve_struct_field_leaves( + accesses: &[StructFieldAccess], + file_schema: &Schema, + schema_descr: &SchemaDescriptor, +) -> Vec { + let mut leaf_indices = Vec::new(); + + for access in accesses { + let root_name = file_schema.field(access.root_index).name(); + let prefix = std::iter::once(root_name.as_str()) + .chain(access.field_path.iter().map(|p| p.as_str())) + .collect::>(); + + for leaf_idx in 0..schema_descr.num_columns() { + let col = schema_descr.column(leaf_idx); + let col_path = col.path().parts(); + + // A leaf matches if its path starts with our prefix. + // e.g., prefix=["s", "value"] matches leaf path ["s", "value"] + // prefix=["s", "outer"] matches ["s", "outer", "inner"] + + // a leaf matches if its path starts with our prefix + // for example: prefix=["s", "value"] matches leaf path ["s", "value"] + // prefix=["s", "outer"] matches ["s", "outer", "inner"] + let leaf_matches_path = col_path.len() >= prefix.len() + && col_path.iter().zip(prefix.iter()).all(|(a, b)| a == b); + + if leaf_matches_path { + leaf_indices.push(leaf_idx); + } + } + } + + leaf_indices +} + +/// Builds a filter schema that includes only the fields actually accessed by the +/// filter expression. +/// +/// For regular (non-struct) columns, the full field type is used. +/// For struct columns accessed via `get_field`, a pruned struct type is created +/// containing only the fields along the access path. Note: it must match the schema +/// that the Parquet reader produces when projecting specific struct leaves +pub(crate) fn build_filter_schema( + file_schema: &Schema, + regular_indices: &[usize], + struct_field_accesses: &[StructFieldAccess], +) -> SchemaRef { + let regular_set: BTreeSet = regular_indices.iter().copied().collect(); + let paths_by_root = group_access_paths_by_root(struct_field_accesses); + + let all_indices = regular_indices + .iter() + .copied() + .chain(paths_by_root.keys().copied()) + .collect::>(); + + let fields = all_indices + .iter() + .map(|&idx| { + let field = file_schema.field(idx); + + // if this column appears as a regular (whole-column) reference, + // keep the full type + // + // Pruning is only valid when the column is accessed exclusively + // through struct field accesses + if regular_set.contains(&idx) { + return Arc::new(field.clone()); + } + + let Some(field_paths) = paths_by_root.get(&idx) else { + return Arc::new(field.clone()); + }; + + let pruned_data_type = prune_struct_type(field.data_type(), field_paths); + Arc::new(Field::new( + field.name(), + pruned_data_type, + field.is_nullable(), + )) + }) + .collect::>(); + + Arc::new(Schema::new_with_metadata( + fields, + file_schema.metadata().clone(), + )) +} + +/// Groups struct field access paths once for the root schema level. +/// +/// Each map entry contains the complete field paths accessed below a root +/// column. Recursive pruning groups these paths by their next component at each +/// nested struct level. +fn group_access_paths_by_root( + struct_field_accesses: &[StructFieldAccess], +) -> BTreeMap> { + let mut paths_by_root: BTreeMap> = BTreeMap::new(); + for StructFieldAccess { + root_index, + field_path, + } in struct_field_accesses + { + paths_by_root + .entry(*root_index) + .or_default() + .push(field_path.as_slice()); + } + + paths_by_root +} + +/// Groups access paths once for the current struct level. +/// +/// The map key is the field name at this level. The map value is the list of +/// remaining path suffixes below that field. An empty suffix means the access +/// path terminates at that field, so the full field must be preserved. +fn group_paths_by_next_field<'a>( + paths: &'a [&'a [String]], +) -> BTreeMap<&'a str, Vec<&'a [String]>> { + let mut paths_by_field: BTreeMap<&str, Vec<&[String]>> = BTreeMap::new(); + for path in paths { + if let Some((field, sub_path)) = path.split_first() { + paths_by_field + .entry(field.as_str()) + .or_default() + .push(sub_path); + } + } + + paths_by_field +} + +fn prune_struct_type(dt: &DataType, paths: &[&[String]]) -> DataType { + let DataType::Struct(fields) = dt else { + return dt.clone(); + }; + + let paths_by_field = group_paths_by_next_field(paths); + + let pruned_fields = fields + .iter() + .filter_map(|f| { + let sub_paths = paths_by_field.get(f.name().as_str())?; + + let out = if sub_paths.iter().any(|sub| sub.is_empty()) { + // Leaf of access path — keep the field as-is. + Arc::clone(f) + } else { + // Recurse into nested struct. + let pruned = prune_struct_type(f.data_type(), sub_paths); + Arc::new(Field::new(f.name(), pruned, f.is_nullable())) + }; + + Some(out) + }) + .collect::>(); + + DataType::Struct(pruned_fields.into()) +} + +#[cfg(test)] +mod test { + use super::*; + use Column as PhysicalColumn; + use arrow::array::{Int32Array, RecordBatch, StringArray, StructArray}; + use arrow::datatypes::Fields; + use datafusion_common::ScalarValue; + use datafusion_expr::{Expr, col}; + use datafusion_functions::core::get_field; + use datafusion_physical_expr::planner::logical2physical; + use parquet::arrow::ArrowWriter; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use tempfile::NamedTempFile; + + #[test] + fn projection_read_plan_preserves_full_struct() { + // Schema: id (Int32), s (Struct{value: Int32, label: Utf8}) + // Parquet leaves: id=0, s.value=1, s.label=2 + let struct_fields: Fields = vec![ + Arc::new(Field::new("value", DataType::Int32, false)), + Arc::new(Field::new("label", DataType::Utf8, false)), + ] + .into(); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("s", DataType::Struct(struct_fields.clone()), false), + ])); + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StructArray::new( + struct_fields, + vec![ + Arc::new(Int32Array::from(vec![10, 20, 30])) as _, + Arc::new(StringArray::from(vec!["a", "b", "c"])) as _, + ], + None, + )), + ], + ) + .unwrap(); + + let file = NamedTempFile::new().expect("temp file"); + let mut writer = + ArrowWriter::try_new(file.reopen().unwrap(), Arc::clone(&schema), None) + .expect("writer"); + writer.write(&batch).expect("write batch"); + writer.close().expect("close writer"); + + let reader_file = file.reopen().expect("reopen file"); + let builder = ParquetRecordBatchReaderBuilder::try_new(reader_file) + .expect("reader builder"); + let metadata = builder.metadata().clone(); + let file_schema = builder.schema().clone(); + let schema_descr = metadata.file_metadata().schema_descr(); + + // Simulate SELECT * output projection: Column("id") and Column("s") + // Plus a get_field(s, 'value') expression from the pushed-down filter + let exprs: Vec> = vec![ + Arc::new(PhysicalColumn::new("id", 0)), + Arc::new(PhysicalColumn::new("s", 1)), + logical2physical( + &get_field().call(vec![ + col("s"), + Expr::Literal(ScalarValue::Utf8(Some("value".to_string())), None), + ]), + &file_schema, + ), + ]; + + let read_plan = + build_projection_read_plan(exprs, &file_schema, schema_descr, true); + + // The projected schema must have the FULL struct type because Column("s") + // is in the projection. It should NOT be narrowed to Struct{value: Int32}. + let s_field = read_plan.projected_schema.field_with_name("s").unwrap(); + assert_eq!( + s_field.data_type(), + &DataType::Struct( + vec![ + Arc::new(Field::new("value", DataType::Int32, false)), + Arc::new(Field::new("label", DataType::Utf8, false)), + ] + .into() + ), + ); + + // all3 Parquet leaves should be in the projection mask + let expected_mask = ProjectionMask::leaves(schema_descr, [0, 1, 2]); + assert_eq!(read_plan.projection_mask, expected_mask,); + } +} diff --git a/datafusion/datasource-parquet/src/row_filter.rs b/datafusion/datasource-parquet/src/row_filter.rs index ef0478f3159bc..cc60a18321e11 100644 --- a/datafusion/datasource-parquet/src/row_filter.rs +++ b/datafusion/datasource-parquet/src/row_filter.rs @@ -65,11 +65,10 @@ //! - `WHERE s['value'] > 5` — pushed down (accesses a primitive leaf) //! - `WHERE s IS NOT NULL` — not pushed down (references the whole struct) -use std::collections::{BTreeMap, BTreeSet}; use std::sync::Arc; use arrow::array::BooleanArray; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::datatypes::{DataType, Schema, SchemaRef}; use arrow::error::{ArrowError, Result as ArrowResult}; use arrow::record_batch::RecordBatch; use datafusion_functions::core::file_row_index::FileRowIndexFunc; @@ -77,20 +76,25 @@ use datafusion_functions::core::getfield::GetFieldFunc; use parquet::arrow::ProjectionMask; use parquet::arrow::arrow_reader::{ArrowPredicate, RowFilter}; use parquet::file::metadata::ParquetMetaData; -use parquet::schema::types::SchemaDescriptor; use datafusion_common::Result; use datafusion_common::cast::as_boolean_array; +use datafusion_common::nested_struct::requires_nested_struct_cast; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor}; use datafusion_physical_expr::ScalarFunctionExpr; -use datafusion_physical_expr::expressions::{Column, Literal}; -use datafusion_physical_expr::utils::{collect_columns, reassign_expr_columns}; +use datafusion_physical_expr::expressions::{CastExpr, Column, Literal}; +use datafusion_physical_expr::utils::reassign_expr_columns; use datafusion_physical_expr::{PhysicalExpr, split_conjunction}; use datafusion_physical_plan::metrics; use super::ParquetFileMetrics; use super::supported_predicates::supports_list_predicates; +use crate::nested_schema_pruning::CastColumnAccess; +use crate::projection_read_plan::{ + ParquetReadPlan, StructFieldAccess, build_filter_schema, leaf_indices_for_roots, + resolve_struct_field_leaves, +}; /// A "compiled" predicate passed to `ParquetRecordBatchStream` to perform /// row-level filtering during parquet decoding. @@ -189,22 +193,6 @@ pub(crate) struct FilterCandidate { read_plan: ParquetReadPlan, } -/// The result of resolving which Parquet leaf columns and Arrow schema fields -/// are needed to evaluate an expression against a Parquet file -/// -/// This is the shared output of the column resolution pipeline used by both -/// the row filter to build `ArrowPredicate`s and the opener to build `ProjectionMask`s -#[derive(Debug, Clone)] -pub(crate) struct ParquetReadPlan { - /// Projection mask built from leaf column indices in the Parquet schema. - /// Using a `ProjectionMask` directly (rather than raw indices) prevents - /// bugs from accidentally mixing up root vs leaf indices. - pub projection_mask: ProjectionMask, - /// The projected Arrow schema containing only the columns/fields required - /// Struct types are pruned to include only the accessed sub-fields - pub projected_schema: SchemaRef, -} - /// Helper to build a `FilterCandidate`. /// /// This will do several things: @@ -256,7 +244,7 @@ impl FilterCandidateBuilder { /// /// Struct field access via `get_field` is supported when the resolved leaf type /// is primitive (e.g. `get_field(struct_col, 'field') > 5`). -struct PushdownChecker<'schema> { +pub(crate) struct PushdownChecker<'schema> { /// Does the expression require any non-primitive columns (like structs)? non_primitive_columns: bool, /// Does the expression reference any columns not present in the file schema? @@ -269,6 +257,13 @@ struct PushdownChecker<'schema> { required_columns: Vec, /// Struct field accesses via `get_field`. struct_field_accesses: Vec, + /// Whole-column casts to a narrower nested type + /// (`CAST(col AS narrower_struct)`). Only collected when + /// [`Self::with_cast_collection`] enables it (projection analysis); + /// filter pushdown leaves this off. + cast_accesses: Vec, + /// Whether to collect [`Self::cast_accesses`]. + collect_cast_accesses: bool, /// Whether nested list columns are supported by the predicate semantics. allow_list_columns: bool, /// The Arrow schema of the parquet file. @@ -276,18 +271,26 @@ struct PushdownChecker<'schema> { } impl<'schema> PushdownChecker<'schema> { - fn new(file_schema: &'schema Schema, allow_list_columns: bool) -> Self { + pub(crate) fn new(file_schema: &'schema Schema, allow_list_columns: bool) -> Self { Self { non_primitive_columns: false, projected_columns: false, has_unpushable_udfs: false, required_columns: Vec::new(), struct_field_accesses: Vec::new(), + cast_accesses: Vec::new(), + collect_cast_accesses: false, allow_list_columns, file_schema, } } + /// Enable collection of whole-column casts to narrower nested types. + pub(crate) fn with_cast_collection(mut self, collect: bool) -> Self { + self.collect_cast_accesses = collect; + self + } + /// Checks whether a struct's root column exists in the file schema and, if so, /// records its index so the entire struct is decoded for filter evaluation. /// @@ -386,12 +389,13 @@ impl<'schema> PushdownChecker<'schema> { /// This method sorts the column indices and removes duplicates. The sort /// is required because downstream code relies on column indices being in /// ascending order for correct schema projection. - fn into_sorted_columns(mut self) -> PushdownColumns { + pub(crate) fn into_sorted_columns(mut self) -> PushdownColumns { self.required_columns.sort_unstable(); self.required_columns.dedup(); PushdownColumns { required_columns: self.required_columns, struct_field_accesses: self.struct_field_accesses, + cast_accesses: self.cast_accesses, } } } @@ -483,6 +487,28 @@ impl TreeNodeVisitor<'_> for PushdownChecker<'_> { } } + // Handle whole-column casts to a narrower nested type, e.g. + // `CAST(events AS List>)` as inserted by the + // physical expression adapter when the logical file schema declares a + // nested column narrower than the physical file. Recording the cast + // target lets the projection read only the leaves the cast consumes + // (see `crate::nested_schema_pruning`). + if self.collect_cast_accesses + && let Some(cast) = node.downcast_ref::() + && let Some(column) = cast.expr().downcast_ref::() + && let Ok(idx) = self.file_schema.index_of(column.name()) + && requires_nested_struct_cast( + self.file_schema.field(idx).data_type(), + cast.cast_type(), + ) + { + self.cast_accesses.push(CastColumnAccess { + root_index: idx, + target_type: cast.cast_type().clone(), + }); + return Ok(TreeNodeRecursion::Jump); + } + if let Some(column) = node.downcast_ref::() && let Some(recursion) = self.check_single_column(column.name()) { @@ -506,27 +532,17 @@ impl TreeNodeVisitor<'_> for PushdownChecker<'_> { /// with respect to nested column handling during Parquet decoding. /// Result of checking which columns are required for filter pushdown. #[derive(Debug)] -struct PushdownColumns { +pub(crate) struct PushdownColumns { /// Sorted, unique column indices into the file schema required to evaluate /// the filter expression. Must be in ascending order for correct schema /// projection matching. Does not include struct columns accessed via `get_field`. - required_columns: Vec, + pub(crate) required_columns: Vec, /// Struct field accesses via `get_field`. Each entry records the root struct /// column index and the field path being accessed. - struct_field_accesses: Vec, -} - -/// Records a struct field access via `get_field(struct_col, 'field1', 'field2', ...)`. -/// -/// This allows the row filter to project only the specific Parquet leaf columns -/// needed by the filter, rather than all leaves of the struct. -#[derive(Debug, Clone)] -struct StructFieldAccess { - /// Arrow root column index of the struct in the file schema. - root_index: usize, - /// Field names forming the path into the struct. - /// e.g., `["value"]` for `s['value']`, `["outer", "inner"]` for `s['outer']['inner']`. - field_path: Vec, + pub(crate) struct_field_accesses: Vec, + /// Whole-column casts to a narrower nested type. Empty unless cast + /// collection was enabled on the checker. + pub(crate) cast_accesses: Vec, } /// Checks if a given expression can be pushed down to the parquet decoder. @@ -604,323 +620,6 @@ pub(crate) fn build_parquet_read_plan( ))) } -/// Builds a unified [`ParquetReadPlan`] for a set of projection expressions -/// -/// Unlike [`build_parquet_read_plan`] (which is used for filter pushdown and -/// returns `None` when an expression references unsupported nested types or -/// missing columns), this function always succeeds. It collects every column -/// that *can* be resolved in the file and produces a leaf-level projection -/// mask. Columns missing from the file are silently skipped since the projection -/// layer handles those by inserting nulls. -pub(crate) fn build_projection_read_plan( - exprs: impl IntoIterator>, - file_schema: &Schema, - schema_descr: &SchemaDescriptor, -) -> ParquetReadPlan { - // fast path: if every expression is a plain Column reference, skip all - // struct analysis and use root-level projection directly - let exprs = exprs.into_iter().collect::>(); - let all_plain_columns = exprs.iter().all(|e| e.downcast_ref::().is_some()); - - if all_plain_columns { - let mut root_indices: Vec = exprs - .iter() - .map(|e| e.downcast_ref::().unwrap().index()) - .collect(); - root_indices.sort_unstable(); - root_indices.dedup(); - - let projection_mask = - ProjectionMask::roots(schema_descr, root_indices.iter().copied()); - let projected_schema = Arc::new( - file_schema - .project(&root_indices) - .expect("valid column indices"), - ); - - return ParquetReadPlan { - projection_mask, - projected_schema, - }; - } - - // secondary fast path: if the schema has no struct columns, we can skip - // PushdownChecker traversal and use root-level projection - let has_struct_columns = file_schema - .fields() - .iter() - .any(|f| matches!(f.data_type(), DataType::Struct(_))); - - if !has_struct_columns { - let mut root_indices = exprs - .into_iter() - .flat_map(|e| collect_columns(&e).into_iter().map(|col| col.index())) - .collect::>(); - - root_indices.sort_unstable(); - root_indices.dedup(); - - let projection_mask = - ProjectionMask::roots(schema_descr, root_indices.iter().copied()); - - let projected_schema = Arc::new( - file_schema - .project(&root_indices) - .expect("valid column indices"), - ); - - return ParquetReadPlan { - projection_mask, - projected_schema, - }; - } - - let mut all_root_indices = Vec::new(); - let mut all_struct_accesses = Vec::new(); - - for expr in exprs { - let mut checker = PushdownChecker::new(file_schema, true); - let _ = expr.visit(&mut checker); - let columns = checker.into_sorted_columns(); - - all_root_indices.extend_from_slice(&columns.required_columns); - all_struct_accesses.extend(columns.struct_field_accesses); - } - - all_root_indices.sort_unstable(); - all_root_indices.dedup(); - - // when no struct field accesses were found, fall back to root-level projection - // to match the performance of the simple path - if all_struct_accesses.is_empty() { - let projection_mask = - ProjectionMask::roots(schema_descr, all_root_indices.iter().copied()); - let projected_schema = Arc::new( - file_schema - .project(&all_root_indices) - .expect("valid column indices"), - ); - - return ParquetReadPlan { - projection_mask, - projected_schema, - }; - } - - let leaf_indices = { - let mut out = - leaf_indices_for_roots(all_root_indices.iter().copied(), schema_descr); - let struct_leaf_indices = - resolve_struct_field_leaves(&all_struct_accesses, file_schema, schema_descr); - - out.extend_from_slice(&struct_leaf_indices); - out.sort_unstable(); - out.dedup(); - - out - }; - - let projection_mask = - ProjectionMask::leaves(schema_descr, leaf_indices.iter().copied()); - - let projected_schema = - build_filter_schema(file_schema, &all_root_indices, &all_struct_accesses); - - ParquetReadPlan { - projection_mask, - projected_schema, - } -} - -fn leaf_indices_for_roots( - root_indices: I, - schema_descr: &SchemaDescriptor, -) -> Vec -where - I: IntoIterator, -{ - // Always map root (Arrow) indices to Parquet leaf indices via the schema - // descriptor. Arrow root indices only equal Parquet leaf indices when the - // schema has no group columns (Struct, Map, etc.); when group columns - // exist, their children become separate leaves and shift all subsequent - // leaf indices. - // Struct columns are unsupported. - let root_set: BTreeSet<_> = root_indices.into_iter().collect(); - - (0..schema_descr.num_columns()) - .filter(|leaf_idx| { - root_set.contains(&schema_descr.get_column_root_idx(*leaf_idx)) - }) - .collect() -} - -/// Resolves struct field access to specific Parquet leaf column indices -/// -/// For every `StructFieldAccess`, finds the leaf columns in the Parquet schema -/// whose path matches the struct root name + field path. This avoids reading all -/// leaves of a struct when only specific fields are needed -fn resolve_struct_field_leaves( - accesses: &[StructFieldAccess], - file_schema: &Schema, - schema_descr: &SchemaDescriptor, -) -> Vec { - let mut leaf_indices = Vec::new(); - - for access in accesses { - let root_name = file_schema.field(access.root_index).name(); - let prefix = std::iter::once(root_name.as_str()) - .chain(access.field_path.iter().map(|p| p.as_str())) - .collect::>(); - - for leaf_idx in 0..schema_descr.num_columns() { - let col = schema_descr.column(leaf_idx); - let col_path = col.path().parts(); - - // A leaf matches if its path starts with our prefix. - // e.g., prefix=["s", "value"] matches leaf path ["s", "value"] - // prefix=["s", "outer"] matches ["s", "outer", "inner"] - - // a leaf matches if its path starts with our prefix - // for example: prefix=["s", "value"] matches leaf path ["s", "value"] - // prefix=["s", "outer"] matches ["s", "outer", "inner"] - let leaf_matches_path = col_path.len() >= prefix.len() - && col_path.iter().zip(prefix.iter()).all(|(a, b)| a == b); - - if leaf_matches_path { - leaf_indices.push(leaf_idx); - } - } - } - - leaf_indices -} - -/// Builds a filter schema that includes only the fields actually accessed by the -/// filter expression. -/// -/// For regular (non-struct) columns, the full field type is used. -/// For struct columns accessed via `get_field`, a pruned struct type is created -/// containing only the fields along the access path. Note: it must match the schema -/// that the Parquet reader produces when projecting specific struct leaves -fn build_filter_schema( - file_schema: &Schema, - regular_indices: &[usize], - struct_field_accesses: &[StructFieldAccess], -) -> SchemaRef { - let regular_set: BTreeSet = regular_indices.iter().copied().collect(); - let paths_by_root = group_access_paths_by_root(struct_field_accesses); - - let all_indices = regular_indices - .iter() - .copied() - .chain(paths_by_root.keys().copied()) - .collect::>(); - - let fields = all_indices - .iter() - .map(|&idx| { - let field = file_schema.field(idx); - - // if this column appears as a regular (whole-column) reference, - // keep the full type - // - // Pruning is only valid when the column is accessed exclusively - // through struct field accesses - if regular_set.contains(&idx) { - return Arc::new(field.clone()); - } - - let Some(field_paths) = paths_by_root.get(&idx) else { - return Arc::new(field.clone()); - }; - - let pruned_data_type = prune_struct_type(field.data_type(), field_paths); - Arc::new(Field::new( - field.name(), - pruned_data_type, - field.is_nullable(), - )) - }) - .collect::>(); - - Arc::new(Schema::new_with_metadata( - fields, - file_schema.metadata().clone(), - )) -} - -/// Groups struct field access paths once for the root schema level. -/// -/// Each map entry contains the complete field paths accessed below a root -/// column. Recursive pruning groups these paths by their next component at each -/// nested struct level. -fn group_access_paths_by_root( - struct_field_accesses: &[StructFieldAccess], -) -> BTreeMap> { - let mut paths_by_root: BTreeMap> = BTreeMap::new(); - for StructFieldAccess { - root_index, - field_path, - } in struct_field_accesses - { - paths_by_root - .entry(*root_index) - .or_default() - .push(field_path.as_slice()); - } - - paths_by_root -} - -/// Groups access paths once for the current struct level. -/// -/// The map key is the field name at this level. The map value is the list of -/// remaining path suffixes below that field. An empty suffix means the access -/// path terminates at that field, so the full field must be preserved. -fn group_paths_by_next_field<'a>( - paths: &'a [&'a [String]], -) -> BTreeMap<&'a str, Vec<&'a [String]>> { - let mut paths_by_field: BTreeMap<&str, Vec<&[String]>> = BTreeMap::new(); - for path in paths { - if let Some((field, sub_path)) = path.split_first() { - paths_by_field - .entry(field.as_str()) - .or_default() - .push(sub_path); - } - } - - paths_by_field -} - -fn prune_struct_type(dt: &DataType, paths: &[&[String]]) -> DataType { - let DataType::Struct(fields) = dt else { - return dt.clone(); - }; - - let paths_by_field = group_paths_by_next_field(paths); - - let pruned_fields = fields - .iter() - .filter_map(|f| { - let sub_paths = paths_by_field.get(f.name().as_str())?; - - let out = if sub_paths.iter().any(|sub| sub.is_empty()) { - // Leaf of access path — keep the field as-is. - Arc::clone(f) - } else { - // Recurse into nested struct. - let pruned = prune_struct_type(f.data_type(), sub_paths); - Arc::new(Field::new(f.name(), pruned, f.is_nullable())) - }; - - Some(out) - }) - .collect::>(); - - DataType::Struct(pruned_fields.into()) -} - /// Checks if a predicate expression can be pushed down to the parquet decoder. /// /// Returns `true` if all columns referenced by the expression: @@ -1194,8 +893,6 @@ mod test { use parquet::file::reader::{FileReader, SerializedFileReader}; use tempfile::NamedTempFile; - use datafusion_physical_expr::expressions::Column as PhysicalColumn; - // List predicates used by the decoder should be accepted for pushdown #[test] fn test_filter_candidate_builder_supports_list_types() { @@ -2052,86 +1749,6 @@ mod test { assert_eq!(file_metrics.pushdown_rows_matched.value(), 2); } - #[test] - fn projection_read_plan_preserves_full_struct() { - // Schema: id (Int32), s (Struct{value: Int32, label: Utf8}) - // Parquet leaves: id=0, s.value=1, s.label=2 - let struct_fields: Fields = vec![ - Arc::new(Field::new("value", DataType::Int32, false)), - Arc::new(Field::new("label", DataType::Utf8, false)), - ] - .into(); - - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Int32, false), - Field::new("s", DataType::Struct(struct_fields.clone()), false), - ])); - - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(Int32Array::from(vec![1, 2, 3])), - Arc::new(StructArray::new( - struct_fields, - vec![ - Arc::new(Int32Array::from(vec![10, 20, 30])) as _, - Arc::new(StringArray::from(vec!["a", "b", "c"])) as _, - ], - None, - )), - ], - ) - .unwrap(); - - let file = NamedTempFile::new().expect("temp file"); - let mut writer = - ArrowWriter::try_new(file.reopen().unwrap(), Arc::clone(&schema), None) - .expect("writer"); - writer.write(&batch).expect("write batch"); - writer.close().expect("close writer"); - - let reader_file = file.reopen().expect("reopen file"); - let builder = ParquetRecordBatchReaderBuilder::try_new(reader_file) - .expect("reader builder"); - let metadata = builder.metadata().clone(); - let file_schema = builder.schema().clone(); - let schema_descr = metadata.file_metadata().schema_descr(); - - // Simulate SELECT * output projection: Column("id") and Column("s") - // Plus a get_field(s, 'value') expression from the pushed-down filter - let exprs: Vec> = vec![ - Arc::new(PhysicalColumn::new("id", 0)), - Arc::new(PhysicalColumn::new("s", 1)), - logical2physical( - &get_field().call(vec![ - col("s"), - Expr::Literal(ScalarValue::Utf8(Some("value".to_string())), None), - ]), - &file_schema, - ), - ]; - - let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); - - // The projected schema must have the FULL struct type because Column("s") - // is in the projection. It should NOT be narrowed to Struct{value: Int32}. - let s_field = read_plan.projected_schema.field_with_name("s").unwrap(); - assert_eq!( - s_field.data_type(), - &DataType::Struct( - vec![ - Arc::new(Field::new("value", DataType::Int32, false)), - Arc::new(Field::new("label", DataType::Utf8, false)), - ] - .into() - ), - ); - - // all3 Parquet leaves should be in the projection mask - let expected_mask = ProjectionMask::leaves(schema_descr, [0, 1, 2]); - assert_eq!(read_plan.projection_mask, expected_mask,); - } - /// Sanity check that the given expression could be evaluated against the given schema without any errors. /// This will fail if the expression references columns that are not in the schema or if the types of the columns are incompatible, etc. fn check_expression_can_evaluate_against_schema( diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index 3443b08475e0d..37e73c055e841 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -639,6 +639,10 @@ impl FileSource for ParquetSource { enable_page_index: self.enable_page_index(), enable_bloom_filter: self.bloom_filter_on_read(), enable_row_group_stats_pruning: self.table_parquet_options.global.pruning, + nested_projection_pruning: self + .table_parquet_options + .global + .nested_projection_pruning, coerce_int96, coerce_int96_tz, #[cfg(feature = "parquet_encryption")] diff --git a/datafusion/proto-common/proto/datafusion_common.proto b/datafusion/proto-common/proto/datafusion_common.proto index 7fff5b6b715ff..5fbea269eebb8 100644 --- a/datafusion/proto-common/proto/datafusion_common.proto +++ b/datafusion/proto-common/proto/datafusion_common.proto @@ -568,6 +568,7 @@ message ParquetOptions { uint64 maximum_parallel_row_group_writers = 24; // default = 1 uint64 maximum_buffered_record_batches_per_stream = 25; // default = 2 bool bloom_filter_on_read = 26; // default = true + bool nested_projection_pruning = 38; // default = true bool bloom_filter_on_write = 27; // default = false bool schema_force_view_types = 28; // default = false bool binary_as_string = 29; // default = false diff --git a/datafusion/proto-common/src/from_proto/mod.rs b/datafusion/proto-common/src/from_proto/mod.rs index 97cc9af230105..8c41dd655fe5c 100644 --- a/datafusion/proto-common/src/from_proto/mod.rs +++ b/datafusion/proto-common/src/from_proto/mod.rs @@ -1063,6 +1063,7 @@ impl TryFrom<&protobuf::ParquetOptions> for ParquetOptions { pushdown_filters: value.pushdown_filters, reorder_filters: value.reorder_filters, force_filter_selections: value.force_filter_selections, + nested_projection_pruning: value.nested_projection_pruning, data_pagesize_limit: value.data_pagesize_limit as usize, write_batch_size: value.write_batch_size as usize, writer_version: value.writer_version.parse().map_err(|e| { diff --git a/datafusion/proto-common/src/generated/pbjson.rs b/datafusion/proto-common/src/generated/pbjson.rs index 963faa5a3e9cb..21c2f3426010b 100644 --- a/datafusion/proto-common/src/generated/pbjson.rs +++ b/datafusion/proto-common/src/generated/pbjson.rs @@ -6388,6 +6388,9 @@ impl serde::Serialize for ParquetOptions { if self.bloom_filter_on_read { len += 1; } + if self.nested_projection_pruning { + len += 1; + } if self.bloom_filter_on_write { len += 1; } @@ -6502,6 +6505,9 @@ impl serde::Serialize for ParquetOptions { if self.bloom_filter_on_read { struct_ser.serialize_field("bloomFilterOnRead", &self.bloom_filter_on_read)?; } + if self.nested_projection_pruning { + struct_ser.serialize_field("nestedProjectionPruning", &self.nested_projection_pruning)?; + } if self.bloom_filter_on_write { struct_ser.serialize_field("bloomFilterOnWrite", &self.bloom_filter_on_write)?; } @@ -6673,6 +6679,8 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "maximumBufferedRecordBatchesPerStream", "bloom_filter_on_read", "bloomFilterOnRead", + "nested_projection_pruning", + "nestedProjectionPruning", "bloom_filter_on_write", "bloomFilterOnWrite", "schema_force_view_types", @@ -6732,6 +6740,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { MaximumParallelRowGroupWriters, MaximumBufferedRecordBatchesPerStream, BloomFilterOnRead, + NestedProjectionPruning, BloomFilterOnWrite, SchemaForceViewTypes, BinaryAsString, @@ -6788,6 +6797,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "maximumParallelRowGroupWriters" | "maximum_parallel_row_group_writers" => Ok(GeneratedField::MaximumParallelRowGroupWriters), "maximumBufferedRecordBatchesPerStream" | "maximum_buffered_record_batches_per_stream" => Ok(GeneratedField::MaximumBufferedRecordBatchesPerStream), "bloomFilterOnRead" | "bloom_filter_on_read" => Ok(GeneratedField::BloomFilterOnRead), + "nestedProjectionPruning" | "nested_projection_pruning" => Ok(GeneratedField::NestedProjectionPruning), "bloomFilterOnWrite" | "bloom_filter_on_write" => Ok(GeneratedField::BloomFilterOnWrite), "schemaForceViewTypes" | "schema_force_view_types" => Ok(GeneratedField::SchemaForceViewTypes), "binaryAsString" | "binary_as_string" => Ok(GeneratedField::BinaryAsString), @@ -6842,6 +6852,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { let mut maximum_parallel_row_group_writers__ = None; let mut maximum_buffered_record_batches_per_stream__ = None; let mut bloom_filter_on_read__ = None; + let mut nested_projection_pruning__ = None; let mut bloom_filter_on_write__ = None; let mut schema_force_view_types__ = None; let mut binary_as_string__ = None; @@ -6952,6 +6963,12 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { } bloom_filter_on_read__ = Some(map_.next_value()?); } + GeneratedField::NestedProjectionPruning => { + if nested_projection_pruning__.is_some() { + return Err(serde::de::Error::duplicate_field("nestedProjectionPruning")); + } + nested_projection_pruning__ = Some(map_.next_value()?); + } GeneratedField::BloomFilterOnWrite => { if bloom_filter_on_write__.is_some() { return Err(serde::de::Error::duplicate_field("bloomFilterOnWrite")); @@ -7106,6 +7123,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { maximum_parallel_row_group_writers: maximum_parallel_row_group_writers__.unwrap_or_default(), maximum_buffered_record_batches_per_stream: maximum_buffered_record_batches_per_stream__.unwrap_or_default(), bloom_filter_on_read: bloom_filter_on_read__.unwrap_or_default(), + nested_projection_pruning: nested_projection_pruning__.unwrap_or_default(), bloom_filter_on_write: bloom_filter_on_write__.unwrap_or_default(), schema_force_view_types: schema_force_view_types__.unwrap_or_default(), binary_as_string: binary_as_string__.unwrap_or_default(), diff --git a/datafusion/proto-common/src/generated/prost.rs b/datafusion/proto-common/src/generated/prost.rs index 93b97c4f1376c..c7fed80871cb2 100644 --- a/datafusion/proto-common/src/generated/prost.rs +++ b/datafusion/proto-common/src/generated/prost.rs @@ -844,6 +844,9 @@ pub struct ParquetOptions { /// default = true #[prost(bool, tag = "26")] pub bloom_filter_on_read: bool, + /// default = true + #[prost(bool, tag = "38")] + pub nested_projection_pruning: bool, /// default = false #[prost(bool, tag = "27")] pub bloom_filter_on_write: bool, diff --git a/datafusion/proto-common/src/to_proto/mod.rs b/datafusion/proto-common/src/to_proto/mod.rs index a6fa13ca7479c..57051dac7a18f 100644 --- a/datafusion/proto-common/src/to_proto/mod.rs +++ b/datafusion/proto-common/src/to_proto/mod.rs @@ -912,6 +912,7 @@ impl TryFrom<&ParquetOptions> for protobuf::ParquetOptions { pushdown_filters: value.pushdown_filters, reorder_filters: value.reorder_filters, force_filter_selections: value.force_filter_selections, + nested_projection_pruning: value.nested_projection_pruning, data_pagesize_limit: value.data_pagesize_limit as u64, write_batch_size: value.write_batch_size as u64, writer_version: value.writer_version.to_string(), diff --git a/datafusion/proto-models/src/generated/datafusion_proto_common.rs b/datafusion/proto-models/src/generated/datafusion_proto_common.rs index 93b97c4f1376c..c7fed80871cb2 100644 --- a/datafusion/proto-models/src/generated/datafusion_proto_common.rs +++ b/datafusion/proto-models/src/generated/datafusion_proto_common.rs @@ -844,6 +844,9 @@ pub struct ParquetOptions { /// default = true #[prost(bool, tag = "26")] pub bloom_filter_on_read: bool, + /// default = true + #[prost(bool, tag = "38")] + pub nested_projection_pruning: bool, /// default = false #[prost(bool, tag = "27")] pub bloom_filter_on_write: bool, diff --git a/datafusion/proto/src/logical_plan/file_formats.rs b/datafusion/proto/src/logical_plan/file_formats.rs index 8940b16bf83f5..80ddc3aafd4c5 100644 --- a/datafusion/proto/src/logical_plan/file_formats.rs +++ b/datafusion/proto/src/logical_plan/file_formats.rs @@ -436,6 +436,9 @@ mod parquet { parquet_options::EncodingOpt::Encoding(encoding) }), bloom_filter_on_read: global_options.global.bloom_filter_on_read, + nested_projection_pruning: global_options + .global + .nested_projection_pruning, bloom_filter_on_write: global_options.global.bloom_filter_on_write, bloom_filter_fpp_opt: global_options.global.bloom_filter_fpp.map(|fpp| { parquet_options::BloomFilterFppOpt::BloomFilterFpp(fpp) @@ -590,6 +593,7 @@ mod parquet { } }), bloom_filter_on_read: proto.bloom_filter_on_read, + nested_projection_pruning: proto.nested_projection_pruning, bloom_filter_on_write: proto.bloom_filter_on_write, bloom_filter_fpp: proto .bloom_filter_fpp_opt diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index bf45564e26333..669d4e05a820d 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -259,6 +259,7 @@ datafusion.execution.parquet.max_row_group_size 1048576 datafusion.execution.parquet.maximum_buffered_record_batches_per_stream 2 datafusion.execution.parquet.maximum_parallel_row_group_writers 1 datafusion.execution.parquet.metadata_size_hint 524288 +datafusion.execution.parquet.nested_projection_pruning true datafusion.execution.parquet.pruning true datafusion.execution.parquet.pushdown_filters false datafusion.execution.parquet.reorder_filters false @@ -418,6 +419,7 @@ datafusion.execution.parquet.max_row_group_size 1048576 (writing) Target maximum datafusion.execution.parquet.maximum_buffered_record_batches_per_stream 2 (writing) By default parallel parquet writer is tuned for minimum memory usage in a streaming execution plan. You may see a performance benefit when writing large parquet files by increasing maximum_parallel_row_group_writers and maximum_buffered_record_batches_per_stream if your system has idle cores and can tolerate additional memory usage. Boosting these values is likely worthwhile when writing out already in-memory data, such as from a cached data frame. datafusion.execution.parquet.maximum_parallel_row_group_writers 1 (writing) By default parallel parquet writer is tuned for minimum memory usage in a streaming execution plan. You may see a performance benefit when writing large parquet files by increasing maximum_parallel_row_group_writers and maximum_buffered_record_batches_per_stream if your system has idle cores and can tolerate additional memory usage. Boosting these values is likely worthwhile when writing out already in-memory data, such as from a cached data frame. datafusion.execution.parquet.metadata_size_hint 524288 (reading) If specified, the parquet reader will try and fetch the last `size_hint` bytes of the parquet file optimistically. If not specified, two reads are required: One read to fetch the 8-byte parquet footer and another to fetch the metadata length encoded in the footer Default setting to 512 KiB, which should be sufficient for most parquet files, it can reduce one I/O operation per parquet file. If the metadata is larger than the hint, two reads will still be performed. +datafusion.execution.parquet.nested_projection_pruning true (reading) If true, when a projected nested column (struct, or struct nested in lists) is read through a cast to a narrower nested type — as happens when the table schema declares fewer struct fields than the parquet file contains — the reader only fetches and decodes the leaf columns the cast retains, instead of the entire column datafusion.execution.parquet.pruning true (reading) If true, the parquet reader attempts to skip entire row groups based on the predicate in the query and the metadata (min/max values) stored in the parquet file datafusion.execution.parquet.pushdown_filters false (reading) If true, filter expressions are be applied during the parquet decoding operation to reduce the number of rows decoded. This optimization is sometimes called "late materialization". datafusion.execution.parquet.reorder_filters false (reading) If true, filter expressions evaluated during the parquet decoding operation will be reordered heuristically to minimize the cost of evaluation. If false, the filters are applied in the same order as written in the query diff --git a/datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt b/datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt new file mode 100644 index 0000000000000..e9aa093639e31 --- /dev/null +++ b/datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt @@ -0,0 +1,105 @@ +# 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. + +########## +# Nested projection pruning: a table whose declared nested type is narrower +# than the parquet file's physical type reads only the declared leaves +# (`datafusion.execution.parquet.nested_projection_pruning`, default true). +# Results must be identical with the setting on and off; the bytes-scanned +# assertions live in the Rust tests +# (datafusion/core/tests/parquet/expr_adapter.rs). +########## + +# The file contains events: ARRAY> and +# s: STRUCT; the tables below declare narrower nested types. +statement ok +COPY ( + SELECT id, events, s + FROM (VALUES + (1, [named_struct('x', 10, 'y', 'a1', 'pad_a', 'p', 'pad_b', 'q')], + named_struct('x', 100, 'y', 's1', 'pad', 'sp1')), + (2, [named_struct('x', 20, 'y', 'b1', 'pad_a', 'p', 'pad_b', 'q'), + named_struct('x', 21, 'y', 'b2', 'pad_a', 'p', 'pad_b', 'q')], + named_struct('x', 200, 'y', 's2', 'pad', 'sp2')), + (3, NULL, + NULL) + ) AS t(id, events, s) +) TO 'test_files/scratch/parquet_nested_schema_pruning/wide.parquet' +STORED AS PARQUET; + +# Declared schema drops pad_a/pad_b from the list elements and pad from the +# struct, declares x as BIGINT (the file has INT), and adds a z column that +# does not exist in the file. +statement ok +CREATE EXTERNAL TABLE narrow ( + id INT, + events ARRAY>, + s STRUCT +) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/wide.parquet'; + +query I?? +SELECT id, events, s FROM narrow ORDER BY id; +---- +1 [{x: 10, y: a1, z: NULL}] {x: 100, y: s1} +2 [{x: 20, y: b1, z: NULL}, {x: 21, y: b2, z: NULL}] {x: 200, y: s2} +3 NULL NULL + +# Struct-level nullability is preserved: row 3's struct is NULL, not a +# struct of NULLs. +query IBB +SELECT id, events IS NULL, s IS NULL FROM narrow ORDER BY id; +---- +1 false false +2 false false +3 true true + +query II +SELECT id, s['x'] FROM narrow ORDER BY id; +---- +1 100 +2 200 +3 NULL + +query II +SELECT id, e['x'] FROM (SELECT id, unnest(events) AS e FROM narrow) ORDER BY id, e['x']; +---- +1 10 +2 20 +2 21 + +# Same results with the setting disabled (the unpruned read path). +statement ok +SET datafusion.execution.parquet.nested_projection_pruning = false; + +query I?? +SELECT id, events, s FROM narrow ORDER BY id; +---- +1 [{x: 10, y: a1, z: NULL}] {x: 100, y: s1} +2 [{x: 20, y: b1, z: NULL}, {x: 21, y: b2, z: NULL}] {x: 200, y: s2} +3 NULL NULL + +query IBB +SELECT id, events IS NULL, s IS NULL FROM narrow ORDER BY id; +---- +1 false false +2 false false +3 true true + +statement ok +SET datafusion.execution.parquet.nested_projection_pruning = true; diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 03340c366d70f..d71b06b6d6d63 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -92,6 +92,7 @@ The following configuration settings are available: | datafusion.execution.parquet.coerce_int96 | NULL | (reading) If true, parquet reader will read columns of physical type int96 as originating from a different resolution than nanosecond. This is useful for reading data from systems like Spark which stores microsecond resolution timestamps in an int96 allowing it to write values with a larger date range than 64-bit timestamps with nanosecond resolution. | | datafusion.execution.parquet.coerce_int96_tz | NULL | (reading) Optional timezone applied to INT96 columns when `coerce_int96` is set. When `Some`, INT96 columns coerce to `Timestamp(, Some())` instead of the default `Timestamp(, None)`. Spark and other systems write INT96 values as UTC-adjusted instants, so callers that need the resulting Arrow type to be timezone-aware (e.g. for Spark `TimestampType` semantics) should set this to `"UTC"`. No effect when `coerce_int96` is `None`. | | datafusion.execution.parquet.bloom_filter_on_read | true | (reading) Use any available bloom filters when reading parquet files | +| datafusion.execution.parquet.nested_projection_pruning | true | (reading) If true, when a projected nested column (struct, or struct nested in lists) is read through a cast to a narrower nested type — as happens when the table schema declares fewer struct fields than the parquet file contains — the reader only fetches and decodes the leaf columns the cast retains, instead of the entire column | | datafusion.execution.parquet.max_predicate_cache_size | NULL | (reading) The maximum predicate cache size, in bytes. When `pushdown_filters` is enabled, sets the maximum memory used to cache the results of predicate evaluation between filter evaluation and output generation. Decreasing this value will reduce memory usage, but may increase IO and CPU usage. None means use the default parquet reader setting. 0 means no caching. | | datafusion.execution.parquet.data_pagesize_limit | 1048576 | (writing) Sets best effort maximum size of data page in bytes | | datafusion.execution.parquet.write_batch_size | 1024 | (writing) Sets write_batch_size in rows |