diff --git a/datafusion/datasource-parquet/src/nested_schema_pruning.rs b/datafusion/datasource-parquet/src/nested_schema_pruning.rs index f84ae788842e5..864651402240b 100644 --- a/datafusion/datasource-parquet/src/nested_schema_pruning.rs +++ b/datafusion/datasource-parquet/src/nested_schema_pruning.rs @@ -139,9 +139,9 @@ pub(crate) fn clip_for_cast( /// This is the second half of merging several nested accesses to one root: /// callers union the offsets required by each access, then use this function /// to derive the Arrow type the Parquet reader emits for that union. Partial -/// projection is supported through the same struct, list, and large-list -/// shapes as [`clip_for_cast`]. Any partial selection below another wrapper -/// returns `None`, preserving the total-fallback property of cast clipping. +/// projection is supported through the same struct and wrapper shapes as +/// [`clip_for_cast`]. Any partial selection below another wrapper returns +/// `None`, preserving the total-fallback property of cast clipping. pub(crate) fn type_for_leaf_subset( physical: &DataType, kept: &[usize], @@ -352,6 +352,28 @@ fn project_type_for_leaves( Ok(project_type_for_leaves(item.data_type(), kept, next_leaf)? .map(|projected| DataType::LargeList(field_with_type(item, projected)))) } + DataType::ListView(item) => { + Ok(project_type_for_leaves(item.data_type(), kept, next_leaf)? + .map(|projected| DataType::ListView(field_with_type(item, projected)))) + } + DataType::LargeListView(item) => Ok(project_type_for_leaves( + item.data_type(), + kept, + next_leaf, + )? + .map(|projected| DataType::LargeListView(field_with_type(item, projected)))), + DataType::FixedSizeList(item, size) => Ok(project_type_for_leaves( + item.data_type(), + kept, + next_leaf, + )? + .map(|projected| { + DataType::FixedSizeList(field_with_type(item, projected), *size) + })), + DataType::Dictionary(key, value) => { + Ok(project_type_for_leaves(value, kept, next_leaf)? + .map(|projected| DataType::Dictionary(key.clone(), Box::new(projected)))) + } _ => { let start = *next_leaf; let end = start + count_leaves(physical); @@ -602,17 +624,48 @@ mod tests { /// must fall back instead of predicting a type the reader may not emit. #[test] fn type_for_leaf_subset_rejects_unsupported_partial_wrapper() { - let physical = DataType::FixedSizeList( + let physical = DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), Arc::new(Field::new( - "item", + "values", struct_of(vec![int64("x"), utf8("pad")]), true, )), - 2, ); assert!(type_for_leaf_subset(&physical, &[0]).is_none()); } + /// The read plan keeps only [`clip_for_cast`]'s leaf offsets and derives + /// the emitted type from them, so [`type_for_leaf_subset`] must reproduce + /// the clipped type for every wrapper the clip descends through. + #[test] + fn type_for_leaf_subset_matches_clip_through_wrappers() { + let item = |fields| Arc::new(Field::new("item", struct_of(fields), true)); + let wrappers: [fn(FieldRef) -> DataType; 6] = [ + DataType::List, + DataType::LargeList, + DataType::ListView, + DataType::LargeListView, + |item| DataType::FixedSizeList(item, 2), + |item| { + DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(item.data_type().clone()), + ) + }, + ]; + for wrap in wrappers { + let physical = wrap(item(vec![int64("x"), utf8("pad")])); + let target = wrap(item(vec![int64("x")])); + let (kept, clipped) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!( + type_for_leaf_subset(&physical, &kept), + Some(clipped), + "{physical}" + ); + } + } + /// 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 diff --git a/datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt b/datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt index 2f65afa3eab40..4265896251d33 100644 --- a/datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt +++ b/datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt @@ -432,6 +432,61 @@ SELECT id, s FROM mixed_files ORDER BY id; statement ok DROP TABLE mixed_files; +########## +# Fixed-size lists clip like lists: a table declaring a narrower element +# struct reads only the leaves it declares. +########## + +statement ok +COPY ( + SELECT id, CAST(f AS STRUCT[2]) AS f + FROM (VALUES + (1, [named_struct('x', 10, 'pad', 'p1'), named_struct('x', 11, 'pad', 'p2')]), + (2, [named_struct('x', 20, 'pad', 'p3'), named_struct('x', 21, 'pad', 'p4')]), + (3, NULL) + ) AS t(id, f) +) TO 'test_files/scratch/parquet_nested_schema_pruning/fixed_size.parquet' +STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE fixed_size_narrow ( + id INT, + f STRUCT[2] +) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/fixed_size.parquet'; + +statement ok +CREATE EXTERNAL TABLE fixed_size_full ( + id INT, + f STRUCT[2] +) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/fixed_size.parquet'; + +query I? +SELECT id, f FROM fixed_size_narrow ORDER BY id; +---- +1 [{x: 10}, {x: 11}] +2 [{x: 20}, {x: 21}] +3 NULL + +query TT +explain analyze select f from fixed_size_narrow; +---- +Plan with Metrics DataSourceExec: metrics=[output_rows=3, bytes_scanned=83.0 B] + +query TT +explain analyze select f from fixed_size_full; +---- +Plan with Metrics DataSourceExec: metrics=[output_rows=3, bytes_scanned=174.0 B] + +statement ok +DROP TABLE fixed_size_narrow; + +statement ok +DROP TABLE fixed_size_full; + ########## # 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