This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23397-4184b0778e0c691ed490c8ab0fa2b0f6f157722c in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 840da055be3809dd994b6942d4295223264ce81f Author: Adrian Garcia Badaracco <[email protected]> AuthorDate: Mon Jul 20 14:55:22 2026 -0500 bench: parquet scan with a table schema narrower than a nested column (#23397) ## Which issue does this PR close? - Related to https://github.com/apache/datafusion-comet/issues/4859. Second of a 3-PR stack: #23396 (refactor), this benchmark, and the feature PR (nested schema pruning for the parquet reader). ## Rationale for this change When a table's declared schema is narrower than a parquet file's nested column (logical `events: LIST<STRUCT<x, y>>` over a physical `LIST<STRUCT<x, y, +8 pads>>`), the reader currently fetches and decodes **every** leaf of the column and discards the extra subfields in memory via the adapter-inserted cast. This is how engines like Spark (via Comet) communicate nested projection pruning to the scan โ as a clipped read schema โ and it is where Comet measured reading 1.35 TB where Spark read 30.9 GB for the same pruned `ReadSchema`. This PR adds a benchmark that documents the current behavior as a checked-in baseline, independent of any fix: ``` list_struct_narrow_schema: bytes_scanned=25.19 MB 3.45 ms list_struct_full_schema: bytes_scanned=25.19 MB 3.36 ms <- narrow == full today list_struct_physically_narrow: bytes_scanned= 3.32 KB 164 ยตs <- the floor ``` ## What changes are included in this PR? A criterion benchmark, `datafusion/core/benches/parquet_nested_schema_pruning.rs`, that registers the same wide `list<struct>` (and top-level struct) parquet file with both its full schema and a narrower declared schema, plus a physically-narrow file as the floor, and prints the scans' `bytes_scanned` at setup so the IO pattern is visible alongside wall time. ## Are these changes tested? It is a benchmark; it compiles under `cargo bench --no-run` and runs green. ## Are there any user-facing changes? No. ๐ค Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01KuMaRtFSPDQesuzjN5Koyd --------- Co-authored-by: Claude Fable 5 <[email protected]> --- datafusion/core/Cargo.toml | 5 + .../core/benches/parquet_nested_schema_pruning.rs | 449 +++++++++++++++++++++ 2 files changed, 454 insertions(+) diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index f1d7dc703e..8679dad9f9 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 0000000000..8db67de9ff --- /dev/null +++ b/datafusion/core/benches/parquet_nested_schema_pruning.rs @@ -0,0 +1,449 @@ +// 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<STRUCT<x, y>>` while the file contains +//! `events: LIST<STRUCT<x, y, pad_0, ..., pad_7>>`. 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 reads the parquet scan's `bytes_scanned` metric for +//! (1), (2) and (3) so the IO pattern is visible in addition to wall time, and +//! asserts the current baseline: today a narrow declared schema scans the same +//! bytes as the full schema. When nested projection pruning lands, that +//! assertion is expected to fail, which is the signal to flip it to +//! `narrow < full` (see [`assert_scan_baseline`]). + +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 criterion::{Criterion, criterion_group, criterion_main}; +use datafusion::datasource::listing::{ + ListingTable, ListingTableConfig, ListingTableConfigExt, +}; +use datafusion::physical_plan::metrics::MetricsSet; +use datafusion::physical_plan::{ExecutionPlan, collect}; +use datafusion::prelude::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. +/// +/// Derived from [`narrow_item_fields`] so the shared columns (`x`, `y`) match +/// by construction โ same names, types and nullability โ and only the extra +/// pad fields distinguish the two. +fn wide_item_fields() -> Fields { + let mut fields: Vec<Field> = narrow_item_fields() + .iter() + .map(|f| f.as_ref().clone()) + .collect(); + 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<String> = (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<ArrayRef> { + 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, + // `seed + i` keeps each pad column's values distinct from the + // others (and matches the additive seeding used above); a + // multiplier like `seed * (i + 1)` collapses to the same seed for + // every column when `seed == 0` (the first batch). + _ => pad_values(count, seed + i), + }) + .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()); +} + +/// Recursively collect the metrics of every node in `plan` into `out`. +fn gather_metrics(plan: &Arc<dyn ExecutionPlan>, out: &mut MetricsSet) { + if let Some(metrics) = plan.metrics() { + for metric in metrics.iter() { + out.push(Arc::clone(metric)); + } + } + for child in plan.children() { + gather_metrics(child, out); + } +} + +/// Execute `sql` and return the parquet scan's `bytes_scanned` metric, read +/// from the typed metrics API rather than scraped from display output (which +/// would silently break if the format ever changed). +fn scan_bytes(ctx: &SessionContext, rt: &Runtime, sql: &str) -> usize { + let df = rt.block_on(ctx.sql(sql)).unwrap(); + let plan = rt.block_on(df.create_physical_plan()).unwrap(); + // Fully drive the plan so the scan populates its metrics. + black_box( + rt.block_on(collect(Arc::clone(&plan), ctx.task_ctx())) + .unwrap(), + ); + + let mut metrics = MetricsSet::new(); + gather_metrics(&plan, &mut metrics); + metrics + .aggregate_by_name() + .sum_by_name("bytes_scanned") + .map(|v| v.as_usize()) + .expect("parquet scan should report a bytes_scanned metric") +} + +/// Report and assert the `bytes_scanned` baseline for one dataset shape. +/// +/// `narrow` selects from a wide file through a narrow declared schema, `full` +/// through the full schema, and `floor` from a physically-narrow file. Today +/// the extra leaves are fetched and discarded, so `narrow == full`; that +/// equality is the checked-in baseline. When nested projection pruning lands, +/// `narrow` should drop toward `floor` and this assertion is expected to fail โ +/// the signal to flip it to `assert!(narrow < full)`. +fn assert_scan_baseline( + ctx: &SessionContext, + rt: &Runtime, + label: &str, + narrow_sql: &str, + full_sql: &str, + floor_sql: &str, +) { + let narrow = scan_bytes(ctx, rt, narrow_sql); + let full = scan_bytes(ctx, rt, full_sql); + let floor = scan_bytes(ctx, rt, floor_sql); + println!( + "{label}: bytes_scanned narrow_schema={narrow} full_schema={full} \ + physically_narrow={floor}" + ); + assert_eq!( + narrow, full, + "{label}: narrow declared schema scanned {narrow} bytes vs {full} for \ + the full schema. The baseline is that a narrow schema still reads \ + every leaf, so these should be equal; if narrow is now smaller, \ + nested projection pruning has likely landed โ flip this to \ + `assert!(narrow < full)`." + ); +} + +struct Fixture { + ctx: SessionContext, + rt: Runtime, + _files: Vec<NamedTempFile>, +} + +/// Tables: +/// `<name>_narrow_schema`: wide file, narrow declared schema +/// `<name>_full_schema`: wide file, full declared schema +/// `<name>_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 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()), + ); + + Fixture { + ctx, + 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); + + assert_scan_baseline( + ctx, + rt, + "list_struct", + "SELECT events FROM list_struct_narrow_schema", + "SELECT events FROM list_struct_full_schema", + "SELECT events FROM list_struct_physically_narrow", + ); + + 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")) + }); + + // 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); + + assert_scan_baseline( + ctx, + rt, + "top_level_struct", + "SELECT s FROM struct_narrow_schema", + "SELECT s FROM struct_full_schema", + "SELECT s FROM struct_physically_narrow", + ); + + 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")) + }); + + // 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); --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
