Skip to content

feat: expose UDF field access and input requirements for Parquet pruning - #25013

Open
peterxcli wants to merge 5 commits into
apache:mainfrom
peterxcli:feat/struct-field-access-capability
Open

feat: expose UDF field access and input requirements for Parquet pruning#25013
peterxcli wants to merge 5 commits into
apache:mainfrom
peterxcli:feat/struct-field-access-capability

Conversation

@peterxcli

@peterxcli peterxcli commented Sep 7, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

Closes #21306.

Rationale for this change

A custom accessor can miss Parquet column pruning and filter pushdown even when it reads only one field:

SELECT my_field(person, 'age')
FROM people
WHERE my_field(person, 'age') > 30;

The affected optimization paths recognize the built-in GetFieldFunc. For an arbitrary UDF, they cannot assume that the other fields are unnecessary: the function could inspect them, validate them or change null handling.

The existing placement() method describes where a function should execute. It does not describe which nested input fields it requires. These are separate facts needed by the optimizer.

There are also two distinct field-related guarantees. An ordinary accessor can promise exact field extraction. A function such as downstream variant_get may instead need several physical fields, decode fallback values and convert its result. Treating that computation as exact field extraction would be incorrect.

This draft exposes both contracts. DataFusion manages generic field paths and Parquet planning. Downstream UDFs interpret their own layouts and declare their requirements; Variant path, shredding, decoding and fallback logic stay in datafusion-variant.

What changes are included in this PR?

  • Add optional struct_field_access() for exact extraction, describing the source argument and literal field path. Implement it for GetFieldFunc and use it in Parquet planning and schema adaptation.
  • Add optional, schema-aware required_input_fields(ReturnFieldArgs) returning InputFieldRequirement { arg_index, field_paths }. It permits pruning fields while retaining the original UDF computation. Forward both hooks through ScalarUDF and aliases.
  • Validate argument indices and paths, reuse Parquet's existing union of consumers' requirements, and preserve evaluation of other arguments and explicit casts.

The input-requirements contract promises identical values, output field and errors after pruning, including when requirements are combined with other consumers. Required validation fields, metadata and encoded fallback values must be included. Selected fields and their ancestors retain metadata and validity. Paths traverse structs and may select entire nested subtrees; invalid declarations retain full inputs.

Input requirements do not establish field equivalence, output statistics, or permission for arbitrary filter movement. Exact extraction remains a stronger, separate promise. Existing implementations and FFI wrappers retain the default fallback.

What is the testing strategy for this PR?

Custom accessor tests cover reversed arguments, aliases, chained paths, literal dots, null parents and children, schema reordering, missing fields, integer widening, explicit cast errors and Map fallback. Leaf-selection assertions and decoder metrics demonstrate pruning, with capability-disabled and pushdown-disabled controls.

Additional generic-requirement tests cover a computed UDF requiring multiple fields and another argument, unions with other consumers, stale column indices in computed arguments, invalid declarations and preservation of cast errors.

Query Main then patch Patch then main
Nested equality 13.234 / 13.527 13.300 / 13.502
Top-level filter control 5.386 / 5.408 5.350 / 5.441
Nested range 15.156 / 15.399 15.030 / 14.692

Downstream validation in peterxcli/datafusion-variant#1 exercises variant_get and variant_get_field with these APIs, including encoded fallback and null semantics. Literal accessor placement also exposes aggregate arguments to scan projection pruning. The integration tests and controlled benchmarks demonstrate input pruning for projection, filters and aggregates without Variant-specific logic in DataFusion.

Are there any user-facing changes?

Custom Rust UDFs can opt into Parquet input pruning and decoder-filter evaluation through additive APIs. DataFusion gains no Variant-specific dependencies or semantics. The methods document their correctness contracts; unsupported calls keep conservative behavior.

@github-actions github-actions Bot added logical-expr Logical plan and expressions physical-expr Changes to the physical-expr crates core Core DataFusion crate sqllogictest SQL Logic Tests (.slt) functions Changes to functions implementation datasource Changes to the datasource crate labels Sep 7, 2026
@codecov-commenter

codecov-commenter commented Sep 7, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 93.29073% with 42 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.95%. Comparing base (8a92281) to head (42d4589).
⚠️ Report is 68 commits behind head on main.

