jayzhan211 commented on code in PR #25013:
URL: https://github.com/apache/datafusion/pull/25013#discussion_r3978882438
##########
datafusion/datasource-parquet/src/projection_read_plan.rs:
##########
@@ -394,90 +385,99 @@ impl TreeNodeVisitor<'_> for PushdownChecker<'_> {
type Node = Arc<dyn PhysicalExpr>;
fn f_down(&mut self, node: &Self::Node) -> Result<TreeNodeRecursion> {
- // Handle struct field access like `s['foo']['bar'] > 10`.
- //
- // DataFusion represents nested field access as
`get_field(Column("s"), "foo")`
- // (or chained: `get_field(get_field(Column("s"), "foo"), "bar")`).
- //
- // We intercept the outermost `get_field` on the way *down* the tree so
- // the visitor never reaches the raw `Column("s")` node. Without this,
- // `check_single_column` would see that `s` is a Struct and reject it.
- //
- // The strategy:
- // 1. Match `get_field` whose first arg is a `Column` (the struct
root).
- // 2. Check that the *resolved* return type is primitive — meaning
we've
- // drilled all the way to a leaf (e.g. `s['foo']` → Utf8).
- // 3. Record the root column index via `check_struct_field_column`
and
- // return `Jump` to skip visiting the children (the Column and the
- // literal field-name args), since we've already handled them.
- //
- // If the return type is still nested (e.g. `s['nested_struct']` →
Struct),
- // we fall through and let normal traversal continue, which will
- // eventually reject the expression when it hits the struct Column.
- if let Some(func) =
-
ScalarFunctionExpr::try_downcast_func::<GetFieldFunc>(node.as_ref())
- {
- if let Some(recursion) = self.check_cast_struct_field_access(func)
{
+ // Resolve capability-declaring accessors, including chains with
+ // different UDFs and argument layouts. Do not look through casts.
+ let mut source = node;
+ let mut paths = Vec::new();
+ while let Some(function) = source.downcast_ref::<ScalarFunctionExpr>()
{
+ let Some(access) = function.struct_field_access() else {
+ break;
+ };
+ paths.push(access.field_path);
+ source = &function.args()[access.source_arg];
+ }
+ let field_path = paths.into_iter().rev().flatten().collect::<Vec<_>>();
+ if !field_path.is_empty() {
+ let return_type = node.data_type(self.file_schema)?;
+ if let Some(recursion) =
+ self.check_cast_struct_field_access(source, &field_path,
&return_type)
+ {
return Ok(recursion);
}
- let args = func.args();
-
- if let Some(column) = args.first().and_then(|a|
a.downcast_ref::<Column>()) {
- // for Map columns, get_field performs a runtime key lookup
rather than a
- // schema-level field access so the entire Map column must be
read,
- // we skip the struct field optimization and defer to normal
Column traversal
- let is_map_column = self
+ if let Some(column) = source.downcast_ref::<Column>() {
+ // Resolve by name: physical column indices may still refer to
+ // an unprojected schema. Map/List paths must remain opaque.
+ let leaf_type = self
.file_schema
- .index_of(column.name())
+ .field_with_name(column.name())
.ok()
- .map(|idx| {
- matches!(
- self.file_schema.field(idx).data_type(),
- DataType::Map(_, _)
- )
- })
- .unwrap_or(false);
-
- let return_type = func.return_type();
+ .and_then(|root| {
+ field_path.iter().try_fold(root.data_type(), |ty,
name| {
+ let DataType::Struct(fields) = ty else {
+ return None;
+ };
+ let mut matches =
+ fields.iter().filter(|field| field.name() ==
name);
+ let field = matches.next()?;
+ if matches.next().is_some() {
+ return None;
+ }
+ Some(field.data_type())
+ })
+ });
+ if leaf_type.is_some()
+ && (!return_type.is_nested()
+ || self.is_nested_type_supported(&return_type))
+ {
+ if let Some(recursion) =
+ self.check_struct_field_column(column.name(),
field_path)
+ {
+ return Ok(recursion);
+ }
+ return Ok(TreeNodeRecursion::Jump);
+ }
+ }
+ }
- if !is_map_column
- && (!DataType::is_nested(return_type)
- || self.is_nested_type_supported(return_type))
+ if let Some(function) = node.downcast_ref::<ScalarFunctionExpr>()
+ && let Some(requirements) =
function.required_input_fields(self.file_schema)
+ && !requirements.is_empty()
+ {
+ for (index, argument) in function.args().iter().enumerate() {
Review Comment:
Thanks @peterxcli , here is a suggestion:
**`required_input_fields` turns off the nested-column pushdown gate**
`projection_read_plan.rs:443-480`. The requirements branch returns
`TreeNodeRecursion::Jump`
for every argument it claims, which skips `check_single_column` →
`handle_nested_type`. That's
the check that sets `non_primitive_columns` for List/Map columns unless
`allow_list_columns`
(i.e. `supports_list_predicates`, the verified `array_has*` / `IS NULL`
allow-list) is true.
The accessor branch immediately above is guarded — `!return_type.is_nested()
||
self.is_nested_type_supported(&return_type)` — but this one has no type
check at all, so
`prevents_pushdown()` stays `false`.
Confirmed with a scratch probe against `PushdownChecker::new(&schema,
/*allow_list_columns=*/ false, false)`:
```rust
// schema: s: List<Int32>
fn required_input_fields(&self, _: ReturnFieldArgs) ->
Option<Vec<InputFieldRequirement>> {
Some(vec![InputFieldRequirement { arg_index: 0, field_paths:
vec![vec![]] }])
}
// => prevents_pushdown() == false (expected: true)
```
Same result for `schema: s: Struct<events: List<Int32>>` with `field_paths:
vec![vec!["events".into()]]`,
and the same hole applies to `Map`. So any downstream UDF can opt itself
past the gate by
declaring a requirement that prunes nothing — which contradicts this PR's
own contract text:
"does not ... authorize moving the function across arbitrary operators".
Suggested fix — resolve each declared path's leaf type and apply the
existing policy before
taking the shortcut, falling back to normal traversal otherwise:
```diff
if let Some(function) = node.downcast_ref::<ScalarFunctionExpr>()
&& let Some(requirements) =
function.required_input_fields(self.file_schema)
&& !requirements.is_empty()
+ && requirements.iter().all(|requirement| {
+ // Reading a nested column into the row filter follows the same
policy
+ // as an accessor: Struct subtrees are fine, other nested types only
+ // when the predicate set supports them. A declaration must not
widen it.
+ function.args()[requirement.arg_index]
+ .return_field(self.file_schema)
+ .is_ok_and(|field| {
+ requirement.field_paths.iter().all(|path| {
+ resolve_leaf_type(field.data_type(),
path).is_some_and(|leaf| {
+ matches!(leaf, DataType::Struct(_))
+ || !DataType::is_nested(leaf)
+ || self.is_nested_type_supported(leaf)
+ })
+ })
+ })
+ })
{
```
(`resolve_leaf_type` being the same name-resolving `Struct`-only walk
already inlined in the
accessor branch — worth factoring out, since it now has three copies.)
A regression test in the same style as
`custom_struct_accessor_does_not_prune_map_entries`
asserting `prevents_pushdown()` for a List-bearing requirement would lock
this down.
--
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]
