This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24090-ce2f153ebffabc659f79fbace6d9f42959470b40 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 5eba27f5bd4da21507cfaccaa94686929c3e1f7c Author: Matt Butrovich <[email protected]> AuthorDate: Wed Aug 5 16:19:34 2026 -0400 feat: prune unread Parquet leaves when a nested column is cast to a narrower type (#24090) ## Which issue does this PR close? - Related to https://github.com/apache/datafusion-comet/issues/4859. No DataFusion issue is filed for this one. ## Rationale for this change When a table declares a nested column narrower than the Parquet file's physical type, DataFusion reads every leaf of the column and drops the extra subfields in memory instead of skipping them at read time. - DefaultPhysicalExprAdapter rewrites the projected column into CAST(col AS narrow_type). - Projection mask derivation only understands literal get_field chains, so the cast's inner column expands to every physical leaf via ProjectionMask::roots. - Comet reported a production query reading 1.35 TB where plain Spark read 30.9 GB for the same pruned ReadSchema (datafusion-comet#4859). - Any embedder that hands DataFusion a pre-pruned schema (Comet, delta-rs, Iceberg integrations) hits the same gap. This is a port of #23398 onto current main. #23398 (stacked on the merged #23396 and #23397, superseding an earlier attempt at #23392) implements the fix and was reviewed favorably, but has merge conflicts against main since a follow-up refactor moved PushdownChecker and PushdownColumns into projection_read_plan.rs, and has four unanswered review comments. This PR reimplements the same approach against current main and resolves those four comments by construction: - Drops the nested_projection_pruning config flag; the clip's total fallback design makes a kill switch unnecessary. - Drops an unreachable zero-name-overlap code path; validate_struct_compatibility already rejects that case during physical planning. - Drops the untested union of a cast clip with a sibling get_field access on the same root, in favor of a full read fallback for that case. - Notes that ListView, LargeListView, and Dictionary wrappers are conservatively left unclipped, a candidate follow up. ## What changes are included in this PR? - New module datafusion/datasource-parquet/src/nested_schema_pruning.rs. clip_for_cast walks the physical and cast target type trees together, matches struct fields by name, recurses through List and LargeList, and returns the kept leaf offsets plus the pruned Arrow type in one pass. The clip is total: maps, dictionaries, wrapper-kind mismatches, and any shape it does not understand keep every leaf, so the worst case is today's full read. - PushdownChecker in projection_read_plan.rs now also collects CastColumnAccess entries for a CastExpr over a plain Column where the cast requires nested struct handling. Collection is opt in and only enabled for projection analysis, so filter pushdown is unchanged. - build_projection_read_plan routes to a new build_read_plan_with_cast_clipping when any cast access survives. It partitions referenced roots into whole column reads, cast clipped reads, get_field only reads, and full read fallbacks, including a fallback for a root reached by both a cast and a separate get_field access. - Fixed a pre-existing bug in the has_struct_columns fast path: it only tested the top level field type against Struct, so a LIST STRUCT root such as events was misclassified as containing no struct and skipped PushdownChecker entirely, meaning cast based clipping never fired for that shape. Replaced with a recursive contains_struct check through List, LargeList, ListView, LargeListView, FixedSizeList, Map, Dictionary, and RunEndEncoded. - No new configuration option. ## Are these changes tested? - 15 unit tests in nested_schema_pruning.rs: struct subset, reordering, leaf promotion, missing field null fill, nested struct in struct, list of struct, two levels of list of struct nesting, maps, dictionaries, wrapper mismatches, and an arrow-rs roundtrip test pinning that ProjectionMask::leaves over a subset of List Struct leaves emits exactly the predicted type. - 2 new unit tests in projection_read_plan.rs, plus the pre-existing struct preservation test, unchanged. - 8 integration tests in datafusion/core/tests/parquet/expr_adapter.rs, asserting identical results and a bytes_scanned drop of more than 2x: list of struct narrowing, top level struct, struct level nullability, get_field on a narrowed struct, mixed whole column and subfield access, filter pushdown enabled, a scan mixing a physically narrow and a wide file, and a regression test built from the ReadSchema and InputSchema shapes reported in datafusion-comet#4859 (a two level list STRUCT column with a dropped struct sibling, a dropped map sibling, and dropped top level columns). - New datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt covering the end to end SQL path. - The existing benchmark (already merged as part of #23397) had an assertion documenting the pre-fix baseline, narrow equals full. That assertion now fails as expected and has been flipped to assert narrow reads less than half of full. Measured for the top level struct case: narrow_schema=1081 bytes vs full_schema=8398289 bytes, matching the physically_narrow floor of 1081 bytes. ## Are there any user-facing changes? No API changes and no new configuration option. Behavior is IO reduction only, results are unchanged. --------- Co-authored-by: Adrian Garcia Badaracco <[email protected]> Co-authored-by: Claude Opus 5 <[email protected]> --- .../core/benches/parquet_nested_schema_pruning.rs | 36 +- datafusion/datasource-parquet/src/mod.rs | 1 + .../src/nested_schema_pruning.rs | 775 +++++++++++++++++++++ .../datasource-parquet/src/projection_read_plan.rs | 725 ++++++++++++++++++- .../test_files/parquet_nested_schema_pruning.slt | 565 +++++++++++++++ 5 files changed, 2069 insertions(+), 33 deletions(-) diff --git a/datafusion/core/benches/parquet_nested_schema_pruning.rs b/datafusion/core/benches/parquet_nested_schema_pruning.rs index 8db67de9ff..de4f0a57a5 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/datasource-parquet/src/mod.rs b/datafusion/datasource-parquet/src/mod.rs index 25b79a6188..35f831230b 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 0000000000..9768282c3b --- /dev/null +++ b/datafusion/datasource-parquet/src/nested_schema_pruning.rs @@ -0,0 +1,775 @@ +// 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<Struct<x, y, z, ...>>` but the expression is +//! `CAST(events AS List<Struct<x, y>>)`, 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. That holds for *any* +//! [`CastExpr`](datafusion_physical_expr::expressions::CastExpr) over a +//! nested type, not just the ones the schema adapter inserts: +//! `ColumnarValue::cast_to` routes every +//! cast for which +//! [`requires_nested_struct_cast`](datafusion_common::nested_struct::requires_nested_struct_cast) +//! holds, the same predicate the projection analysis gates on, through +//! `cast_column`. +//! +//! Struct-level nullability is preserved because the Parquet reader +//! reconstructs ancestor validity from the definition levels of any surviving +//! leaf, so every struct level that is clipped must keep at least one leaf. +//! A struct cast with zero field-name overlap at *any* nesting depth would +//! break that: the reader drops a field whose leaves are all masked out, so +//! the emitted type would not match the one predicted here. Such a cast is +//! rejected during physical planning +//! (`datafusion_common::nested_struct::validate_struct_compatibility`, called +//! recursively from `DefaultPhysicalExprAdapter::rewrite`) and by the logical +//! planner's own castability check, so it should never reach this module; if +//! one does anyway (a custom `PhysicalExprAdapter` could build one), +//! [`clip_for_cast`] detects the empty level and declines to clip. +//! +//! 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::collections::HashMap; +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<usize>, DataType)> { + let total = count_leaves(physical); + let mut kept = Vec::new(); + let mut next_leaf = 0; + let mut unclippable = false; + let pruned_type = clip_type( + physical, + cast_target, + &mut next_leaf, + &mut kept, + &mut unclippable, + ); + debug_assert_eq!(next_leaf, total, "leaf accounting must cover the type"); + if unclippable || 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) +} + +/// Above this many target fields, matching physical children against them one +/// by one turns into a quadratic string comparison; build a name lookup +/// instead. Below it the map's allocation costs more than the linear scan it +/// saves (Spark's `ParquetReadSupport.clipParquetGroupFields` builds the map +/// unconditionally; struct widths in practice are small enough that the +/// threshold is worth the branch). +const LINEAR_FIELD_SCAN_MAX: usize = 8; + +/// Find `name` among `fields`, using `by_name` when it was worth building. +/// Duplicate names resolve to the first occurrence either way. +fn lookup_field<'a>( + fields: &'a Fields, + by_name: &Option<HashMap<&'a str, &'a FieldRef>>, + name: &str, +) -> Option<&'a FieldRef> { + match by_name { + Some(map) => map.get(name).copied(), + None => fields.iter().find(|f| f.name() == name), + } +} + +/// 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. +/// +/// `unclippable` is set when a shape is encountered whose emitted type this +/// module cannot predict; the caller must then read the whole column. The walk +/// still runs to completion so `next_leaf` stays a valid leaf count. +fn clip_type( + physical: &DataType, + target: &DataType, + next_leaf: &mut usize, + kept: &mut Vec<usize>, + unclippable: &mut bool, +) -> DataType { + match (physical, target) { + (DataType::Struct(p_children), DataType::Struct(t_children)) => { + let t_by_name = (t_children.len() > LINEAR_FIELD_SCAN_MAX).then(|| { + let mut map = HashMap::with_capacity(t_children.len()); + for tc in t_children.iter() { + map.entry(tc.name().as_str()).or_insert(tc); + } + map + }); + let kept_children: Fields = p_children + .iter() + .filter_map(|pc| { + let Some(tc) = lookup_field(t_children, &t_by_name, pc.name()) else { + skip_leaves(pc.data_type(), next_leaf); + return None; + }; + let before = kept.len(); + let pruned = clip_type( + pc.data_type(), + tc.data_type(), + next_leaf, + kept, + unclippable, + ); + if kept.len() == before { + // This child matched by name but kept no leaves at + // all, which only happens when a nested struct level + // below it shares no field name with its target. The + // reader drops a field whose leaves are all masked + // out, so the emitted type could not be predicted; + // give up on clipping this column entirely rather + // than promise a type the decoder will not produce. + // (`DefaultPhysicalExprAdapter` never builds such a + // cast — `validate_struct_compatibility` rejects a + // zero-overlap struct level at planning time — but a + // custom `PhysicalExprAdapter` could.) + *unclippable = true; + } + 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, + unclippable, + ); + 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, + unclippable, + ); + 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<usize>, +) -> 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<Field>) -> 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); + // Wrapper kinds must be descended through, not counted as one leaf. + // A dictionary or run-end-encoded *value* that is itself a struct has + // as many leaves as the struct: counting it as 1 would misalign every + // later leaf index in the mask. + assert_eq!( + count_leaves(&DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(struct_of(vec![utf8("a"), int64("b")])) + )), + 2 + ); + assert_eq!( + count_leaves(&DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(Field::new( + "values", + struct_of(vec![utf8("a"), int64("b")]), + true + )) + )), + 2 + ); + } + + /// [`contains_struct`] gates the projection fast path, so it has to agree + /// with [`count_leaves`] about which wrappers are descended through. + #[test] + fn contains_struct_shapes() { + assert!(!contains_struct(&DataType::Int32)); + assert!(!contains_struct(&list_of(DataType::Int32))); + assert!(contains_struct(&struct_of(vec![int64("a")]))); + assert!(contains_struct(&list_of(struct_of(vec![int64("a")])))); + assert!(contains_struct(&DataType::LargeList(Arc::new(Field::new( + "item", + struct_of(vec![int64("a")]), + true + ))))); + assert!(contains_struct(&DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(struct_of(vec![int64("a")])) + ))); + assert!(!contains_struct(&DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Utf8) + ))); + // A map's entries are a struct, so a map always contains one. + assert!(contains_struct(&DataType::Map( + Arc::new(Field::new( + "entries", + struct_of(vec![utf8("key"), int64("value")]), + false + )), + false + ))); + } + + /// `{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<Struct>, 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<struct>` nesting, the inner one also narrowed, + /// the `events: array<struct<..., items: array<struct<...>>>>` 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()); + } + + /// A *nested* struct level with zero field-name overlap must not be + /// clipped, even when a sibling keeps leaves. The reader drops a field + /// whose leaves are all masked out (pinned by + /// [`reader_drops_struct_child_with_no_selected_leaves`]), so predicting + /// `{inner: Struct[], c}` here would be a schema the decoder never + /// produces. Read the whole column instead. + #[test] + fn no_clip_when_nested_struct_level_has_no_overlap() { + let physical = struct_of(vec![ + struct_of(vec![int64("a"), int64("b")]).into_field("inner"), + int64("c"), + ]); + let target = struct_of(vec![ + struct_of(vec![int64("z")]).into_field("inner"), + int64("c"), + ]); + assert!(clip_for_cast(&physical, &target).is_none()); + } + + /// Same, one level deeper and behind a list wrapper. + #[test] + fn no_clip_when_nested_list_struct_level_has_no_overlap() { + let physical = struct_of(vec![ + list_of(struct_of(vec![int64("a"), int64("b")])).into_field("items"), + int64("c"), + ]); + let target = struct_of(vec![ + list_of(struct_of(vec![int64("z")])).into_field("items"), + int64("c"), + ]); + assert!(clip_for_cast(&physical, &target).is_none()); + } + + /// Wide structs take the name-map matching path rather than the linear + /// scan; both must produce the same clip. + #[test] + fn clip_wide_struct_matches_by_name() { + let width = LINEAR_FIELD_SCAN_MAX * 4; + let physical = struct_of((0..width).map(|i| int64(&format!("f{i}"))).collect()); + // Even fields only, declared in reverse order: the emitted type is + // still in physical order. + let target = struct_of( + (0..width) + .rev() + .filter(|i| i % 2 == 0) + .map(|i| int64(&format!("f{i}"))) + .collect(), + ); + let (kept, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, (0..width).filter(|i| i % 2 == 0).collect::<Vec<_>>()); + assert_eq!( + emitted, + struct_of( + (0..width) + .filter(|i| i % 2 == 0) + .map(|i| int64(&format!("f{i}"))) + .collect() + ) + ); + } + + /// Duplicate physical field names both match the single target field and + /// are both kept, which is what the reader emits for that mask. + #[test] + fn clip_keeps_duplicate_physical_field_names() { + let physical = struct_of(vec![int64("a"), utf8("pad"), int64("a")]); + let target = struct_of(vec![int64("a")]); + let (kept, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0, 2]); + assert_eq!(emitted, struct_of(vec![int64("a"), int64("a")])); + } + + /// Pins the arrow-rs behavior the empty-level guard above depends on: a + /// struct child none of whose leaves are selected disappears from the + /// type the reader emits, rather than surviving as an empty struct. + #[test] + fn reader_drops_struct_child_with_no_selected_leaves() { + use arrow::array::{ArrayRef, Int64Array, StructArray}; + use arrow::record_batch::RecordBatch; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use parquet::arrow::{ArrowWriter, ProjectionMask}; + + let inner_fields = Fields::from(vec![int64("a"), int64("b")]); + let outer_fields = Fields::from(vec![ + Field::new("inner", DataType::Struct(inner_fields.clone()), true), + int64("c"), + ]); + let inner: ArrayRef = Arc::new(StructArray::new( + inner_fields, + vec![ + Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef, + Arc::new(Int64Array::from(vec![3, 4])) as ArrayRef, + ], + None, + )); + let outer = StructArray::new( + outer_fields.clone(), + vec![inner, Arc::new(Int64Array::from(vec![5, 6])) as ArrayRef], + None, + ); + let schema = Arc::new(arrow::datatypes::Schema::new(vec![Field::new( + "s", + DataType::Struct(outer_fields), + true, + )])); + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(outer)]).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(); + + let builder = + ParquetRecordBatchReaderBuilder::try_new(file.reopen().unwrap()).unwrap(); + assert_eq!(builder.parquet_schema().num_columns(), 3); + // Keep only s.c (leaf 2): every leaf of s.inner is masked out. + let mask = ProjectionMask::leaves(builder.parquet_schema(), [2usize]); + let reader = builder.with_projection(mask).build().unwrap(); + let out: Vec<RecordBatch> = reader.map(|b| b.unwrap()).collect(); + assert_eq!( + out[0].schema().field(0).data_type(), + &struct_of(vec![int64("c")]), + "the fully masked `inner` child is dropped, not emitted as an empty struct" + ); + } + + /// Pins the arrow-rs behavior this module relies on: selecting a subset + /// of leaves under a `List<Struct>` 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<ArrayRef> = 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<RecordBatch> = 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::<ListArray>().unwrap(); + assert!(events.is_valid(0)); + assert!(events.is_null(1)); + assert!(events.is_valid(2)); + let structs = events + .values() + .as_any() + .downcast_ref::<StructArray>() + .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::<Int64Array>() + .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 96c99ab207..350ed9596b 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<usize>, /// Struct field accesses via `get_field`. struct_field_accesses: Vec<StructFieldAccess>, + /// 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<CastColumnAccess>, + /// 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<Struct<subset of fields>>)` 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::<CastExpr>() + && let Some(column) = cast.expr().downcast_ref::<Column>() + && 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::<Column>() && 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<StructFieldAccess>, + /// Whole-column casts to a narrower nested type. Empty unless cast + /// collection was enabled on the checker. + pub(crate) cast_accesses: Vec<CastColumnAccess>, } /// Builds a unified [`ParquetReadPlan`] for a set of projection expressions @@ -369,19 +415,29 @@ 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 - let has_struct_columns = file_schema - .fields() - .iter() - .any(|f| matches!(f.data_type(), DataType::Struct(_))); + // secondary fast path: if none of the *projected* columns contains a + // struct at any nesting level, there are no leaves to prune and we can + // skip the PushdownChecker traversal and use root-level projection. + // + // Gating on the projected roots rather than on every field of the file + // schema keeps this step O(projected columns): a wide file with a nested + // column the projection never touches should not push the whole + // projection through the slower, name-resolving path. Any column whose + // `index` does not line up with the file schema (a stale `Column` from an + // earlier rewrite) falls through to that path, which resolves by name. + let projected_columns = exprs.iter().flat_map(collect_columns).collect::<Vec<_>>(); + let all_resolvable_and_struct_free = projected_columns.iter().all(|col| { + file_schema + .fields() + .get(col.index()) + .is_some_and(|f| f.name() == col.name() && !contains_struct(f.data_type())) + }); - if !has_struct_columns { - let mut root_indices = exprs - .into_iter() - .flat_map(|e| collect_columns(&e).into_iter().map(|col| col.index())) + if all_resolvable_and_struct_free { + let mut root_indices = projected_columns + .iter() + .map(|c| c.index()) .collect::<Vec<_>>(); - root_indices.sort_unstable(); root_indices.dedup(); @@ -390,19 +446,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 +493,192 @@ 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`), a root reached by two casts +/// with *different* targets (a projection can consume the same column +/// through more than one narrowing cast, e.g. +/// `SELECT CAST(s AS STRUCT(a)), CAST(s AS STRUCT(b)) FROM t`; clipping to +/// either target alone would starve the other), 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<usize> = whole_root_indices.iter().copied().collect(); + let struct_access_roots: BTreeSet<usize> = + 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<usize, (Vec<usize>, DataType)> = BTreeMap::new(); + // Roots with a cast access that must fall back to a full read. + let mut fallback_roots: BTreeSet<usize> = BTreeSet::new(); + // The cast target already clipped for a root, so a second cast on the + // same root can be recognised as either a repeat (same target: nothing to + // do) or a conflict (different target: neither clip is valid on its own). + let mut clipped_target_by_root: BTreeMap<usize, &DataType> = BTreeMap::new(); + + for access in cast_accesses { + let root = access.root_index; + if whole_roots.contains(&root) || fallback_roots.contains(&root) { + continue; + } + if let Some(previous) = clipped_target_by_root.get(&root) { + if **previous != access.target_type { + // The projection consumes this root through two different + // narrowing casts. Each cast only needs its own leaves, but + // the mask is per column: clipping to the first target would + // silently null-fill whatever the second one needs. Read the + // whole root instead. + clipped_by_root.remove(&root); + clipped_target_by_root.remove(&root); + fallback_roots.insert(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)); + clipped_target_by_root.insert(root, &access.target_type); + } + // Nothing prunable for this cast: every leaf is consumed. + None => { + fallback_roots.insert(root); + } + } + } + + // `get_field` accesses on roots not already read in full (as a whole + // column, or as a cast that fell back) keep the existing (non-cast) leaf + // resolution. + let get_field_accesses: Vec<StructFieldAccess> = struct_accesses + .iter() + .filter(|a| { + // A root carrying a `get_field` access is put into + // `fallback_roots` before any clip is attempted (see the loop + // above), so it can never also be clipped. Assert that rather + // than re-testing it here, so a future reordering trips the + // assert instead of silently changing which leaves are read. + debug_assert!(!clipped_by_root.contains_key(&a.root_index)); + !whole_roots.contains(&a.root_index) + && !fallback_roots.contains(&a.root_index) + }) + .cloned() + .collect(); + + let mut leaf_indices: Vec<usize> = Vec::new(); + let mut fields: BTreeMap<usize, Arc<Field>> = BTreeMap::new(); + + for root in whole_roots.iter().chain(fallback_roots.iter()) { + // A root with no parquet leaves contributes nothing to the mask; + // `ProjectionMask::roots` handles that case the same way, so match it + // rather than indexing and panicking. + if let Some(leaves) = leaves_by_root.get(root) { + leaf_indices.extend(leaves.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<usize> = + get_field_accesses.iter().map(|a| a.root_index).collect(); + // `build_filter_schema` emits one field per accessed root in + // ascending root order, which is the order `get_field_roots` iterates + // in, so the two line up positionally. Pairing them beats looking each + // one up by name: no repeated linear scans, and no ambiguity if two + // roots happen to share a name. + debug_assert_eq!(get_field_roots.len(), get_field_schema.fields().len()); + for (root, field) in get_field_roots.iter().zip(get_field_schema.fields()) { + fields.insert(*root, Arc::clone(field)); + } + } + + 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::<Vec<_>>(), + 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<usize, Vec<usize>> { + let mut by_root: BTreeMap<usize, Vec<usize>> = 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 +937,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 +1019,442 @@ 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<ParquetMetaData>) { + 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()) + } + + /// Writes a two-struct-root fixture so tests can combine a cast on one + /// root with an access on another. + /// + /// Schema: a (Struct{p: Int32, q: Utf8}), b (Struct{m: Int32, n: Utf8}). + /// Parquet leaves: a.p=0, a.q=1, b.m=2, b.n=3. + fn write_two_struct_file() -> (SchemaRef, Arc<ParquetMetaData>) { + let group = |first: &str, second: &str| -> Fields { + vec![ + Arc::new(Field::new(first, DataType::Int32, false)), + Arc::new(Field::new(second, DataType::Utf8, false)), + ] + .into() + }; + let (a_fields, b_fields) = (group("p", "q"), group("m", "n")); + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Struct(a_fields.clone()), false), + Field::new("b", DataType::Struct(b_fields.clone()), false), + ])); + + let values = |fields: Fields, ints: [i32; 2], strs: [&str; 2]| { + Arc::new(StructArray::new( + fields, + vec![ + Arc::new(Int32Array::from(ints.to_vec())) as _, + Arc::new(StringArray::from(strs.to_vec())) as _, + ], + None, + )) as _ + }; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + values(a_fields, [1, 2], ["a0", "a1"]), + values(b_fields, [3, 4], ["b0", "b1"]), + ], + ) + .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()) + } + + /// Builds `CAST(Column(name, index) AS Struct{fields})`. + fn cast_to_struct( + name: &str, + index: usize, + fields: Vec<(&str, DataType)>, + ) -> Arc<dyn PhysicalExpr> { + let target = DataType::Struct( + fields + .into_iter() + .map(|(n, dt)| Arc::new(Field::new(n, dt, true))) + .collect::<Vec<_>>() + .into(), + ); + Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new(name, index)), + target, + None, + )) + } + + /// Builds `get_field(Column(name, index), field)`. + fn get_field_of( + file_schema: &Schema, + name: &str, + field: &str, + ) -> Arc<dyn PhysicalExpr> { + logical2physical( + &get_field().call(vec![ + col(name), + Expr::Literal(ScalarValue::Utf8(Some(field.to_string())), None), + ]), + file_schema, + ) + } + + /// Clipping a cast whose only surviving field is *not* the struct's first + /// one: the kept offsets are relative to the root's first leaf and must be + /// rebased onto it. With `s` starting at leaf 1 and `label` at offset 1, + /// getting the arithmetic wrong reads `id` (leaf 0) instead of `s.label`. + #[test] + fn build_projection_read_plan_clips_cast_to_a_non_leading_field() { + let (file_schema, metadata) = write_id_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + let exprs = vec![cast_to_struct("s", 1, vec![("label", DataType::Utf8)])]; + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + assert_eq!( + read_plan.projection_mask, + ProjectionMask::leaves(schema_descr, [2]) + ); + 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("label", DataType::Utf8, false))].into() + ), + ); + } + + /// A cast on one root and a `get_field` on a *different* root: each root + /// keeps only what it needs, and both appear in the projected schema in + /// root order. + #[test] + fn build_projection_read_plan_clips_cast_beside_get_field_on_another_root() { + let (file_schema, metadata) = write_two_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + let exprs = vec![ + cast_to_struct("a", 0, vec![("p", DataType::Int32)]), + get_field_of(&file_schema, "b", "n"), + ]; + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + // a.p (leaf 0) from the clip, b.n (leaf 3) from the field access. + assert_eq!( + read_plan.projection_mask, + ProjectionMask::leaves(schema_descr, [0, 3]) + ); + let field_types = read_plan + .projected_schema + .fields() + .iter() + .map(|f| (f.name().clone(), f.data_type().clone())) + .collect::<Vec<_>>(); + assert_eq!( + field_types, + vec![ + ( + "a".to_string(), + DataType::Struct( + vec![Arc::new(Field::new("p", DataType::Int32, false))].into() + ) + ), + ( + "b".to_string(), + DataType::Struct( + vec![Arc::new(Field::new("n", DataType::Utf8, false))].into() + ) + ), + ] + ); + } + + /// Once conflicting cast targets have demoted a root to a full read, a + /// *third* cast on it must not resurrect the clip. + #[test] + fn build_projection_read_plan_keeps_full_read_after_a_third_cast() { + let (file_schema, metadata) = write_id_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + let exprs = vec![ + cast_to_struct("s", 1, vec![("value", DataType::Int32)]), + cast_to_struct("s", 1, vec![("label", DataType::Utf8)]), + cast_to_struct("s", 1, vec![("value", DataType::Int32)]), + ]; + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + assert_eq!( + read_plan.projection_mask, + ProjectionMask::leaves(schema_descr, [1, 2, 3]) + ); + let s_field = read_plan.projected_schema.field_with_name("s").unwrap(); + assert_eq!(s_field.data_type(), file_schema.field(1).data_type()); + } + + /// A whole-column reference wins over a `get_field` access on the same + /// root even when another root is being clipped: `a` keeps every leaf and + /// its full type, `b` keeps only the cast target's. + #[test] + fn build_projection_read_plan_whole_column_beats_get_field_beside_a_clip() { + let (file_schema, metadata) = write_two_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + let exprs: Vec<Arc<dyn PhysicalExpr>> = vec![ + Arc::new(PhysicalColumn::new("a", 0)), + get_field_of(&file_schema, "a", "p"), + cast_to_struct("b", 1, vec![("m", DataType::Int32)]), + ]; + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + // Every leaf of `a` (0, 1) plus b.m (leaf 2). + assert_eq!( + read_plan.projection_mask, + ProjectionMask::leaves(schema_descr, [0, 1, 2]) + ); + let a_field = read_plan.projected_schema.field_with_name("a").unwrap(); + assert_eq!( + a_field.data_type(), + file_schema.field(0).data_type(), + "the whole-column reference must keep `a`'s full type" + ); + } + + /// Columns are resolved by *name*: a `Column` whose index points at a + /// different field (a stale index left by an earlier rewrite) must not be + /// taken at face value by the struct fast-path gate. + #[test] + fn build_projection_read_plan_resolves_stale_column_indices_by_name() { + let (file_schema, metadata) = write_id_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + // `s` is at index 1; this claims index 0, which is `id`. + let exprs = vec![cast_to_struct("s", 0, vec![("value", DataType::Int32)])]; + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + assert_eq!( + read_plan.projection_mask, + ProjectionMask::leaves(schema_descr, [1]), + "the cast must resolve to `s`, not to whatever sits at index 0" + ); + } + + /// 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<Arc<dyn PhysicalExpr>> = 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() + ), + ); + } + + /// Two casts on the same root with the *same* target still clip: this is + /// the shape the expression adapter produces when one column is + /// referenced several times (`SELECT s, s FROM narrowed`). + #[test] + fn build_projection_read_plan_clips_repeated_identical_casts() { + 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 cast = || -> Arc<dyn PhysicalExpr> { + Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("s", 1)), + narrow.clone(), + None, + )) + }; + + let read_plan = + build_projection_read_plan(vec![cast(), cast()], &file_schema, schema_descr); + + assert_eq!( + read_plan.projection_mask, + ProjectionMask::leaves(schema_descr, [1]) + ); + } + + /// Two casts on the same root with *different* targets cannot both be + /// served by one mask: clipping to either target alone would null-fill + /// whatever the other one needs (or fail its runtime struct-compatibility + /// check outright). Read the whole root instead. + #[test] + fn build_projection_read_plan_falls_back_on_conflicting_cast_targets() { + let (file_schema, metadata) = write_id_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + let narrow = |name: &str, dt: DataType| -> Arc<dyn PhysicalExpr> { + Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("s", 1)), + DataType::Struct(vec![Arc::new(Field::new(name, dt, true))].into()), + None, + )) + }; + let exprs = vec![ + narrow("value", DataType::Int32), + narrow("label", DataType::Utf8), + ]; + + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + assert_eq!( + read_plan.projection_mask, + ProjectionMask::leaves(schema_descr, [1, 2, 3]), + "every leaf of `s` must be read so both casts see their fields" + ); + let s_field = read_plan.projected_schema.field_with_name("s").unwrap(); + assert_eq!(s_field.data_type(), file_schema.field(1).data_type()); + } + + /// The struct fast-path gate looks at the *projected* columns, not at + /// every field of the file schema: projecting only `id` produces the same + /// root-level plan it would for a schema with no struct in it at all. + #[test] + fn build_projection_read_plan_ignores_unprojected_struct_columns() { + let (file_schema, metadata) = write_id_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + // Not a bare column, so the all-plain-columns fast path does not apply. + let exprs: Vec<Arc<dyn PhysicalExpr>> = vec![Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("id", 0)), + DataType::Int64, + None, + ))]; + + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + assert_eq!( + read_plan.projection_mask, + ProjectionMask::roots(schema_descr, [0]) + ); + assert_eq!(read_plan.projected_schema.fields().len(), 1); + } + + /// 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<Arc<dyn PhysicalExpr>> = 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 0000000000..d936a89beb --- /dev/null +++ b/datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt @@ -0,0 +1,565 @@ +# 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. +# +# This file covers both halves of that claim: the results are correct, and +# the scan really did read less. Each `explain analyze` below pins a literal +# bytes_scanned against a same-context baseline table declaring the file's +# own physical schema, so no cast is inserted and every leaf is read. A +# change that silently widens a clipped read shows up as a mismatch here. +########## + +# The file contains events: ARRAY<STRUCT<x, y, pad_a, pad_b>> and +# s: STRUCT<x, y, pad>; 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<STRUCT<x BIGINT, y VARCHAR, z BIGINT>>, + s STRUCT<x BIGINT, y VARCHAR> +) +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_schema` names every field the file has, so nothing can be clipped away +# and the scan always reads every leaf: a same-context baseline for the +# bytes_scanned comparison below. (A cast is still inserted — the declared +# leaf types differ from the file's, e.g. VARCHAR maps to Utf8View here while +# the file holds Utf8 — but it is not a *narrowing* one, so `clip_for_cast` +# keeps all the leaves.) +statement ok +CREATE EXTERNAL TABLE full_schema ( + id INT, + events ARRAY<STRUCT<x INT, y VARCHAR, pad_a VARCHAR, pad_b VARCHAR>>, + s STRUCT<x INT, y VARCHAR, pad VARCHAR> +) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/wide.parquet'; + +# bytes_scanned is a literal (not <slt:ignore>) 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: <slt:ignore>metrics=[output_rows=3, <slt:ignore>bytes_scanned=172<slt:ignore>] + +query TT +explain analyze select events from full_schema; +---- +Plan with Metrics DataSourceExec: <slt:ignore>metrics=[output_rows=3, <slt:ignore>bytes_scanned=312<slt:ignore>] + +# Same for the top-level struct column: the clipped read drops `pad`. +query TT +explain analyze select s from narrow; +---- +Plan with Metrics DataSourceExec: <slt:ignore>metrics=[output_rows=3, <slt:ignore>bytes_scanned=146<slt:ignore>] + +query TT +explain analyze select s from full_schema; +---- +Plan with Metrics DataSourceExec: <slt:ignore>metrics=[output_rows=3, <slt:ignore>bytes_scanned=219<slt:ignore>] + +# `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 "nothing was clipped" +# baseline is therefore reading every physical leaf of `s` +# (`select s from full_schema` above), not the same `get_field` query against +# `full_schema` -- that one needs no cast at all and takes `get_field`'s own, +# more precise, single-leaf pushdown path. +query TT +explain analyze select s['x'] from narrow; +---- +Plan with Metrics DataSourceExec: <slt:ignore>metrics=[output_rows=3, <slt:ignore>bytes_scanned=146<slt:ignore>] + +# Mixed access -- the whole (narrowed) column and a subfield of it -- still +# reads only the narrow schema's leaves. +query TT +explain analyze select s, s['y'] from narrow; +---- +Plan with Metrics DataSourceExec: <slt:ignore>metrics=[output_rows=3, <slt:ignore>bytes_scanned=146<slt:ignore>] + +query TT +explain analyze select s, s['y'] from full_schema; +---- +Plan with Metrics DataSourceExec: <slt:ignore>metrics=[output_rows=3, <slt:ignore>bytes_scanned=219<slt:ignore>] + + +# `SELECT *` goes through the same clipped read as an explicit projection. +query I?? +SELECT * 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 + +# Referencing the narrowed column as a whole *and* through a field access in +# the same query. +query I?T +SELECT id, s, s['y'] FROM narrow ORDER BY id; +---- +1 {x: 100, y: s1} s1 +2 {x: 200, y: s2} s2 +3 NULL NULL + +# Aggregating over a clipped nested column. +query IIT +SELECT count(*), sum(s['x']), string_agg(s['y'], ',' ORDER BY id) FROM narrow; +---- +3 300 s1,s2 + +# Filtering on a field of a clipped nested column. +query I? +SELECT id, s FROM narrow WHERE s['x'] = 200 ORDER BY id; +---- +2 {x: 200, y: s2} + +# A declared schema whose fields are in a different order from the file's: +# the values follow the declared order, not the physical one. +statement ok +CREATE EXTERNAL TABLE reordered ( + id INT, + events ARRAY<STRUCT<y VARCHAR, x BIGINT>>, + s STRUCT<y VARCHAR, x BIGINT> +) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/wide.parquet'; + +query I?? +SELECT id, events, s FROM reordered ORDER BY id; +---- +1 [{y: a1, x: 10}] {y: s1, x: 100} +2 [{y: b1, x: 20}, {y: b2, x: 21}] {y: s2, x: 200} +3 NULL NULL + +statement ok +DROP TABLE reordered; + +# A declared struct sharing no field name with the file's is rejected rather +# than silently null-filled: `clip_for_cast` never sees a zero-overlap cast. +statement ok +CREATE EXTERNAL TABLE no_overlap ( + id INT, + s STRUCT<not_in_file BIGINT> +) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/wide.parquet'; + +statement error DataFusion error: Execution error: Cannot cast column 's' +SELECT s FROM no_overlap; + +statement ok +DROP TABLE no_overlap; + +########## +# Struct nested inside a struct: both levels are clipped, and the reader's +# reconstruction of struct validity survives at both levels. +########## + +statement ok +COPY ( + SELECT id, n + FROM (VALUES + (1, named_struct('inner', named_struct('a', 1, 'pad_i', 'pi1'), 'c', 'c1', 'pad_o', 'po1')), + (2, named_struct('inner', named_struct('a', 2, 'pad_i', 'pi2'), 'c', 'c2', 'pad_o', 'po2')), + (3, NULL) + ) AS t(id, n) +) TO 'test_files/scratch/parquet_nested_schema_pruning/nested_struct.parquet' +STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE nested_narrow ( + id INT, + n STRUCT<inner STRUCT<a BIGINT>, c VARCHAR> +) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/nested_struct.parquet'; + +query I? +SELECT id, n FROM nested_narrow ORDER BY id; +---- +1 {inner: {a: 1}, c: c1} +2 {inner: {a: 2}, c: c2} +3 NULL + +query IBB +SELECT id, n IS NULL, n['inner'] IS NULL FROM nested_narrow ORDER BY id; +---- +1 false false +2 false false +3 true true + +statement ok +DROP TABLE nested_narrow; + +########## +# The exact shape reported in datafusion-comet#4859: a two-level +# `ARRAY<STRUCT<..., items: ARRAY<STRUCT<...>>>>` 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 as the issue (field names +# 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. +# +# Every dropped sibling carries real data rather than NULLs, so the +# bytes_scanned gap below is attributable to the clip and not to NULL columns +# being cheap. +########## + +statement ok +COPY ( + SELECT id, is_flagged, dimension_id, region_code, events, raw_payload + FROM (VALUES + (1, true, 1001, 'us-east', [named_struct( + 'is_available', true, + 'event_time_ms', 10, + 'event_token', 'token-0', + 'latency_parts', named_struct('queue_time_ms', 5, 'retry_count', 1), + 'items', [named_struct('group_id', 1, 'entity_id', 101, 'metric_value', 1.5, + 'feature_map', MAP {'f1': 0.25}, + 'diagnostics', named_struct('module_id', 'm1', 'trace_id', 't1'), + 'pad', 'pad-0000'), + named_struct('group_id', 2, 'entity_id', 102, 'metric_value', 3.0, + 'feature_map', MAP {'f2': 0.5}, + 'diagnostics', named_struct('module_id', 'm2', 'trace_id', 't2'), + 'pad', 'pad-0001')])], + 'payload-0'), + (2, false, 1002, 'us-west', [named_struct( + 'is_available', false, + 'event_time_ms', 20, + 'event_token', 'token-1', + 'latency_parts', named_struct('queue_time_ms', 7, 'retry_count', 2), + 'items', [named_struct('group_id', 3, 'entity_id', 103, 'metric_value', 4.5, + 'feature_map', MAP {'f3': 0.75}, + 'diagnostics', named_struct('module_id', 'm3', 'trace_id', 't3'), + 'pad', 'pad-0002')])], + 'payload-1') + ) AS t(id, is_flagged, dimension_id, region_code, events, raw_payload) +) TO 'test_files/scratch/parquet_nested_schema_pruning/two_level.parquet' +STORED AS PARQUET; + +# Declares neither the dropped top-level columns nor, inside `events`, +# `event_token`/`latency_parts`, nor, inside `items`, the map, the nested +# struct, or the pad. +statement ok +CREATE EXTERNAL TABLE two_level_narrow ( + id INT, + is_flagged BOOLEAN, + events ARRAY<STRUCT< + is_available BOOLEAN, + event_time_ms BIGINT, + items ARRAY<STRUCT<group_id BIGINT, entity_id BIGINT, metric_value DOUBLE>> + >> +) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/two_level.parquet'; + +# The file's own physical schema, so no cast is inserted: the same-context +# baseline for the bytes comparison. +statement ok +CREATE EXTERNAL TABLE two_level_full +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/two_level.parquet'; + +# Only the declared subfields survive, at *both* nesting levels: the printed +# structs are the emitted Arrow type. +query I? +SELECT id, events FROM two_level_narrow ORDER BY id; +---- +1 [{is_available: true, event_time_ms: 10, items: [{group_id: 1, entity_id: 101, metric_value: 1.5}, {group_id: 2, entity_id: 102, metric_value: 3.0}]}] +2 [{is_available: false, event_time_ms: 20, items: [{group_id: 3, entity_id: 103, metric_value: 4.5}]}] + +# Unnesting twice reaches the inner list's surviving leaves. +query III +SELECT id, i['group_id'], i['entity_id'] +FROM (SELECT id, unnest(e['items']) AS i + FROM (SELECT id, unnest(events) AS e FROM two_level_narrow)) +ORDER BY id, i['group_id']; +---- +1 1 101 +1 2 102 +2 3 103 + +query TT +explain analyze select events from two_level_narrow; +---- +Plan with Metrics DataSourceExec: <slt:ignore>metrics=[output_rows=2, <slt:ignore>bytes_scanned=381<slt:ignore>] + +query TT +explain analyze select events from two_level_full; +---- +Plan with Metrics DataSourceExec: <slt:ignore>metrics=[output_rows=2, <slt:ignore>bytes_scanned=1.05 K<slt:ignore>] + +statement ok +DROP TABLE two_level_narrow; + +statement ok +DROP TABLE two_level_full; + +########## +# A MAP column is never clipped (the runtime cast routes maps through Arrow's +# positional struct cast, which needs every child), but it must not stop a +# struct sibling from being clipped. The declared schema below omits the map +# entirely, leaving it in the file as an unprojected root. +########## + +statement ok +COPY ( + SELECT id, m, s + FROM (VALUES + (1, MAP {'k1': 1, 'k2': 2}, named_struct('x', 10, 'pad', 'p1')), + (2, MAP {'k1': 3}, named_struct('x', 20, 'pad', 'p2')) + ) AS t(id, m, s) +) TO 'test_files/scratch/parquet_nested_schema_pruning/with_map.parquet' +STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE map_sibling ( + id INT, + s STRUCT<x BIGINT> +) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/with_map.parquet'; + +query I? +SELECT id, s FROM map_sibling ORDER BY id; +---- +1 {x: 10} +2 {x: 20} + +statement ok +DROP TABLE map_sibling; + +########## +# One table over two files, one physically narrow (no cast inserted) and one +# wide (clipped). Both must read correctly in the same scan. +########## + +statement ok +COPY ( + SELECT id, s + FROM (VALUES + (10, named_struct('x', 1000, 'y', 'w1', 'pad', 'wp1')) + ) AS t(id, s) +) TO 'test_files/scratch/parquet_nested_schema_pruning/mixed/wide.parquet' +STORED AS PARQUET; + +statement ok +COPY ( + SELECT id, s + FROM (VALUES + (20, named_struct('x', 2000, 'y', 'n1')) + ) AS t(id, s) +) TO 'test_files/scratch/parquet_nested_schema_pruning/mixed/narrow.parquet' +STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE mixed_files ( + id INT, + s STRUCT<x BIGINT, y VARCHAR> +) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/mixed/'; + +query I? +SELECT id, s FROM mixed_files ORDER BY id; +---- +10 {x: 1000, y: w1} +20 {x: 2000, y: n1} + +statement ok +DROP TABLE mixed_files; + +########## +# A predicate on a primitive column with filter pushdown enabled, while the +# projected nested column is clipped: the clip and the row filter have to +# coexist on the same scan. +# +# The predicate is deliberately on `id` and not on a field of the clipped +# column: `WHERE s['x'] = ...` with pushdown enabled is silently dropped +# today (apache/datafusion#24109), which is a pre-existing row-filter bug +# rather than anything this feature does. +########## + +statement ok +set datafusion.execution.parquet.pushdown_filters = true; + +query I? +SELECT id, s FROM narrow WHERE id >= 2 ORDER BY id; +---- +2 {x: 200, y: s2} +3 NULL + +query TT +explain analyze select s from narrow where id >= 2; +---- +Plan with Metrics DataSourceExec: <slt:ignore>metrics=[output_rows=2, <slt:ignore>bytes_scanned=219<slt:ignore>] + +query TT +explain analyze select s from full_schema where id >= 2; +---- +Plan with Metrics DataSourceExec: <slt:ignore>metrics=[output_rows=2, <slt:ignore>bytes_scanned=292<slt:ignore>] + +statement ok +set datafusion.execution.parquet.pushdown_filters = false; + +########## +# Query-level casts. `ProjectionExec` is merged into the scan, so a `CAST` +# written in the query reaches the same read-plan analysis as an +# adapter-inserted one — including one column consumed through two *different* +# cast targets, which no single clipped read can serve. Clipping to one +# target's leaves would leave the other cast reading a struct that is missing +# the fields it names, which `cast_column` either null-fills (wrong results) +# or, for disjoint targets, rejects outright. +# +# `exact` infers its schema from the file, so no adapter cast is interposed +# and the casts below are the only ones the scan sees. +########## + +statement ok +CREATE EXTERNAL TABLE exact +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/wide.parquet'; + +# A single query-level cast is clipped like an adapter-inserted one. +query ? +SELECT CAST(s AS STRUCT<y VARCHAR>) FROM exact ORDER BY id; +---- +{y: s1} +{y: s2} +NULL + +# Disjoint targets. +query ?? +SELECT CAST(s AS STRUCT<x BIGINT>) AS q0, + CAST(s AS STRUCT<pad VARCHAR>) AS q1 +FROM exact ORDER BY id; +---- +{x: 100} {pad: sp1} +{x: 200} {pad: sp2} +NULL NULL + +# Overlapping targets: q1 needs a leaf q0's clip would have dropped. +query ?? +SELECT CAST(s AS STRUCT<x BIGINT>) AS q0, + CAST(s AS STRUCT<x BIGINT, y VARCHAR>) AS q1 +FROM exact ORDER BY id; +---- +{x: 100} {x: 100, y: s1} +{x: 200} {x: 200, y: s2} +NULL NULL + +# Repeated identical targets still clip. +query ?? +SELECT CAST(s AS STRUCT<x BIGINT>) AS q0, + CAST(s AS STRUCT<x BIGINT>) AS q1 +FROM exact ORDER BY id; +---- +{x: 100} {x: 100} +{x: 200} {x: 200} +NULL NULL + +# A cast alongside a whole-column reference: the whole-column read wins. +query ?? +SELECT CAST(s AS STRUCT<x BIGINT>) AS q0, s +FROM exact ORDER BY id; +---- +{x: 100} {x: 100, y: s1, pad: sp1} +{x: 200} {x: 200, y: s2, pad: sp2} +NULL NULL + +# The conflicting-target fallback reads the whole column -- exactly what a +# scan with no clipping at all reads, and never more. These two must match: +# the first falls back, the second never clips in the first place. +query TT +explain analyze select CAST(s AS STRUCT<x BIGINT>) AS q0, CAST(s AS STRUCT<pad VARCHAR>) AS q1 from exact; +---- +Plan with Metrics DataSourceExec: <slt:ignore>metrics=[output_rows=3, <slt:ignore>bytes_scanned=219<slt:ignore>] + +query TT +explain analyze select s from exact; +---- +Plan with Metrics DataSourceExec: <slt:ignore>metrics=[output_rows=3, <slt:ignore>bytes_scanned=219<slt:ignore>] + +statement ok +DROP TABLE exact; + +# A query cast stacked on top of the adapter's cast for a narrowed table. +query ? +SELECT CAST(s AS STRUCT<y VARCHAR>) FROM narrow ORDER BY id; +---- +{y: s1} +{y: s2} +NULL + +statement ok +DROP TABLE narrow; + +statement ok +DROP TABLE full_schema; --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