Files with missing lines Patch % Lines
...ion/datasource-parquet/src/projection_read_plan.rs 93.02% 24 Missing and 8 partials ⚠️
...usion/physical-expr-adapter/src/schema_rewriter.rs 81.48% 4 Missing and 1 partial ⚠️
datafusion/physical-expr/src/scalar_function.rs 94.68% 2 Missing and 3 partials ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main   #25013      +/-   ##
==========================================
+ Coverage   81.72%   81.95%   +0.23%     
==========================================
  Files        1127     1133       +6     
  Lines      416519   424078    +7559     
  Branches   416519   424078    +7559     
==========================================
+ Hits       340401   347572    +7171     
+ Misses      56115    55905     -210     
- Partials    20003    20601     +598     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@peterxcli peterxcli changed the title feat: expose struct field access for custom scalar functions feat: enable Parquet pruning and filter pushdown for custom struct accessors Sep 8, 2026
@peterxcli
peterxcli force-pushed the feat/struct-field-access-capability branch from de7d5a3 to 1180ae3 Compare September 8, 2026 05:59
@peterxcli peterxcli changed the title feat: enable Parquet pruning and filter pushdown for custom struct accessors feat: expose UDF field access and input requirements for Parquet pruning Sep 8, 2026
@peterxcli

peterxcli commented Sep 8, 2026

Copy link
Copy Markdown
Member Author

I validated this API in peterxcli/datafusion-variant#1. Variant projection, filtering and aggregate input pruning now work while preserving encoded fallback.

Warm-cache median execution time on a wide, partially shredded dataset (1,048,576 rows; 12 trials per case; planning excluded):

Query Before Disabled control Enabled Speedup
variant_get projection 768.94 ms 755.45 ms 179.32 ms 4.29x
variant_get filter 827.26 ms 818.26 ms 183.94 ms 4.50x
SUM(variant_get(...)) 758.71 ms 762.19 ms 165.49 ms 4.58x

Projection/filter compare the original downstream implementation with the input-field capability; their disabled controls use the new dependencies with that capability off. The aggregate is a separate comparison against the previous downstream PR head, with matching dependencies and the capability enabled throughout; its control disables expression placement. Moving the accessor into the scan projection lets the aggregate prune unused inputs.

Reader bytes fall about 80%, despite retaining the encoded root value. All 405 original and 249 follow-up executions passed result checks, including encoded-fallback cases. Raw-input and follow-up projection/filter controls remain approximately unchanged. These gains are specific to this wide synthetic workload.

peterxcli/datafusion#2 separately builds on struct_field_access() for row-group statistics pruning (#20871);

