Skip to content

Commit ec4b3f2

Browse files
sunchaoadriangb
andcommitted
perf: prune retained Struct cast reads in Parquet filters
Reuse projection read clipping for retained Struct casts while preserving every conversion named by the cast target and sizing only selected leaves. Add SQL coverage for explicit sibling errors and all-null scalar/List decimal adaptation, strengthen partial-clipping coverage, and clarify why Struct ancestors remain around the covered decimal conversions. Co-authored-by: Adrian Garcia Badaracco <[email protected]>
1 parent e79e912 commit ec4b3f2

4 files changed

Lines changed: 286 additions & 43 deletions

File tree

‎datafusion/datasource-parquet/src/projection_read_plan.rs‎

Lines changed: 51 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -178,15 +178,16 @@ pub(crate) struct PushdownChecker<'schema> {
178178
/// Struct field accesses via `get_field`.
179179
struct_field_accesses: Vec<StructFieldAccess>,
180180
/// Whole-column casts to a narrower nested type
181-
/// (`CAST(col AS narrower_struct)`). Only collected when
182-
/// [`Self::with_cast_collection`] enables it (projection analysis);
183-
/// filter pushdown leaves this off.
181+
/// (`CAST(col AS narrower_struct)`), collected either when
182+
/// [`Self::with_cast_collection`] enables it (projection analysis) or when
183+
/// [`Self::allow_struct_casts`] accepts a retained cast under a `get_field`
184+
/// (filter pushdown).
184185
cast_accesses: Vec<CastColumnAccess>,
185186
/// Whether to collect [`Self::cast_accesses`].
186187
collect_cast_accesses: bool,
187188
/// Allow `get_field(CAST(struct_column AS Struct(...)), 'field', ...)`
188-
/// after schema adaptation, preserving the cast and reading the full source
189-
/// Struct. Both source and target must be Struct types.
189+
/// after schema adaptation, preserving the cast and reading the leaves its
190+
/// target names. Both source and target must be Struct types.
190191
/// Planning keeps this disabled so explicit casts retain a residual filter.
191192
allow_struct_casts: bool,
192193
/// Whether nested list columns are supported by the predicate semantics.
@@ -258,8 +259,17 @@ impl<'schema> PushdownChecker<'schema> {
258259
None
259260
}
260261

261-
/// Preserve a Struct cast retained by schema adaptation and read its full
262-
/// root. Pruning siblings or moving the cast could change errors or nulls.
262+
/// Preserve a Struct cast retained by schema adaptation and record the
263+
/// leaves its target consumes.
264+
///
265+
/// The cast is kept intact — moving it could change errors or nulls — but
266+
/// the target itself names every field the conversion touches, so the read
267+
/// can be clipped to those leaves. `cast_struct_column` resolves source
268+
/// children by name and ignores the rest, so a leaf the target does not
269+
/// name cannot affect the result. Note this clips by the *cast target*, not
270+
/// by the `get_field` key: an explicit query cast names every field the
271+
/// user asked to convert, so its siblings stay in the read and their
272+
/// conversions still run.
263273
fn check_cast_struct_field_access(
264274
&mut self,
265275
func: &ScalarFunctionExpr,
@@ -304,7 +314,10 @@ impl<'schema> PushdownChecker<'schema> {
304314
.data_type();
305315
}
306316

307-
self.required_columns.push(index);
317+
self.cast_accesses.push(CastColumnAccess {
318+
root_index: index,
319+
target_type: cast.cast_type().clone(),
320+
});
308321
Some(TreeNodeRecursion::Jump)
309322
}
310323

@@ -518,8 +531,8 @@ pub(crate) struct PushdownColumns {
518531
/// Struct field accesses via `get_field`. Each entry records the root struct
519532
/// column index and the field path being accessed.
520533
pub(crate) struct_field_accesses: Vec<StructFieldAccess>,
521-
/// Whole-column casts to a narrower nested type. Empty unless cast
522-
/// collection was enabled on the checker.
534+
/// Whole-column casts to a narrower nested type, collected for projections
535+
/// or retained Struct casts accepted by the runtime filter checker.
523536
pub(crate) cast_accesses: Vec<CastColumnAccess>,
524537
}
525538

