Support Schema Field Metadata in User-Defined Aggregate Functions (UDAFs) #17085
+256
−56
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
github-actions
bot
added
logical-expr
- Improve documentation for AccumulatorArgs.schema and exprs: - Add example showing how to retrieve field metadata and return field. - Explain synthesized schema behavior for literal-only inputs. - Clarify precedence when inputs are mixed (physical schema metadata wins; synthesized metadata used only when physical schema is empty). - Update AggregateFunctionExpr::args_schema docs: - Explain field order guarantees, synthesized schema usage, and that std::borrow::Cow is used to avoid allocations when possible. - Add a TODO to factor AccumulatorArgs construction into a private helper. Documentation-only changes; no behavioral changes.
… in AggregateFunctionExpr
kosiew
force-pushed
the
udaf-schema-16997
branch
from
fb6c9a8 to
2991acc
Compare
- Updated the `DummyUdf::new` function to accept a `Signature` parameter for initialization. - Modified test cases to reflect the new initialization method for `DummyUdf`. - Improved the `args_schema` method to better manage when to borrow the existing schema or create a new one, ensuring correct correspondence between expressions and schema fields. - Added comments and examples for improved clarity on the schema handling behavior.
- Introduced a `build_acc_args` method to encapsulate the logic for building `AccumulatorArgs` and executing a closure with them. - Updated the `create_accumulator`, `create_sliding_accumulator`, `groups_accumulator_supported`, and `create_groups_accumulator` methods to utilize the new method for better readability and maintainability.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
AccumulatorArgs.schemais empty when passing in scalar input #16997Rationale for this change
Previously, aggregate UDFs could not easily access field-level metadata (e.g., Arrow extension types) when invoked with literal arguments only, because the
AccumulatorArgs.schemawas always derived from the physical schema — which is empty for literal-only inputs.This change ensures that in such cases, a schema is synthesized from the literal expressions, preserving metadata and enabling richer accumulator behavior. It also clarifies API documentation for
AccumulatorArgsandAggregateUDFImpl.What changes are included in this PR?
args_schema()helper toAggregateFunctionExprto return either the physical input schema or a synthesized schema from literals when the physical schema is empty.create_accumulator,create_sliding_accumulator,groups_accumulator_supported, andcreate_groups_accumulatorto use the new schema logic viamake_acc_args().AccumulatorArgsandAggregateUDFImpldocumentation to explain how to access input field metadata and when synthesized schemas are used.SchemaBasedAggregateUdftest inuser_defined_aggregates.rsto validate metadata handling in literal-only aggregates.aggregate.rsto verify correct schema behavior for both literal-only and physical-schema-present cases.Are these changes tested?
Yes.
test_schema_based_aggregate_udf_metadataensures that metadata from literals is accessible in the accumulator.aggregate.rsvalidates thatargs_schema()returns an owned schema for literal-only inputs and a borrowed schema for non-empty physical schemas.Are there any user-facing changes?
Yes:
AccumulatorArgs.schemafor literal-only inputs.