You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "kosiew (via GitHub)" <gi...@apache.org> on 2025/08/08 10:33:04 UTC
[PR] Support Schema Field Metadata in User-Defined Aggregate Functions (UDAFs) [datafusion]
kosiew opened a new pull request, #17085:
URL: https://github.com/apache/datafusion/pull/17085
## Which issue does this PR close?
Closes #16997
## Rationale for this change
Currently, when a user-defined aggregate function (UDAF) is invoked with literal values, any associated field-level metadata (such as Arrow extension types or custom tags) is lost due to the lack of a proper schema. This PR ensures that UDAFs can access metadata for literal inputs by synthesizing a schema from those literals. This empowers UDAFs to exhibit metadata-driven behavior.
## What changes are included in this PR?
* Synthesized a schema from literal inputs inside `AggregateFunctionExpr::args_schema` to retain metadata.
* Updated `AccumulatorArgs` documentation and usage across `datafusion` to reflect and utilize the synthesized schema.
* Added a new test `test_scalar_udaf_schema_metadata` to validate metadata-driven behavior in UDAFs.
* Improved internal comments and docs to clarify the relationship between `AccumulatorArgs`, metadata, and field inference.
## Are these changes tested?
Yes 🎉
* The new functionality is validated with a dedicated test: `test_scalar_udaf_schema_metadata`, ensuring that metadata is preserved and interpreted correctly in user-defined accumulators.
## Are there any user-facing changes?
Yes, but non-breaking:
* Users writing UDAFs now have access to field metadata even when arguments are literals.
* Enhanced documentation and behavior for advanced UDAF use cases relying on metadata.
These changes are backwards compatible and additive, providing new capabilities without breaking existing functionality.
---
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Support Schema Field Metadata in User-Defined Aggregate Functions (UDAFs) [datafusion]
Posted by "Jefffrey (via GitHub)" <gi...@apache.org>.
Jefffrey commented on PR #17085:
URL: https://github.com/apache/datafusion/pull/17085#issuecomment-3379521826
> Not at all 😄
>
> * `acc_args.schema.field(i)` — returns the raw Arrow `Field` from the (physical) input schema at position `i` (name, type, nullability, metadata exactly as in that schema).
>
> * `acc_args.exprs[i].return_field(&acc_args.schema)?` — asks the expression for the effective `FieldRef` for argument `i` given the full schema. It incorporates expression semantics (casts, literals, computed types, extension metadata, nullability changes, etc.) and returns an owned/Arc `FieldRef` (and can fail), not just a borrowed `&Field`.
Thank you for the explanations, I'm still trying to wrap my head around all the parts involved here 😅
So I think my main confusion lies around the difference between the physical input schema, and the effective `FieldRef` argument; is there a reason we provide the ability to access both? This fix only synthesizes a schema if the physical schema is missing as you mentioned, but would it be incorrect behaviour to instead _always_ synthesize the schema from the physical schema (whether present or not)?
If we look at scalar & window functions, I don't see them having equivalent logic around providing direct access to the physical schema, instead they provide methods to directly access `Field`s:
- https://docs.rs/datafusion/50.0.0/datafusion/logical_expr/function/struct.ExpressionArgs.html
- https://docs.rs/datafusion/50.0.0/datafusion/logical_expr/function/struct.PartitionEvaluatorArgs.html
- https://docs.rs/datafusion/50.0.0/datafusion/logical_expr/struct.ScalarFunctionArgs.html
I'm trying to understand why `AccumulatorArgs` seems to be the odd one out here; I'm sure there's some historical reason but the limited existing documentation on `AccumulatorArgs` makes it hard for me to reason that this fix is the correct approach 🤔
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Support Schema Field Metadata in User-Defined Aggregate Functions (UDAFs) [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #17085:
URL: https://github.com/apache/datafusion/pull/17085#discussion_r2394701655
##########
datafusion/expr/src/udaf.rs:
##########
@@ -762,7 +762,24 @@ pub trait AggregateUDFImpl: Debug + DynEq + DynHash + Send + Sync {
/// group during query execution.
///
/// acc_args: [`AccumulatorArgs`] contains information about how the
- /// aggregate function was called.
+ /// aggregate function was called. Use `acc_args.exprs` together with
+ /// `acc_args.schema` to inspect the [`FieldRef`] of each input.
+ ///
+ /// Example: retrieving metadata and return field for input `i`:
+ /// ```ignore
+ /// let metadata = acc_args.schema.field(i).metadata();
+ /// let field = acc_args.exprs[i].return_field(&acc_args.schema)?;
+ /// ```
Review Comment:
Not at all 😄
* `acc_args.schema.field(i)` — returns the raw Arrow `Field` from the (physical) input schema at position `i` (name, type, nullability, metadata exactly as in that schema).
* `acc_args.exprs[i].return_field(&acc_args.schema)?` — asks the expression for the effective `FieldRef` for argument `i` given the full schema. It incorporates expression semantics (casts, literals, computed types, extension metadata, nullability changes, etc.) and returns an owned/Arc `FieldRef` (and can fail), not just a borrowed `&Field`.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Support Schema Field Metadata in User-Defined Aggregate Functions (UDAFs) [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #17085:
URL: https://github.com/apache/datafusion/pull/17085#discussion_r2394519916
##########
datafusion/expr/src/udaf.rs:
##########
@@ -762,7 +762,24 @@ pub trait AggregateUDFImpl: Debug + DynEq + DynHash + Send + Sync {
/// group during query execution.
///
/// acc_args: [`AccumulatorArgs`] contains information about how the
- /// aggregate function was called.
+ /// aggregate function was called. Use `acc_args.exprs` together with
+ /// `acc_args.schema` to inspect the [`FieldRef`] of each input.
+ ///
+ /// Example: retrieving metadata and return field for input `i`:
+ /// ```ignore
+ /// let metadata = acc_args.schema.field(i).metadata();
+ /// let field = acc_args.exprs[i].return_field(&acc_args.schema)?;
+ /// ```
Review Comment:
The snippet is meant to illustrate the sentence immediately above it: you pair `acc_args.exprs` with `acc_args.schema` to recover the full `FieldRef` for argument *i*.
Pulling the metadata out of `schema.field(i)` is one common use case, and the follow-up line shows how you would then obtain the complete `FieldRef` (name, type, metadata) via:
```rust
exprs[i].return_field(&acc_args.schema)
```
...using the same pairing.
I'll tweak the wording.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Support Schema Field Metadata in User-Defined Aggregate Functions (UDAFs) [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on PR #17085:
URL: https://github.com/apache/datafusion/pull/17085#issuecomment-3413735133
Closing this.
#18100 has a [better approach](https://github.com/apache/datafusion/pull/18100#issuecomment-3413462730).
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Support Schema Field Metadata in User-Defined Aggregate Functions (UDAFs) [datafusion]
Posted by "Jefffrey (via GitHub)" <gi...@apache.org>.
Jefffrey commented on PR #17085:
URL: https://github.com/apache/datafusion/pull/17085#issuecomment-3389773308
Sorry this doesn't clear it up for me; the example of `SUM(a + b)` needing the physical schema because it must resolve both `a` and `b` doesn't make sense to me given scalar and window functions can similarly accept multiple columns like that (e.g. `SIN(a + b)`). I'll see if I can dig deeper into the related code to try bring my own understanding into this 🤔
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Support Schema Field Metadata in User-Defined Aggregate Functions (UDAFs) [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on PR #17085:
URL: https://github.com/apache/datafusion/pull/17085#issuecomment-3356523033
> Is this a band-aid fix? Is there a root cause we should be looking for instead?
> There's a heavy emphasis on the word "synthesize" throughout this PR but I don't know what it means to "synthesize" a schema from literal expressions 🤔
AggregateExprBuilder already captures a FieldRef for every argument (including literals) by calling each physical expression’s return_field during construction, so we retain the full Arrow metadata for those inputs in input_fields.
The new args_schema helper detects when the physical input schema is empty—something that legitimately happens when an aggregate is invoked with literals only because the child plan has no columns—and in that case reconstitutes a Schema from the stored input_fields so the accumulator can still see that metadata.
We then hand that schema to every AccumulatorArgs we build, so UDAFs observe the same field information whether their inputs were columns or literals. In other words, “synthesize” means “wrap the already-computed argument fields in a temporary Schema when the physical schema is empty”; there isn’t another layer hiding the real root cause.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Support Schema Field Metadata in User-Defined Aggregate Functions (UDAFs) [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #17085:
URL: https://github.com/apache/datafusion/pull/17085#discussion_r2394741289
##########
datafusion/physical-expr/src/aggregate.rs:
##########
@@ -24,6 +24,92 @@ pub(crate) mod groups_accumulator {
accumulate::NullState, GroupsAccumulatorAdapter,
};
}
+
+#[cfg(test)]
Review Comment:
I will move the tests.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Support Schema Field Metadata in User-Defined Aggregate Functions (UDAFs) [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew closed pull request #17085: Support Schema Field Metadata in User-Defined Aggregate Functions (UDAFs)
URL: https://github.com/apache/datafusion/pull/17085
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Support Schema Field Metadata in User-Defined Aggregate Functions (UDAFs) [datafusion]
Posted by "Jefffrey (via GitHub)" <gi...@apache.org>.
Jefffrey commented on code in PR #17085:
URL: https://github.com/apache/datafusion/pull/17085#discussion_r2390273336
##########
datafusion/expr/src/udaf.rs:
##########
@@ -762,7 +762,24 @@ pub trait AggregateUDFImpl: Debug + DynEq + DynHash + Send + Sync {
/// group during query execution.
///
/// acc_args: [`AccumulatorArgs`] contains information about how the
- /// aggregate function was called.
+ /// aggregate function was called. Use `acc_args.exprs` together with
+ /// `acc_args.schema` to inspect the [`FieldRef`] of each input.
+ ///
+ /// Example: retrieving metadata and return field for input `i`:
+ /// ```ignore
+ /// let metadata = acc_args.schema.field(i).metadata();
+ /// let field = acc_args.exprs[i].return_field(&acc_args.schema)?;
+ /// ```
Review Comment:
I'm having some trouble understanding this example; I can understand the part for getting the metadata of a field given the context of the PR, but why do we also include an example for getting the return field?
##########
datafusion/physical-expr/src/aggregate.rs:
##########
@@ -24,6 +24,92 @@ pub(crate) mod groups_accumulator {
accumulate::NullState, GroupsAccumulatorAdapter,
};
}
+
+#[cfg(test)]
+mod tests {
Review Comment:
Tests should be at the bottom of the file
##########
datafusion/physical-expr/src/aggregate.rs:
##########
@@ -24,6 +24,92 @@ pub(crate) mod groups_accumulator {
accumulate::NullState, GroupsAccumulatorAdapter,
};
}
+
+#[cfg(test)]
Review Comment:
Tests should be moved to the bottom of the file
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Support Schema Field Metadata in User-Defined Aggregate Functions (UDAFs) [datafusion]
Posted by "Jefffrey (via GitHub)" <gi...@apache.org>.
Jefffrey commented on code in PR #17085:
URL: https://github.com/apache/datafusion/pull/17085#discussion_r2394652327
##########
datafusion/expr/src/udaf.rs:
##########
@@ -762,7 +762,24 @@ pub trait AggregateUDFImpl: Debug + DynEq + DynHash + Send + Sync {
/// group during query execution.
///
/// acc_args: [`AccumulatorArgs`] contains information about how the
- /// aggregate function was called.
+ /// aggregate function was called. Use `acc_args.exprs` together with
+ /// `acc_args.schema` to inspect the [`FieldRef`] of each input.
+ ///
+ /// Example: retrieving metadata and return field for input `i`:
+ /// ```ignore
+ /// let metadata = acc_args.schema.field(i).metadata();
+ /// let field = acc_args.exprs[i].return_field(&acc_args.schema)?;
+ /// ```
Review Comment:
> The snippet is meant to illustrate the sentence immediately above it: you pair `acc_args.exprs` with `acc_args.schema` to recover the full `FieldRef` for argument _i_.
This may be a silly question, but what's the difference between `acc_args.exprs[i].return_field(&acc_args.schema)` and `acc_args.schema.field(i)`?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Support Schema Field Metadata in User-Defined Aggregate Functions (UDAFs) [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on PR #17085:
URL: https://github.com/apache/datafusion/pull/17085#issuecomment-3389122318
<html><head></head><body><h1>Why <code inline="">AccumulatorArgs</code> Differs from Other Function Args</h1>
<p>Your review raises an excellent question about API consistency. The answer lies in <strong>when and how</strong> these functions resolve their input fields.</p>
<hr>
<h2>Scalar & Window Functions: Pre-computed Fields</h2>
<p>Both scalar and window functions receive <strong>pre-computed</strong> <code inline="">FieldRef</code>s at creation time:</p>
<pre><code class="language-rust">// From create_udwf_window_expr (windows/mod.rs:178-180)
let input_fields: Vec<_> = args
.iter()
.map(|arg| arg.return_field(input_schema)) // Computed once at planning
.collect::<Result<_>>()?;
</code></pre>
<p>These are then passed to <code inline="">ExpressionArgs</code>, <code inline="">PartitionEvaluatorArgs</code>, or <code inline="">ScalarFunctionArgs</code>.<br>
The schema itself is never exposed because <strong>the fields have already been computed</strong>.</p>
<hr>
<h2>Aggregate Functions: Runtime Schema Access</h2>
<p>Aggregates, by contrast, receive the <strong>physical schema</strong> and compute fields <strong>on demand</strong>:</p>
<pre><code class="language-rust">// From AggregateFunctionExpr::create_accumulator (aggregate.rs:429-432)
let schema = self.args_schema();
let acc_args = self.make_acc_args(schema.as_ref()); // Schema passed here
self.fun.accumulator(acc_args)
</code></pre>
<hr>
<h1>Why Can't Aggregates Use Pre-computed Fields?</h1>
<p>The expressions need <strong>schema access at runtime</strong>. Consider:</p>
<pre><code>SUM(a + b)
</code></pre>
<p>The expression <code inline="">a + b</code> references <strong>two columns</strong> from the physical schema, but produces <strong>one logical argument</strong> for the accumulator.<br>
The <code inline="">PhysicalExpr</code> for <code inline="">a + b</code> needs to:</p>
<ol>
<li>
<p>Resolve columns <code inline="">a</code> and <code inline="">b</code> from the physical schema</p>
</li>
<li>
<p>Evaluate the addition</p>
</li>
<li>
<p>Pass the result to the accumulator</p>
</li>
</ol>
<p>If we only passed pre-computed fields (like scalar/window functions do), the <code inline="">PhysicalExpr</code> couldn’t resolve its column references — it needs the full <code inline="">Schema</code>.</p>
<hr>
<h3>From the patch (aggregate.rs:391-395):</h3>
<blockquote>
<p>Physical expressions may reference columns that are not one-to-one with the aggregate arguments (for example, <code inline="">SUM(a + b)</code> references both <code inline="">a</code> and <code inline="">b</code>). Replacing the schema with the synthesized argument fields would prevent those expressions from resolving their column references.</p>
</blockquote>
<hr>
<h2>What the Patch Actually Changes</h2>
<p>The patch <strong>doesn't change runtime behavior</strong> — it:</p>
<ol>
<li>
<p><strong>Clarifies documentation:</strong> Makes explicit that <code inline="">schema</code> is the physical input schema</p>
</li>
<li>
<p><strong>Adds convenience methods:</strong> <code inline="">input_field()</code> and <code inline="">input_fields()</code> wrap the common pattern of calling <code inline="">expr.return_field(&schema)</code></p>
</li>
<li>
<p><strong>Aligns ergonomics:</strong> Matches the convenience of <code inline="">ScalarFunctionArgs.arg_fields</code> and <code inline="">PartitionEvaluatorArgs.input_fields()</code>, without losing the schema access that aggregate expressions require</p>
</li>
</ol>
<hr>
<h1>Comparison Table</h1>
Aspect | Scalar / Window | Accumulator (Before) | Accumulator (After)
-- | -- | -- | --
Fields | Pre-computed at planning | Computed on-demand | Computed on-demand
Schema | Hidden (not needed) | Exposed (confusing) | Exposed (documented)
Field access | Direct: args.arg_fields[i] | Manual: exprs[i].return_field(schema)? | Helper: input_field(i)?
Why different? | Simple 1:1 args | Expressions need schema for column resolution | Same (clarified)
<hr>
<h2>Why This Fix Is Correct</h2>
<p>The patch acknowledges that <code inline="">AccumulatorArgs</code> <strong>must remain unique</strong> because aggregates operate at the physical execution layer with complex expressions.<br>
Rather than force-fitting them into the scalar/window model (which would break column resolution), it:</p>
<ul>
<li>
<p>Preserves the necessary schema access</p>
</li>
<li>
<p>Documents <strong>why</strong> it’s different</p>
</li>
<li>
<p>Adds ergonomic helpers for the common case</p>
</li>
</ul>
<hr>
<p><strong>This is the right approach:</strong><br>
Clarify intent and align ergonomics where possible, while preserving the functionality that makes aggregates work correctly.</p></body></html>
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org