kosiew commented on code in PR #23391:
URL: https://github.com/apache/datafusion/pull/23391#discussion_r3548595384
##########
datafusion/datasource-parquet/src/row_filter.rs:
##########
@@ -2141,4 +2292,77 @@ mod test {
let batch = RecordBatch::new_empty(Arc::clone(table_schema));
expr.evaluate(&batch).is_ok()
}
+
+ fn list_of(inner: DataType) -> DataType {
+ DataType::List(Arc::new(Field::new("element", inner, true)))
+ }
+
+ fn struct_of(fields: Vec<Field>) -> DataType {
+ DataType::Struct(fields.into_iter().map(Arc::new).collect::<Fields>())
+ }
+
+ #[test]
Review Comment:
The added tests exercise `prune_and_collect` directly, which is helpful, but
they do not quite cover the main behavior this PR changes: recognizing a real
`CastExpr(Column, narrow_type)` projection and turning it into the correct read
plan.
A regression in `try_nested_projection_leaves` dispatch,
`expr.data_type(file_schema)`, column rebasing, `ProjectionMask::leaves`, or
`projected_schema` construction would not be caught by the current tests.
Could you please add at least one integration-level regression test through
`build_projection_read_plan` or `DecoderProjection::try_new` using the real
`CastExpr(Column, narrow_type)` shape, and assert the leaf mask and projected
schema for `array<struct<a,b>> -> array<struct<a>>`?
Since the PR explicitly handles maps, it would also be good to include a
narrowed map value struct case. Otherwise, we should probably defer map support
if it is not intended to be covered here.
##########
datafusion/datasource-parquet/src/row_filter.rs:
##########
@@ -732,6 +743,147 @@ pub(crate) fn build_projection_read_plan(
}
}
+/// `field` rebuilt with a new data type, preserving its name, nullability,
and metadata.
+fn with_type(field: &Field, dt: DataType) -> Arc<Field> {
+ Arc::new(field.clone().with_data_type(dt))
+}
+
+/// Element type of `pruned` if it is a list variant.
+fn pruned_list_element(pruned: Option<&DataType>) -> Option<&DataType> {
+ match pruned {
+ Some(
+ DataType::List(f) | DataType::LargeList(f) |
DataType::FixedSizeList(f, _),
+ ) => Some(f.data_type()),
+ _ => None,
+ }
+}
+
+/// Prune `full` to the subtree present in `pruned`, advancing `*leaf` across
every
+/// primitive of `full` in parquet leaf order and recording the kept indices
in `out`.
+/// Names and nullability come from `full` so the result matches the decoder's
output.
+/// Returns `None` when no leaf under `full` is kept, which lets callers
detect a
+/// request that is not a structural subtree (e.g. a `get_field` extraction)
and defer.
+fn prune_and_collect(
+ full: &DataType,
+ pruned: Option<&DataType>,
+ leaf: &mut usize,
+ out: &mut Vec<usize>,
+) -> Option<DataType> {
+ match full {
+ DataType::Struct(fields) => {
+ let pruned = match pruned {
+ Some(DataType::Struct(p)) => Some(p),
+ _ => None,
+ };
+ let kept: Fields = fields
+ .iter()
+ .filter_map(|f| {
+ let child = pruned
+ .and_then(|p| p.iter().find(|x| x.name() == f.name()))
+ .map(|x| x.data_type());
+ Some(with_type(
+ f,
+ prune_and_collect(f.data_type(), child, leaf, out)?,
+ ))
+ })
+ .collect();
+ (!kept.is_empty()).then_some(DataType::Struct(kept))
+ }
+ DataType::List(f) => {
+ prune_and_collect(f.data_type(), pruned_list_element(pruned),
leaf, out)
+ .map(|dt| DataType::List(with_type(f, dt)))
+ }
+ DataType::LargeList(f) => {
+ prune_and_collect(f.data_type(), pruned_list_element(pruned),
leaf, out)
+ .map(|dt| DataType::LargeList(with_type(f, dt)))
+ }
+ DataType::FixedSizeList(f, n) => {
+ prune_and_collect(f.data_type(), pruned_list_element(pruned),
leaf, out)
+ .map(|dt| DataType::FixedSizeList(with_type(f, dt), *n))
+ }
+ DataType::Map(entries, sorted) => {
+ let child = match pruned {
+ Some(DataType::Map(p, _)) => Some(p.data_type()),
+ _ => None,
+ };
+ prune_and_collect(entries.data_type(), child, leaf, out)
+ .map(|dt| DataType::Map(with_type(entries, dt), *sorted))
+ }
+ _ => {
+ if pruned.is_some() {
+ out.push(*leaf);
+ }
+ *leaf += 1;
+ pruned.map(|_| full.clone())
+ }
+ }
+}
+
+/// Build a leaf-pruned read plan for projections that narrow nested columns,
or `None`
+/// to defer to the generic path. Fires only when every expression projects a
single
+/// distinct file column and at least one requests a nested type narrower than
the
+/// file's; then only the leaves reached by each requested type are read.
Keyed off
+/// `expr.data_type`, so it is independent of the concrete projection
expression.
+fn try_nested_projection_leaves(
+ exprs: &[Arc<dyn PhysicalExpr>],
+ file_schema: &Schema,
+ schema_descr: &SchemaDescriptor,
+) -> Option<ParquetReadPlan> {
+ // First parquet leaf index of each root column. Iterating leaves high to
low, each
+ // root's slot ends holding its lowest leaf index, where prune_and_collect
starts.
+ let mut offsets = vec![0usize; file_schema.fields().len()];
+ for leaf in (0..schema_descr.num_columns()).rev() {
+ offsets[schema_descr.get_column_root_idx(leaf)] = leaf;
+ }
+
+ let mut leaves = Vec::new();
+ let mut projected: Vec<(usize, Arc<Field>)> = Vec::new();
+ let mut roots = std::collections::HashSet::new();
Review Comment:
Small style thought: consider importing `HashSet` with the existing
collection imports, or using `BTreeSet` consistently in this file. Not a
blocker, but it would keep this helper visually aligned with the surrounding
projection-pruning code.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
