mbutrovich commented on code in PR #23397: URL: https://github.com/apache/datafusion/pull/23397#discussion_r3608558771
########## datafusion/core/benches/parquet_nested_schema_pruning.rs: ########## @@ -0,0 +1,390 @@ +// 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 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::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<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, + _ => 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=") Review Comment: `bytes_scanned` is a real `Count` metric, so we can read it in process instead of scraping `pretty_format_batches` output. The current parse returns `<not found>` if the display format ever shifts, which is easy to miss. ```rust let plan = df.create_physical_plan().await?; let _ = collect(plan.clone(), task_ctx).await?; let bytes = plan.metrics().unwrap() .aggregate_by_name() .sum_by_name("bytes_scanned") .map(|v| v.as_usize()); ``` `plan.metrics().unwrap().aggregate_by_name()` is already used this way at `tests/sql/explain_analyze.rs:162`. It hands back a typed `usize`, which also makes it easy to assert the narrow == full baseline instead of just printing it. ########## datafusion/core/benches/parquet_nested_schema_pruning.rs: ########## @@ -0,0 +1,390 @@ +// 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 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::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<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, + _ => 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) { Review Comment: The baseline claim, that narrow scans the same bytes as full today, is printed but never enforced, so a later change could shift it and nobody would notice. The sibling bench asserts its physical-read invariants (`parquet_struct_projection.rs:167-178`). Can we do the same on the relationship here? ```rust assert_eq!(narrow_bytes, full_bytes, "narrow declared schema still reads the full column today"); assert!(physically_narrow_bytes < narrow_bytes); ``` The nice part: when PR-3 lands, the bench fails right here and prompts you to flip it to `assert!(narrow_bytes < full_bytes)`, so the fix documents itself. Asserting the relationship instead of an absolute byte count keeps it from breaking on formatting changes. (See the related note on line 233 about reading the value from the metrics API, which gives you a typed number to assert on.) ########## datafusion/core/benches/parquet_nested_schema_pruning.rs: ########## @@ -0,0 +1,390 @@ +// 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 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::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<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, + _ => pad_values(count, seed * (i + 1)), Review Comment: `pad_values` already makes columns distinct through its per-element index, so the `seed * (i + 1)` multiplier isn't doing anything load-bearing. `seed + i` matches the additive convention used elsewhere and reads more plainly. An explicit `name if name.starts_with("pad_")` arm would also state the intent instead of leaning on the `_` catch-all. Minor. ########## datafusion/core/benches/parquet_nested_schema_pruning.rs: ########## @@ -0,0 +1,390 @@ +// 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 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::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; Review Comment: `NUM_BATCHES` and `ROW_GROUP_ROW_COUNT` (line 63) match `parquet_struct_projection.rs:45,47`, and `ROWS_PER_BATCH` is that file's `WRITE_RECORD_BATCH_SIZE`. Natural to fold into a shared bench module if you decide to dedupe the helpers. ########## datafusion/core/benches/parquet_nested_schema_pruning.rs: ########## @@ -0,0 +1,390 @@ +// 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 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::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<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, + _ => 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( Review Comment: This is close to `parquet_struct_projection.rs:137-188`, and there's a third copy in `parquet_struct_query.rs`. Benches can share code across criterion binaries here: `benches/data_utils/mod.rs` is reused by around nine benches via `mod data_utils; use data_utils::...`. A small shared module for `generate_file`, `query`, and a metrics-based bytes reader would keep the three `parquet_struct*` benches in sync. Only worth it if the duplication bugs you; just noting the option. ########## datafusion/core/benches/parquet_nested_schema_pruning.rs: ########## @@ -0,0 +1,390 @@ +// 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 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::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<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, + _ => 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) { Review Comment: This helper is duplicated in both sibling benches (`parquet_struct_projection.rs`, `parquet_struct_query.rs`), so it's a candidate for the shared bench module too. If it does get hoisted, keep this borrowing version rather than the cloning one. ########## datafusion/core/benches/parquet_nested_schema_pruning.rs: ########## @@ -0,0 +1,390 @@ +// 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 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::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), Review Comment: `x` is `Int64, false` here but `Int64, true` in `narrow_item_fields` (line 71), while `y` is `true` in both. Reads like unintended drift for the same logical column. It's harmless today, but deriving wide from narrow keeps the shared prefix aligned by construction: ```rust 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) } ``` ########## datafusion/core/Cargo.toml: ########## @@ -247,6 +247,11 @@ harness = false name = "parquet_struct_query" required-features = ["parquet"] +[[bench]] Review Comment: More of a question than a request. This measures the same thing as `parquet_struct_projection.rs` at a complementary altitude, schema-level narrowing versus expression-level pruning, so a `schema_narrowed` group inside that file would put the whole story in one place and land PR-3's fix next to the assert it flips. The trade-off is a longer file and generalizing the helpers (`register_parquet` versus `ListingTableConfig::with_schema`). Fine either way; just worth a look before committing to a second binary. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