@@ -607,13 +620,14 @@ pub(crate) fn build_projection_read_plan(
607620
all_cast_accesses.retain(|c| all_root_indices.binary_search(&c.root_index).is_err());
608621

609622
if !all_cast_accesses.is_empty() {
610-
return build_read_plan_with_cast_clipping(
623+
let (read_plan, _leaf_indices) = build_read_plan_with_cast_clipping(
611624
file_schema,
612625
schema_descr,
613626
&all_root_indices,
614627
&all_struct_accesses,
615628
&all_cast_accesses,
616629
);
630+
return read_plan;
617631
}
618632

619633
// when no struct field accesses were found, fall back to root-level projection
@@ -645,8 +659,7 @@ enum RootRead {
645659
///
646660
/// Per root, in ascending root-index order:
647661
/// - roots referenced as whole columns keep every leaf and their full
648-
/// physical field (whole-column reads take precedence; cast accesses on
649-
/// such roots were already dropped by the caller);
662+
/// physical field (whole-column reads take precedence over cast accesses);
650663
/// - roots consumed through one or more casts keep the union of the leaves
651664
/// their targets name (see `crate::nested_schema_pruning`);
652665
/// - a root consumed through both casts and `get_field` accesses keeps the
@@ -657,13 +670,16 @@ enum RootRead {
657670
/// `nested_schema_pruning::clip_for_cast`), an access that resolves to no
658671
/// leaf at all, or a merged leaf set whose emitted Arrow type can't be
659672
/// derived safely, falls back to a full read of that root.
660-
fn build_read_plan_with_cast_clipping(
673+
///
674+
/// Also returns the resolved Parquet leaf indices, sorted and deduplicated, so
675+
/// callers can size the columns the decoder will read.
676+
pub(crate) fn build_read_plan_with_cast_clipping(
661677
file_schema: &Schema,
662678
schema_descr: &SchemaDescriptor,
663679
whole_root_indices: &[usize],
664680
struct_accesses: &[StructFieldAccess],
665681
cast_accesses: &[CastColumnAccess],
666-
) -> ParquetReadPlan {
682+
) -> (ParquetReadPlan, Vec<usize>) {
667683
// Every referenced root's Parquet leaves, grouped in one pass over the
668684
// schema descriptor rather than one `leaf_indices_for_roots` scan per
669685
// root (this function may look up several roots).
@@ -767,17 +783,24 @@ fn build_read_plan_with_cast_clipping(
767783
fields.push(Arc::new(field.clone()));
768784
}
769785
// `ProjectionMask::leaves` only flips flags in a `vec![false; num_columns]`,
770-
// so `leaf_indices` needs no sorting or deduplication here.
771-
ParquetReadPlan {
772-
projection_mask: ProjectionMask::leaves(
773-
schema_descr,
774-
leaf_indices.iter().copied(),
775-
),
776-
projected_schema: Arc::new(Schema::new_with_metadata(
777-
fields,
778-
file_schema.metadata().clone(),
779-
)),
780-
}
786+
// so the mask itself needs neither ordering nor deduplication. Callers that
787+
// size the read do care, so normalize before handing the indices back:
788+
// `size_of_columns` sums per index and would double-count a repeat.
789+
leaf_indices.sort_unstable();
790+
leaf_indices.dedup();
791+
(
792+
ParquetReadPlan {
793+
projection_mask: ProjectionMask::leaves(
794+
schema_descr,
795+
leaf_indices.iter().copied(),
796+
),
797+
projected_schema: Arc::new(Schema::new_with_metadata(
798+
fields,
799+
file_schema.metadata().clone(),
800+
)),
801+
},
802+
leaf_indices,
803+
)
781804
}
782805

783806
/// Groups every Parquet leaf index by its root (Arrow) column index, in one
@@ -1757,7 +1780,7 @@ mod test {
17571780
vec![Arc::new(Field::new("p", DataType::Int32, true))].into(),
17581781
),
17591782
};
1760-
let read_plan = build_read_plan_with_cast_clipping(
1783+
let (read_plan, _leaf_indices) = build_read_plan_with_cast_clipping(
17611784
&file_schema,
17621785
schema_descr,
17631786
&[],
@@ -1927,7 +1950,7 @@ mod test {
19271950
Schema::new(vec![divergent("a", "p", "q"), divergent("b", "m", "n")]);
19281951

19291952
// `a` is reached by a narrowing cast, `b` only by `get_field`.
1930-
let read_plan = build_read_plan_with_cast_clipping(
1953+
let (read_plan, _leaf_indices) = build_read_plan_with_cast_clipping(
19311954
&file_schema,
19321955
schema_descr,
19331956
&[],

‎datafusion/datasource-parquet/src/row_filter.rs‎

Lines changed: 101 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,7 @@ use super::ParquetFileMetrics;
8787
use super::supported_predicates::supports_list_predicates;
8888
use crate::projection_read_plan::{
8989
ParquetReadPlan, PushdownChecker, PushdownColumns, assemble_read_plan,
90+
build_read_plan_with_cast_clipping,
9091
};
9192

9293
/// A "compiled" predicate passed to `ParquetRecordBatchStream` to perform
@@ -281,12 +282,26 @@ pub(crate) fn build_parquet_read_plan(
281282
return Ok(None);
282283
};
283284

284-
let (read_plan, leaf_indices) = assemble_read_plan(
285-
&required_columns.required_columns,
286-
&required_columns.struct_field_accesses,
287-
file_schema,
288-
schema_descr,
289-
);
285+
// A retained Struct cast names the fields its conversion touches, so the
286+
// read is clipped to those leaves rather than decoding the whole root.
287+
// A cast whose target covers every leaf, or that cannot be clipped safely,
288+
// falls back to a full read of that root inside the helper.
289+
let (read_plan, leaf_indices) = if required_columns.cast_accesses.is_empty() {
290+
assemble_read_plan(
291+
&required_columns.required_columns,
292+
&required_columns.struct_field_accesses,
293+
file_schema,
294+
schema_descr,
295+
)
296+
} else {
297+
build_read_plan_with_cast_clipping(
298+
file_schema,
299+
schema_descr,
300+
&required_columns.required_columns,
301+
&required_columns.struct_field_accesses,
302+
&required_columns.cast_accesses,
303+
)
304+
};
290305

291306
let required_bytes = size_of_columns(&leaf_indices, metadata)?;
292307

@@ -1237,11 +1252,12 @@ mod test {
12371252
fn get_field_filter_candidate_has_correct_leaf_indices() {
12381253
use arrow::array::{Int32Array, StringArray, StructArray};
12391254

1240-
// Schema: id (Int32), s (Struct{value: Int32, label: Utf8})
1241-
// Parquet leaves: id=0, s.value=1, s.label=2
1255+
// Schema: id (Int32), s (Struct{value: Int32, label: Utf8, unused: Utf8})
1256+
// Parquet leaves: id=0, s.value=1, s.label=2, s.unused=3
12421257
let struct_fields: Fields = vec![
12431258
Arc::new(Field::new("value", DataType::Int32, false)),
12441259
Arc::new(Field::new("label", DataType::Utf8, false)),
1260+
Arc::new(Field::new("unused", DataType::Utf8, false)),
12451261
]
12461262
.into();
12471263
let schema = Arc::new(Schema::new(vec![
@@ -1258,6 +1274,9 @@ mod test {
12581274
vec![
12591275
Arc::new(Int32Array::from(vec![10, 20, 30])) as _,
12601276
Arc::new(StringArray::from(vec!["a", "b", "c"])) as _,
1277+
Arc::new(StringArray::from(vec![
1278+
"unused-a", "unused-b", "unused-c",
1279+
])) as _,
12611280
],
12621281
None,
12631282
)),
@@ -1293,16 +1312,20 @@ mod test {
12931312
.expect("get_field filter on struct should be pushable");
12941313

12951314
// The filter accesses only s.value, so only Parquet leaf 1 is needed.
1296-
// Leaf 2 (s.label) is not read, reducing unnecessary I/O.
1315+
// Neither sibling is read, reducing unnecessary I/O.
12971316
let expected_mask =
12981317
ProjectionMask::leaves(metadata.file_metadata().schema_descr(), [1]);
12991318
assert_eq!(
13001319
candidate.read_plan.projection_mask, expected_mask,
13011320
"projection_mask should select only the accessed struct field leaf"
13021321
);
13031322

1304-
// Schema adaptation can leave a Struct cast intact. Its runtime filter
1305-
// must read every sibling, while planning still rejects explicit casts.
1323+
// Schema adaptation retains Struct ancestors for some decimal conversions
1324+
// so an all-null parent can skip child conversion. Runtime filters must
1325+
// preserve every conversion named by the retained cast target.
1326+
// This target includes `label`, so its conversion must still run even
1327+
// though get_field selects only `value`. Planning keeps a residual filter
1328+
// for explicit Struct casts.
13061329
let cast_type = DataType::Struct(
13071330
vec![
13081331
Field::new("value", DataType::Int32, false),
@@ -1325,12 +1348,29 @@ mod test {
13251348
.build(&metadata)
13261349
.expect("building cast candidate")
13271350
.expect("an adapted struct cast must remain evaluable");
1351+
// Clip the unused sibling, but preserve the failing `label` conversion.
13281352
let expected_mask =
1329-
ProjectionMask::roots(metadata.file_metadata().schema_descr(), [1]);
1353+
ProjectionMask::leaves(metadata.file_metadata().schema_descr(), [1, 2]);
13301354
assert_eq!(candidate.read_plan.projection_mask, expected_mask);
1355+
let DataType::Struct(physical_fields) = file_schema.field(1).data_type() else {
1356+
unreachable!("s is a struct")
1357+
};
1358+
let clipped_field =
1359+
file_schema
1360+
.field(1)
1361+
.clone()
1362+
.with_data_type(DataType::Struct(
1363+
physical_fields.iter().take(2).cloned().collect(),
1364+
));
13311365
assert_eq!(
13321366
candidate.read_plan.projected_schema.as_ref(),
1333-
&file_schema.project(&[1]).unwrap()
1367+
&Schema::new(vec![clipped_field])
1368+
);
1369+
assert_eq!(
1370+
candidate.required_bytes,
1371+
(metadata.row_group(0).column(1).compressed_size()
1372+
+ metadata.row_group(0).column(2).compressed_size()) as usize,
1373+
"filter cost must count only the leaves the cast target reads"
13341374
);
13351375

13361376
// A simultaneous direct access must not prune siblings that the cast
@@ -1362,6 +1402,54 @@ mod test {
13621402
.unwrap();
13631403
let error = row_filter.evaluate(batch).unwrap_err().to_string();
13641404
datafusion_common::assert_contains!(error, "While casting struct field 'label'");
1405+
1406+
// A retained cast whose target names only the selected field — the
1407+
// shape `retain_field_path` produces for an evolved decimal — clips the
1408+
// read to that field's leaf instead of decoding the whole root.
1409+
let narrow_cast_type =
1410+
DataType::Struct(vec![Field::new("value", DataType::Int32, false)].into());
1411+
let narrow_field = get_field().call(vec![
1412+
datafusion_expr::cast(col("s"), narrow_cast_type.clone()),
1413+
lit("value"),
1414+
]);
1415+
let narrow_predicate = logical2physical(&narrow_field.gt(lit(5)), &file_schema);
1416+
let candidate =
1417+
FilterCandidateBuilder::new(narrow_predicate, Arc::clone(&file_schema))
1418+
.build(&metadata)
1419+
.expect("building narrow cast candidate")
1420+
.expect("a clipped struct cast must remain evaluable");
1421+
assert_eq!(
1422+
candidate.read_plan.projection_mask,
1423+
ProjectionMask::leaves(metadata.file_metadata().schema_descr(), [1]),
1424+
"the read must be clipped to the leaf the cast target names"
1425+
);
1426+
assert_eq!(candidate.read_plan.projected_schema.fields().len(), 1);
1427+
assert_eq!(
1428+
candidate.read_plan.projected_schema.field(0).data_type(),
1429+
&narrow_cast_type,
1430+
"sibling leaves must be pruned from the filter schema"
1431+
);
1432+
1433+
// The clipped schema must still evaluate: every row has value > 5.
1434+
let mut row_filter = DatafusionArrowPredicate::try_new(
1435+
candidate,
1436+
Count::new(),
1437+
Count::new(),
1438+
Time::new(),
1439+
)
1440+
.unwrap();
1441+
let batch = ParquetRecordBatchReaderBuilder::try_new(file.reopen().unwrap())
1442+
.unwrap()
1443+
.with_projection(row_filter.projection().clone())
1444+
.build()
1445+
.unwrap()
1446+
.next()
1447+
.unwrap()
1448+
.unwrap();
1449+
assert_eq!(
1450+
row_filter.evaluate(batch).unwrap(),
1451+
BooleanArray::from(vec![true, true, true])
1452+
);
13651453
}
13661454

13671455
/// Deeply nested get_field: get_field(struct_col, 'outer', 'inner') where the

‎datafusion/physical-expr-adapter/src/schema_rewriter.rs‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -530,7 +530,14 @@ impl DefaultPhysicalExprAdapterRewriter {
530530
// narrowable by a parent get_field.
531531
// Container casts involving Struct values must keep their existing
532532
// dispatch: Arrow can unwrap a container into a Struct where cast_column
533-
// cannot.
533+
// cannot. That leaves one shape uncovered: unwrapping hands the Struct
534+
// to Arrow's own cast, which has no all-null shortcut, so a decimal
535+
// below it can still fail on an all-null input (for example
536+
// `Dictionary(Int8, Struct<b: Utf8>)` to `Struct<b: Decimal128(10, -1)>`).
537+
// The Parquet reader does not produce dictionary-encoded Struct columns,
538+
// so this is not reachable through a Parquet scan; closing it would mean
539+
// telling "Arrow must unwrap this" apart from "Arrow will convert a
540+
// decimal while unwrapping" rather than dropping the carve-out.
534541
let source_type = physical_struct_field.data_type();
535542
let target_type = logical_struct_field.data_type();
536543
let is_struct = |data_type: &DataType| matches!(data_type, DataType::Struct(_));

0 commit comments

Comments
 (0)