diff --git a/datafusion/core/benches/parquet_nested_schema_pruning.rs b/datafusion/core/benches/parquet_nested_schema_pruning.rs index 8db67de9ffa9b..de4f0a57a5c41 100644 --- a/datafusion/core/benches/parquet_nested_schema_pruning.rs +++ b/datafusion/core/benches/parquet_nested_schema_pruning.rs @@ -36,10 +36,9 @@ //! //! 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`]). +//! asserts that nested projection pruning keeps the narrow declared schema's +//! scan well below the full schema's, close to the physically-narrow floor +//! (see [`assert_scan_prunes`]). use arrow::array::{ ArrayRef, Int32Array, Int64Array, ListArray, StringArray, StructArray, @@ -267,15 +266,14 @@ fn scan_bytes(ctx: &SessionContext, rt: &Runtime, sql: &str) -> usize { .expect("parquet scan should report a bytes_scanned metric") } -/// Report and assert the `bytes_scanned` baseline for one dataset shape. +/// Report and assert the `bytes_scanned` improvement 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( +/// through the full schema, and `floor` from a physically-narrow file. +/// Nested projection pruning clips the narrow read to the declared leaves, so +/// `narrow` should read substantially less than `full`, close to `floor`, +/// the cost of a file that never had the extra leaves to begin with. +fn assert_scan_prunes( ctx: &SessionContext, rt: &Runtime, label: &str, @@ -290,13 +288,11 @@ fn assert_scan_baseline( "{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)`." + assert!( + narrow * 2 < full, + "{label}: expected the narrow declared schema to read less than half \ + of the full schema's {full} bytes (physically-narrow floor is \ + {floor} bytes), but it read {narrow}" ); } @@ -363,7 +359,7 @@ 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( + assert_scan_prunes( ctx, rt, "list_struct", @@ -410,7 +406,7 @@ 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( + assert_scan_prunes( ctx, rt, "top_level_struct", diff --git a/datafusion/core/tests/parquet/expr_adapter.rs b/datafusion/core/tests/parquet/expr_adapter.rs index 535828fa29c2f..f3979af9471d4 100644 --- a/datafusion/core/tests/parquet/expr_adapter.rs +++ b/datafusion/core/tests/parquet/expr_adapter.rs @@ -1204,3 +1204,671 @@ 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, instead of reading the whole column and +// discarding the extra subfields in the adapter-inserted cast. +// +// Each test registers two tables against the *same* physical file: `t_narrow` +// (the declared schema under test) and `t_full` (the file's own physical +// schema, so no cast is inserted and the scan always reads every leaf). That +// gives a same-context upper bound to compare `bytes_scanned` against, +// without needing a config flag to disable pruning. +// --------------------------------------------------------------------------- + +mod nested_projection_pruning { + use super::*; + 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 narrow table schema will not mention. `x` is Int32 in the + /// file (the narrow schema 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 one 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, + ) + } + + fn wide_list_schema() -> SchemaRef { + let item = Arc::new(Field::new( + "item", + DataType::Struct(wide_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`, `events List` (one element per row). + fn wide_list_batch() -> RecordBatch { + let schema = wide_list_schema(); + let item = match schema.field(1).data_type() { + DataType::List(item) => Arc::clone(item), + other => unreachable!("expected List, got {other:?}"), + }; + let events = ListArray::new( + item, + OffsetBuffer::from_lengths(std::iter::repeat_n(1, NUM_ELEMENTS)), + Arc::new(wide_struct_values(None)), + None, + ); + RecordBatch::try_new( + Arc::clone(&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), + ])) + } + + fn wide_struct_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("s", DataType::Struct(wide_item_fields()), 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::>()); + RecordBatch::try_new( + wide_struct_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), + ])) + } + + /// Registers `t_narrow` (the schema under test) and `t_full` (the file's + /// own physical schema, so no cast is inserted) against the same store. + async fn register_narrow_and_full( + ctx: &SessionContext, + store: Arc, + narrow_schema: SchemaRef, + full_schema: SchemaRef, + ) { + let store_url = ObjectStoreUrl::parse("memory://").unwrap(); + ctx.register_object_store(store_url.as_ref(), store); + + for (name, schema) in [("t_narrow", narrow_schema), ("t_full", full_schema)] { + let config = ListingTableConfig::new( + ListingTableUrl::parse("memory:///data/").unwrap(), + ) + .infer_options(&ctx.state()) + .await + .unwrap() + .with_schema(schema) + .with_expr_adapter_factory(Arc::new(DefaultPhysicalExprAdapterFactory)); + let table = ListingTable::try_new(config).unwrap(); + ctx.register_table(name, Arc::new(table)).unwrap(); + } + } + + async fn setup_with_config( + batches: Vec<(&str, RecordBatch)>, + narrow_schema: SchemaRef, + full_schema: SchemaRef, + cfg: SessionConfig, + ) -> 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 ctx = SessionContext::new_with_config(cfg); + register_narrow_and_full(&ctx, store, narrow_schema, full_schema).await; + ctx + } + + async fn setup( + batches: Vec<(&str, RecordBatch)>, + narrow_schema: SchemaRef, + full_schema: SchemaRef, + ) -> SessionContext { + setup_with_config( + batches, + narrow_schema, + full_schema, + SessionConfig::new().with_collect_statistics(false), + ) + .await + } + + 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 `narrow_sql` against `t_narrow` and `full_sql` against `t_full`; + /// assert the narrow scan read strictly less than half of the full + /// scan's bytes (the pads dominate the file), and return the narrow + /// scan's results for correctness assertions. + /// + /// The two SQL strings need not have the same shape: a `get_field` over + /// a narrowed struct clips to the *cast target*, not further down to the + /// specific field accessed (see `prunes_get_field_on_narrowed_struct`), + /// so comparing against the same `get_field` query on `t_full` would + /// unfairly compare this clip against `get_field`'s own, more precise, + /// single-leaf pruning (which only applies when there is no cast in the + /// way). Callers that aren't in that situation can just pass the same + /// query shape with the table name substituted. + async fn assert_prunes( + batches: Vec<(&str, RecordBatch)>, + narrow_schema: SchemaRef, + full_schema: SchemaRef, + narrow_sql: &str, + full_sql: &str, + ) -> Vec { + let ctx = setup(batches, narrow_schema, full_schema).await; + + let (result_narrow, metrics_narrow) = run(&ctx, narrow_sql).await; + let (_result_full, metrics_full) = run(&ctx, full_sql).await; + + let (narrow_bytes, full_bytes) = + (bytes_scanned(&metrics_narrow), bytes_scanned(&metrics_full)); + assert!( + narrow_bytes * 2 < full_bytes, + "expected pruned scan to read less than half of {full_bytes} bytes, \ + read {narrow_bytes}: {narrow_sql}" + ); + result_narrow + } + + #[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(), + wide_list_schema(), + "SELECT events FROM t_narrow ORDER BY id", + "SELECT events FROM t_full 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(), + wide_struct_schema(), + "SELECT s FROM t_narrow ORDER BY id", + "SELECT s FROM t_full 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(), + wide_struct_schema(), + "SELECT id, s IS NULL AS s_null, s FROM t_narrow ORDER BY id", + "SELECT id, s IS NULL AS s_null, s FROM t_full 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 (every + /// field the *narrow* schema declares), not further down to just `x`. + /// The fair "no clipping happened" baseline is therefore reading every + /// physical leaf of `s` (`SELECT s FROM t_full`), not the same + /// `get_field` query against `t_full`. That query needs no cast at all + /// and takes `get_field`'s own, more precise, single-leaf pushdown path. + #[tokio::test] + async fn prunes_get_field_on_narrowed_struct() { + let batches = assert_prunes( + vec![("wide.parquet", wide_struct_batch())], + narrow_struct_table_schema(), + wide_struct_schema(), + "SELECT s['x'] AS x FROM t_narrow ORDER BY id", + "SELECT s FROM t_full 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(), + wide_struct_schema(), + "SELECT s, s['y'] AS y FROM t_narrow ORDER BY id", + "SELECT s, s['y'] AS y FROM t_full 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 mut cfg = SessionConfig::new().with_collect_statistics(false); + cfg.options_mut().execution.parquet.pushdown_filters = true; + let ctx = setup_with_config( + vec![("wide.parquet", wide_list_batch())], + narrow_list_table_schema(), + wide_list_schema(), + cfg, + ) + .await; + + let filter = "WHERE id >= 32 ORDER BY id"; + let (result_narrow, metrics_narrow) = + run(&ctx, &format!("SELECT events FROM t_narrow {filter}")).await; + let (_result_full, metrics_full) = + run(&ctx, &format!("SELECT events FROM t_full {filter}")).await; + + let combined = + concat_batches(&result_narrow[0].schema(), &result_narrow).unwrap(); + assert_eq!(combined.num_rows(), NUM_ELEMENTS / 2); + assert!(bytes_scanned(&metrics_narrow) * 2 < bytes_scanned(&metrics_full)); + } + + /// 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 store = Arc::new(InMemory::new()) as Arc; + write_parquet(wide_list_batch(), Arc::clone(&store), "data/wide.parquet").await; + write_parquet(narrow_batch, Arc::clone(&store), "data/narrow.parquet").await; + + let ctx = test_context(); + register_memory_listing_table( + &ctx, + store, + "memory:///data/", + narrow_list_table_schema(), + ) + .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); + } + + /// Regression test for the exact shape reported in + /// `datafusion-comet#4859`: a two-level `array>>>` column, with a dropped struct sibling + /// (`latency_parts`), a dropped map sibling (`feature_map`), a dropped + /// nested-struct sibling (`diagnostics`), and dropped top-level sibling + /// columns (`dimension_id`, `region_code`, `raw_payload`), structurally + /// the same `ReadSchema`/`InputSchema` pair from the issue (field names + /// kept representative, not verbatim), which let Comet's production + /// query read 1.35 TB where plain Spark, given the same pruned + /// `ReadSchema`, read 30.9 GB. + #[tokio::test] + async fn comet_4859_two_level_nested_list_regression() { + use arrow::array::{Float64Array, new_null_array}; + + const NUM_ROWS: usize = 8; + const EVENTS_PER_ROW: usize = 2; + const ITEMS_PER_EVENT: usize = 2; + const NUM_EVENTS: usize = NUM_ROWS * EVENTS_PER_ROW; + const NUM_ITEMS: usize = NUM_EVENTS * ITEMS_PER_EVENT; + + // items: array> + let map_type = DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Float64, true), + ])), + false, + )), + false, + ); + let diagnostics_type = DataType::Struct(Fields::from(vec![ + Field::new("module_id", DataType::Utf8, true), + Field::new("trace_id", DataType::Utf8, true), + ])); + let wide_item_struct_fields = Fields::from(vec![ + Field::new("group_id", DataType::Int64, false), + Field::new("entity_id", DataType::Int64, false), + Field::new("metric_value", DataType::Float64, false), + Field::new("feature_map", map_type.clone(), true), + Field::new("diagnostics", diagnostics_type.clone(), true), + Field::new("pad", DataType::Utf8, false), + ]); + let pad_base = "x".repeat(PAD_LEN); + let items_struct = StructArray::new( + wide_item_struct_fields.clone(), + vec![ + Arc::new(Int64Array::from_iter_values(0..NUM_ITEMS as i64)), + Arc::new(Int64Array::from_iter_values( + (0..NUM_ITEMS).map(|i| 100 + i as i64), + )), + Arc::new(Float64Array::from_iter_values( + (0..NUM_ITEMS).map(|i| i as f64 * 1.5), + )), + new_null_array(&map_type, NUM_ITEMS), + new_null_array(&diagnostics_type, NUM_ITEMS), + Arc::new(StringArray::from_iter_values( + (0..NUM_ITEMS).map(|i| format!("{i:08}{pad_base}")), + )), + ], + None, + ); + let items_item_field = Arc::new(Field::new( + "item", + DataType::Struct(wide_item_struct_fields), + true, + )); + let items_list = ListArray::new( + Arc::clone(&items_item_field), + OffsetBuffer::from_lengths(std::iter::repeat_n(ITEMS_PER_EVENT, NUM_EVENTS)), + Arc::new(items_struct), + None, + ); + + // events: array> + let latency_type = DataType::Struct(Fields::from(vec![ + Field::new("queue_time_ms", DataType::Int64, true), + Field::new("retry_count", DataType::Int32, true), + ])); + let wide_event_fields = Fields::from(vec![ + Field::new("is_available", DataType::Boolean, false), + Field::new("event_time_ms", DataType::Int64, false), + Field::new("event_token", DataType::Utf8, false), + Field::new("latency_parts", latency_type.clone(), true), + Field::new("items", DataType::List(items_item_field), true), + ]); + let events_struct = StructArray::new( + wide_event_fields.clone(), + vec![ + Arc::new(BooleanArray::from_iter( + (0..NUM_EVENTS).map(|i| Some(i % 2 == 0)), + )), + Arc::new(Int64Array::from_iter_values(0..NUM_EVENTS as i64)), + Arc::new(StringArray::from_iter_values( + (0..NUM_EVENTS).map(|i| format!("token-{i}")), + )), + new_null_array(&latency_type, NUM_EVENTS), + Arc::new(items_list), + ], + None, + ); + let events_item_field = Arc::new(Field::new( + "item", + DataType::Struct(wide_event_fields), + true, + )); + let events_list = ListArray::new( + Arc::clone(&events_item_field), + OffsetBuffer::from_lengths(std::iter::repeat_n(EVENTS_PER_ROW, NUM_ROWS)), + Arc::new(events_struct), + None, + ); + + let wide_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("is_flagged", DataType::Boolean, false), + Field::new("dimension_id", DataType::Int64, true), + Field::new("region_code", DataType::Utf8, true), + Field::new( + "events", + DataType::List(Arc::clone(&events_item_field)), + true, + ), + Field::new("raw_payload", DataType::Utf8, true), + ])); + let wide_batch = RecordBatch::try_new( + Arc::clone(&wide_schema), + vec![ + Arc::new(Int32Array::from_iter_values(0..NUM_ROWS as i32)), + Arc::new(BooleanArray::from_iter( + (0..NUM_ROWS).map(|i| Some(i % 3 == 0)), + )), + new_null_array(&DataType::Int64, NUM_ROWS), + new_null_array(&DataType::Utf8, NUM_ROWS), + Arc::new(events_list), + new_null_array(&DataType::Utf8, NUM_ROWS), + ], + ) + .unwrap(); + + let narrow_item_type = DataType::Struct(Fields::from(vec![ + Field::new("group_id", DataType::Int64, true), + Field::new("entity_id", DataType::Int64, true), + Field::new("metric_value", DataType::Float64, true), + ])); + let narrow_event_type = DataType::Struct(Fields::from(vec![ + Field::new("is_available", DataType::Boolean, true), + Field::new("event_time_ms", DataType::Int64, true), + Field::new( + "items", + DataType::List(Arc::new(Field::new("item", narrow_item_type, true))), + true, + ), + ])); + let narrow_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("is_flagged", DataType::Boolean, false), + Field::new( + "events", + DataType::List(Arc::new(Field::new("item", narrow_event_type, true))), + true, + ), + ])); + + let batches = assert_prunes( + vec![("wide.parquet", wide_batch)], + narrow_schema, + wide_schema, + "SELECT events FROM t_narrow ORDER BY id", + "SELECT events FROM t_full ORDER BY id", + ) + .await; + + let events = batches[0].column(0); + let events_list = events.as_any().downcast_ref::().unwrap(); + let event_structs = events_list + .values() + .as_any() + .downcast_ref::() + .unwrap(); + // event_token and latency_parts are dropped; only is_available, + // event_time_ms, and items survive. + assert_eq!(event_structs.fields().len(), 3); + + let items_list = event_structs + .column_by_name("items") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let item_structs = items_list + .values() + .as_any() + .downcast_ref::() + .unwrap(); + // feature_map, diagnostics, and pad are dropped; only group_id, + // entity_id, and metric_value survive, at the *inner* list + // nested two levels deep inside the outer one. + assert_eq!(item_structs.fields().len(), 3); + assert_eq!(item_structs.len(), NUM_ITEMS); + let group_id = item_structs + .column_by_name("group_id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + group_id.values(), + &(0..NUM_ITEMS as i64).collect::>() + ); + } +} diff --git a/datafusion/datasource-parquet/src/mod.rs b/datafusion/datasource-parquet/src/mod.rs index 25b79a618830c..35f831230b305 100644 --- a/datafusion/datasource-parquet/src/mod.rs +++ b/datafusion/datasource-parquet/src/mod.rs @@ -30,6 +30,7 @@ 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; 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..eed86f91fa952 --- /dev/null +++ b/datafusion/datasource-parquet/src/nested_schema_pruning.rs @@ -0,0 +1,511 @@ +// 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, for example 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, and the Arrow type the +//! reader will emit for them, by walking the physical and target type trees +//! in parallel and 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 every struct level clipped here keeps at least one leaf: a +//! struct cast with zero field-name overlap at *any* nesting depth is +//! rejected during physical planning +//! (`datafusion_common::nested_struct::validate_struct_compatibility`, called +//! recursively from `DefaultPhysicalExprAdapter::rewrite`), so a +//! [`CastColumnAccess`] observed here always has overlap at every struct +//! level it contains. +//! +//! 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. Nor are +//! `ListView`/`LargeListView`/`Dictionary` wrappers clipped here, even though +//! `cast_column` does recurse through them by name. That is a conservative +//! choice (safe, since the worst case is still just a full read) left as a +//! candidate follow-up rather than something this module currently handles. + +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, FieldRef, Fields}; + +/// The single child type one level of container nesting wraps, or `None` for +/// a type this module does not descend through (leaves, `Struct`, `Map`, and +/// wrapper kinds this module intentionally does not clip, see the module +/// doc). Shared by [`count_leaves`] and [`contains_struct`], which otherwise +/// need to agree on the exact same set of container variants. +fn nested_child(dt: &DataType) -> Option<&DataType> { + match dt { + DataType::List(f) + | DataType::LargeList(f) + | DataType::ListView(f) + | DataType::LargeListView(f) + | DataType::FixedSizeList(f, _) + | DataType::Map(f, _) => Some(f.data_type()), + DataType::Dictionary(_, value) => Some(value), + DataType::RunEndEncoded(_, value) => Some(value.data_type()), + _ => None, + } +} + +/// Clip `physical` against `cast_target`, returning the Parquet leaves the +/// cast actually consumes (as offsets relative to the root column's first +/// leaf, sorted ascending and non-empty) together with the Arrow type the +/// reader will emit for exactly those leaves. +/// +/// 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<(Vec, DataType)> { + let total = count_leaves(physical); + let mut kept = Vec::new(); + let mut next_leaf = 0; + let pruned_type = 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, pruned_type)) +} + +/// 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() + } + _ => nested_child(dt).map_or(1, count_leaves), + } +} + +/// Does this type contain a struct at any nesting depth? Used as a fast-path +/// gate: a root with no struct anywhere in its type has no leaves this +/// module could ever clip. +pub(crate) fn contains_struct(dt: &DataType) -> bool { + matches!(dt, DataType::Struct(_)) || nested_child(dt).is_some_and(contains_struct) +} + +/// Recursive walker: advances `next_leaf` across every leaf of `physical`, +/// pushing the offsets the cast target consumes into `kept`, and returns the +/// Arrow type the reader emits for those kept leaves. +fn clip_type( + physical: &DataType, + target: &DataType, + next_leaf: &mut usize, + kept: &mut Vec, +) -> DataType { + match (physical, target) { + (DataType::Struct(p_children), DataType::Struct(t_children)) => { + let kept_children: Fields = p_children + .iter() + .filter_map(|pc| { + let Some(tc) = t_children.iter().find(|tc| tc.name() == pc.name()) + else { + skip_leaves(pc.data_type(), next_leaf); + return None; + }; + let pruned = + clip_type(pc.data_type(), tc.data_type(), next_leaf, kept); + Some(field_with_type(pc, pruned)) + }) + .collect(); + DataType::Struct(kept_children) + } + (DataType::List(p_item), DataType::List(t_item)) => { + let pruned = + clip_type(p_item.data_type(), t_item.data_type(), next_leaf, kept); + DataType::List(field_with_type(p_item, pruned)) + } + (DataType::LargeList(p_item), DataType::LargeList(t_item)) => { + let pruned = + clip_type(p_item.data_type(), t_item.data_type(), next_leaf, kept); + DataType::LargeList(field_with_type(p_item, pruned)) + } + // Anything else, leaf pairs, wrapper-kind mismatches, maps, + // dictionaries, fixed-size lists, views, is kept wholesale. + _ => keep_all_leaves(physical, next_leaf, kept), + } +} + +/// Keep every leaf of `dt` (no pruning below this point); returns `dt` +/// unchanged since nothing was clipped. +fn keep_all_leaves( + dt: &DataType, + next_leaf: &mut usize, + kept: &mut Vec, +) -> DataType { + let n = count_leaves(dt); + kept.extend(*next_leaf..*next_leaf + n); + *next_leaf += n; + dt.clone() +} + +fn skip_leaves(dt: &DataType, next_leaf: &mut usize) { + *next_leaf += count_leaves(dt); +} + +/// 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::*; + + 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, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![1]); + assert_eq!(emitted, 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, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0, 2]); + assert_eq!(emitted, 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, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0]); + assert_eq!(emitted, 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, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0]); + assert_eq!( + emitted, + struct_of(vec![Field::new("x", DataType::Int32, true)]) + ); + } + + /// 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, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0]); + assert_eq!( + emitted, + 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, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0, 1]); + assert_eq!(emitted, list_of(struct_of(vec![int64("x"), utf8("y")]))); + } + + /// Two levels of `list` nesting, the inner one also narrowed, + /// the `events: array>>>` shape + /// reported in `datafusion-comet#4859`, where a sibling struct field at + /// the outer level (`aux`, standing in for that report's + /// `latency_parts`) is dropped entirely rather than clipped. + #[test] + fn clip_two_level_nested_list_of_struct() { + let physical = list_of(struct_of(vec![ + int64("a"), + utf8("pad"), + struct_of(vec![int64("x"), utf8("y")]).into_field("aux"), + list_of(struct_of(vec![int64("g"), utf8("pad2")])).into_field("items"), + ])); + let target = list_of(struct_of(vec![ + int64("a"), + list_of(struct_of(vec![int64("g")])).into_field("items"), + ])); + + let (kept, emitted) = clip_for_cast(&physical, &target).unwrap(); + // a=0, pad=1, aux.x=2, aux.y=3, items.g=4, items.pad2=5: only a and + // items.g survive; pad, all of aux, and items.pad2 are dropped. + assert_eq!(kept, vec![0, 4]); + assert_eq!( + emitted, + list_of(struct_of(vec![ + int64("a"), + list_of(struct_of(vec![int64("g")])).into_field("items"), + ])) + ); + } + + #[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, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0]); + assert_eq!(emitted, 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()); + } + + /// A struct level with zero field-name overlap can't actually reach this + /// code: `validate_struct_compatibility` rejects it during physical + /// planning (see the module doc), so `clip_for_cast` is only ever called + /// with targets that overlap at every nesting level. If it were reached + /// anyway, the generic catch-all keeps every leaf, still safe, just + /// unpruned. + #[test] + fn no_clip_on_zero_overlap() { + let physical = struct_of(vec![utf8("a"), int64("b")]); + let target = struct_of(vec![utf8("z")]); + assert!(clip_for_cast(&physical, &target).is_none()); + } + + /// 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, predicted_type) = 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(), &predicted_type); + + // 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/projection_read_plan.rs b/datafusion/datasource-parquet/src/projection_read_plan.rs index 96c99ab20750e..a73b3a58d7575 100644 --- a/datafusion/datasource-parquet/src/projection_read_plan.rs +++ b/datafusion/datasource-parquet/src/projection_read_plan.rs @@ -35,13 +35,18 @@ use parquet::arrow::ProjectionMask; use parquet::schema::types::SchemaDescriptor; use datafusion_common::Result; +use datafusion_common::nested_struct::requires_nested_struct_cast; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor}; use datafusion_functions::core::file_row_index::FileRowIndexFunc; use datafusion_functions::core::getfield::GetFieldFunc; -use datafusion_physical_expr::expressions::{Column, Literal}; +use datafusion_physical_expr::expressions::{CastExpr, Column, Literal}; use datafusion_physical_expr::utils::collect_columns; use datafusion_physical_expr::{PhysicalExpr, ScalarFunctionExpr}; +use crate::nested_schema_pruning::{ + CastColumnAccess, clip_for_cast, contains_struct, count_leaves, field_with_type, +}; + /// The result of resolving which Parquet leaf columns and Arrow schema fields /// are needed to evaluate an expression against a Parquet file /// @@ -94,6 +99,13 @@ pub(crate) 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. @@ -108,11 +120,19 @@ impl<'schema> PushdownChecker<'schema> { 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) -> Self { + self.collect_cast_accesses = true; + 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. /// @@ -217,6 +237,7 @@ impl<'schema> PushdownChecker<'schema> { PushdownColumns { required_columns: self.required_columns, struct_field_accesses: self.struct_field_accesses, + cast_accesses: self.cast_accesses, } } } @@ -308,6 +329,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()) { @@ -337,6 +380,9 @@ pub(crate) struct PushdownColumns { /// Struct field accesses via `get_field`. Each entry records the root struct /// column index and the field path being accessed. 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, } /// Builds a unified [`ParquetReadPlan`] for a set of projection expressions @@ -369,12 +415,13 @@ pub(crate) fn build_projection_read_plan( return root_level_plan(&root_indices, file_schema, schema_descr); } - // secondary fast path: if the schema has no struct columns, we can skip - // PushdownChecker traversal and use root-level projection + // 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| matches!(f.data_type(), DataType::Struct(_))); + .any(|f| contains_struct(f.data_type())); if !has_struct_columns { let mut root_indices = exprs @@ -390,19 +437,37 @@ pub(crate) fn build_projection_read_plan( 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); + let mut checker = PushdownChecker::new(file_schema, true).with_cast_collection(); 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 a cast + // access on the same root would be overridden anyway: drop those up + // front. `all_root_indices` is already sorted, so a binary search + // avoids building a second set just for this filter. + all_cast_accesses.retain(|c| all_root_indices.binary_search(&c.root_index).is_err()); + + 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() { @@ -419,6 +484,160 @@ pub(crate) fn build_projection_read_plan( read_plan } +/// 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 through a cast, and not also through a `get_field` +/// access on the same root, keep only the leaves the cast target names +/// (see `crate::nested_schema_pruning`); +/// - roots consumed only through `get_field` accesses keep the union of the +/// leaves those accesses reach, as before; +/// - any other referenced root, a cast that can't be safely clipped (see +/// `nested_schema_pruning::clip_for_cast`), or a root reached by both a +/// cast and a `get_field` access (not produced by +/// `DefaultPhysicalExprAdapter`, which always routes a `get_field` over a +/// narrowed column through the same cast rather than a separate access, +/// but a custom `PhysicalExprAdapter` could in principle inject both), +/// falls back to a full read of that root. +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(); + let struct_access_roots: BTreeSet = + struct_accesses.iter().map(|a| a.root_index).collect(); + // Every referenced root's Parquet leaves, grouped in one pass over the + // schema descriptor rather than one `leaf_indices_for_roots` scan per + // root (this function may look up several roots). + let leaves_by_root = leaves_grouped_by_root(schema_descr); + + // Root -> (absolute kept leaf indices, cast-clipped Arrow type) for + // roots successfully clipped via a cast. + let mut clipped_by_root: BTreeMap, DataType)> = BTreeMap::new(); + // Roots with a cast access that must fall back to a full read. + let mut fallback_roots: BTreeSet = BTreeSet::new(); + + for access in cast_accesses { + let root = access.root_index; + if whole_roots.contains(&root) + || fallback_roots.contains(&root) + || clipped_by_root.contains_key(&root) + { + continue; + } + if struct_access_roots.contains(&root) { + fallback_roots.insert(root); + continue; + } + + let physical_type = file_schema.field(root).data_type(); + let root_leaves = leaves_by_root.get(&root).map_or(&[][..], Vec::as_slice); + + // 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); + continue; + } + + match clip_for_cast(physical_type, &access.target_type) { + Some((kept_offsets, pruned_type)) => { + let start = root_leaves[0]; + let absolute = kept_offsets.into_iter().map(|o| start + o).collect(); + clipped_by_root.insert(root, (absolute, pruned_type)); + } + // Nothing prunable for this cast: every leaf is consumed. + None => { + fallback_roots.insert(root); + } + } + } + + // `get_field` accesses on roots not already handled by a cast clip (or + // by a whole-column/fallback full read) keep the existing (non-cast) + // leaf resolution. + let get_field_accesses: Vec = struct_accesses + .iter() + .filter(|a| { + !whole_roots.contains(&a.root_index) + && !fallback_roots.contains(&a.root_index) + && !clipped_by_root.contains_key(&a.root_index) + }) + .cloned() + .collect(); + + let mut leaf_indices: Vec = Vec::new(); + let mut fields: BTreeMap> = BTreeMap::new(); + + for root in whole_roots.iter().chain(fallback_roots.iter()) { + leaf_indices.extend(leaves_by_root[root].iter().copied()); + fields.insert(*root, Arc::new(file_schema.field(*root).clone())); + } + + for (&root, (kept, pruned_type)) in &clipped_by_root { + leaf_indices.extend(kept.iter().copied()); + fields.insert( + root, + field_with_type(file_schema.field(root), pruned_type.clone()), + ); + } + + if !get_field_accesses.is_empty() { + leaf_indices.extend(resolve_struct_field_leaves( + &get_field_accesses, + file_schema, + schema_descr, + )); + let get_field_schema = build_filter_schema(file_schema, &[], &get_field_accesses); + let get_field_roots: BTreeSet = + get_field_accesses.iter().map(|a| a.root_index).collect(); + for root in get_field_roots { + let field = get_field_schema + .field_with_name(file_schema.field(root).name()) + .expect("root name preserved by build_filter_schema"); + fields.insert(root, 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.into_values().collect::>(), + file_schema.metadata().clone(), + )), + } +} + +/// Groups every Parquet leaf index by its root (Arrow) column index, in one +/// pass over the schema descriptor. +fn leaves_grouped_by_root( + schema_descr: &SchemaDescriptor, +) -> BTreeMap> { + let mut by_root: BTreeMap> = BTreeMap::new(); + for leaf_idx in 0..schema_descr.num_columns() { + by_root + .entry(schema_descr.get_column_root_idx(leaf_idx)) + .or_default() + .push(leaf_idx); + } + by_root +} + /// Builds a leaf-level [`ParquetReadPlan`] covering `root_indices` in full plus /// the individual leaves reached by `struct_field_accesses`. /// @@ -677,6 +896,7 @@ mod test { use datafusion_physical_expr::planner::logical2physical; use parquet::arrow::ArrowWriter; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use parquet::file::metadata::ParquetMetaData; use tempfile::NamedTempFile; #[test] @@ -758,4 +978,134 @@ mod test { let expected_mask = ProjectionMask::leaves(schema_descr, [0, 1, 2]); assert_eq!(read_plan.projection_mask, expected_mask,); } + + /// Writes the id/struct fixture and returns the schema and metadata a + /// reader sees for it, so callers don't each repeat the reopen + + /// `ParquetRecordBatchReaderBuilder` boilerplate. + /// + /// Schema: id (Int32), s (Struct{value: Int32, label: Utf8, pad: Utf8}). + /// Parquet leaves: id=0, s.value=1, s.label=2, s.pad=3. + fn write_id_struct_file() -> (SchemaRef, Arc) { + let struct_fields: Fields = vec![ + Arc::new(Field::new("value", DataType::Int32, false)), + Arc::new(Field::new("label", DataType::Utf8, false)), + Arc::new(Field::new("pad", 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 _, + Arc::new(StringArray::from(vec!["p0", "p1", "p2"])) 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 builder = ParquetRecordBatchReaderBuilder::try_new(file.reopen().unwrap()) + .expect("reader builder"); + (builder.schema().clone(), builder.metadata().clone()) + } + + /// A projection consisting solely of a narrowing cast over a struct root + /// clips the read to the cast target's leaves. + #[test] + fn build_projection_read_plan_clips_cast_over_struct() { + let (file_schema, metadata) = write_id_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + let narrow = DataType::Struct( + vec![Arc::new(Field::new("value", DataType::Int32, true))].into(), + ); + let exprs: Vec> = vec![ + Arc::new(PhysicalColumn::new("id", 0)), + Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("s", 1)), + narrow.clone(), + None, + )), + ]; + + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + // Only id's leaf (0) and s.value's leaf (1) should be read: s.label + // and s.pad are clipped away. + let expected_mask = ProjectionMask::leaves(schema_descr, [0, 1]); + assert_eq!(read_plan.projection_mask, expected_mask); + + 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))].into() + ), + ); + } + + /// A root reached by both a narrowing cast and a `get_field` access (not + /// producible by `DefaultPhysicalExprAdapter`, but a custom + /// `PhysicalExprAdapter` could inject both) falls back to a full read of + /// that root rather than attempting to union the two leaf sets. + #[test] + fn build_projection_read_plan_falls_back_when_cast_and_get_field_share_a_root() { + let (file_schema, metadata) = write_id_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + let narrow = DataType::Struct( + vec![Arc::new(Field::new("value", DataType::Int32, true))].into(), + ); + let exprs: Vec> = vec![ + Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("s", 1)), + narrow, + None, + )), + logical2physical( + &get_field().call(vec![ + col("s"), + Expr::Literal(ScalarValue::Utf8(Some("label".to_string())), None), + ]), + &file_schema, + ), + ]; + + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + // Every leaf of `s` is read (full fallback), not just value/label. + let expected_mask = ProjectionMask::leaves(schema_descr, [1, 2, 3]); + assert_eq!(read_plan.projection_mask, expected_mask); + + 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)), + Arc::new(Field::new("pad", DataType::Utf8, false)), + ] + .into() + ), + ); + } } 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..89b40a652d720 --- /dev/null +++ b/datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt @@ -0,0 +1,110 @@ +# 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. +# The bytes-scanned assertions live in the Rust tests +# (datafusion/core/tests/parquet/expr_adapter.rs); this file covers the +# end-to-end SQL correctness path. +########## + +# The file contains events: ARRAY> and +# s: STRUCT; the table below declares 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 + +# `full` declares the file's physical schema exactly, so no cast is inserted +# and the scan always reads every leaf: a same-context baseline for the +# bytes_scanned comparison below. +statement ok +CREATE EXTERNAL TABLE full_schema ( + id INT, + events ARRAY>, + s STRUCT +) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/wide.parquet'; + +# bytes_scanned is a literal (not ) checked-in value: narrow +# reads fewer bytes than full_schema because the cast-clipped leaves drop +# pad_a, pad_b, and pad. A future change that widens the narrow read shows +# up here as a bytes_scanned mismatch. +query TT +explain analyze select events from narrow; +---- +Plan with Metrics DataSourceExec: metrics=[output_rows=3, bytes_scanned=172] + +query TT +explain analyze select events from full_schema; +---- +Plan with Metrics DataSourceExec: metrics=[output_rows=3, bytes_scanned=312] +