&& let Some(requirements) = function.required_input_fields(self.file_schema)
&& !requirements.is_empty()
{
for (index, argument) in function.args().iter().enumerate() {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_columnhandle_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):

// 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:

 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.

@github-actions github-actions Bot removed the sqllogictest SQL Logic Tests (.slt) label Sep 10, 2026
@jayzhan211

Copy link
Copy Markdown
Contributor

@peterxcli , here is a suggestion:

requirement paths ending at a Struct bypass the List/Map pushdown policy

projection_read_plan.rs:440-446 short-circuits on Struct before looking at what the struct
contains, so the policy guard doesn't hold for anything nested below a selected sub-struct:

resolve_struct_field_type(&data_type, path).is_some_and(|leaf| {
    matches!(leaf, DataType::Struct(_))   // <-- accepts the whole subtree unseen
        || !leaf.is_nested()
        || self.is_nested_type_supported(leaf)
})

Repro (dropped into projection_read_plan::test, UDF requires ["selected"] on
s: Struct<selected: Struct<arr: List<Int32>>>, allow_list_columns = false):

PROBE prevents_pushdown (list nested in selected struct, allow_lists=false) = false
PROBE baseline whole-struct prevents_pushdown = true

The same list data reached via get_field(s,'selected') or a whole-column reference to s is
blocked. allow_list_columns comes from supports_list_predicates, i.e. the allowlist of
verified predicates (array_has/array_has_all/array_has_any, IS [NOT] NULL) — a UDF
requirement isn't on it, so declaring required_input_fields becomes a general escape hatch for
getting list/map columns into the row filter. That's precisely what the comment two lines above
says must not happen:

// Declaring dependencies cannot bypass the List/Map pushdown policy.

Map is affected identically, since the subtree is never inspected.

Suggested fix — check every leaf of the selected subtree instead of accepting Struct wholesale:

+    /// Whether every leaf below `data_type` is acceptable under the current
+    /// nested-pushdown policy. A selected sub-struct must not smuggle a List
+    /// or Map past the policy that a direct reference to it would hit.
+    fn subtree_is_pushable(&self, data_type: &DataType) -> bool {
+        match data_type {
+            DataType::Struct(fields) => fields
+                .iter()
+                .all(|field| self.subtree_is_pushable(field.data_type())),
+            other => !other.is_nested() || self.is_nested_type_supported(other),
+        }
+    }
                 data_type.is_some_and(|data_type| {
                     requirement.field_paths.iter().all(|path| {
-                        resolve_struct_field_type(&data_type, path).is_some_and(|leaf| {
-                            matches!(leaf, DataType::Struct(_))
-                                || !leaf.is_nested()
-                                || self.is_nested_type_supported(leaf)
-                        })
+                        resolve_struct_field_type(&data_type, path)
+                            .is_some_and(|leaf| self.subtree_is_pushable(leaf))
                     })
                 })

Existing udf_input_requirements_respect_nested_pushdown_policy cases still pass under this
(Struct<value: Int32> is all-primitive). Worth extending that test's type matrix with a
Struct wrapping each of the List/Map variants so the nested shape is pinned too.

@jayzhan211

Copy link
Copy Markdown
Contributor

@peterxcli , a suggestion:

try_narrow_struct_cast can now fail the query instead of skipping the rewrite

schema_rewriter.rs:565 rebuilds the accessor against the physical file schema:

let extracted = Arc::new(ScalarFunctionExpr::try_new(
    Arc::new(get_field_expr.fun().clone()),
    args,
    &self.physical_file_schema,
    Arc::new(get_field_expr.config_options().clone()),
)?) as Arc<dyn PhysicalExpr>;

try_new runs fields_with_udf and return_field_from_args on the physical types — which differ from the logical ones by construction, since that difference is why the adapter inserted the cast you're about to drop. With the old try_downcast_func::<GetFieldFunc> gate the ? was safe (Signature::user_defined, permissive). It isn't safe for arbitrary UDFs: a function that declares struct_field_access but doesn't accept the physical types now aborts the whole query, where every other path in this rewriter degrades to Ok(None).

Reproduced on this branch using your own FieldAt helper with an Exact signature pinned to the logical struct type. File s: Struct<value: Int32>, registered schema s: Struct<value: Int64>, query SELECT id FROM t WHERE field_at('value', s) > 5:

declare_access=false: Ok([...])            // 1 row, as expected
declare_access=true:  Err(Plan("Failed to coerce arguments to satisfy a call to
  'field_at' function: coercion from Utf8, Struct(\"value\": Int32) to the
  signature Exact(Utf8, Struct(\"value\": Int64)) failed"))

Declaring the capability should never make a working query fail. Make the rebuild best-effort:

-        let extracted = Arc::new(ScalarFunctionExpr::try_new(
-            Arc::new(get_field_expr.fun().clone()),
-            args,
-            &self.physical_file_schema,
-            Arc::new(get_field_expr.config_options().clone()),
-        )?) as Arc<dyn PhysicalExpr>;
+        // A third-party accessor may not accept the physical field types. That
+        // is a missed rewrite, not a query error: keep the cast and bail out.
+        let Ok(extracted) = ScalarFunctionExpr::try_new(
+            Arc::new(get_field_expr.fun().clone()),
+            args,
+            &self.physical_file_schema,
+            Arc::new(get_field_expr.config_options().clone()),
+        ) else {
+            return Ok(None);
+        };
+        let extracted = Arc::new(extracted) as Arc<dyn PhysicalExpr>;

Two follow-ups worth folding in:

  1. FieldAt::signature is currently Signature::any(2, Volatility::Immutable) at all three construction sites, i.e. the one shape that can't hit this. Point one of them at an Exact signature over the logical struct type and assert the evolved query returns the right rows (unpruned) rather than erroring.
  2. The struct_field_access doc promises implementors only "reordered or narrowed source structs". This rewrite does more — it drops the type-adapting cast and re-applies the conversion to the output, so the UDF sees the file's physical field types. Please say so in the contract, since that is the assumption this code actually relies on.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core Core DataFusion crate datasource Changes to the datasource crate functions Changes to functions implementation logical-expr Logical plan and expressions physical-expr Changes to the physical-expr crates

Projects

None yet

Development

Successfully merging this pull request may close these issues.

All struct-aware optimizations are hardcoded to GetFieldFunc

3 participants