feat: Extend the simple UDAF interface with function-level variables #12067

liujiayi771 · 2025-01-12T11:11:35Z

Some aggregate functions implemented based on the simple UDAF interface need to obtain information such as step, argTypes, and resultType in the AccumulatorType struct. This PR adds the capability for user-defined variables to the simple UDAF interface, supporting the passing of step, argTypes, and resultType to AccumulatorType struct when registering the SimpleAggregateAdapter.

The UDAF author add function-level variables as data members in the UDAF class (outside of its AccumulatorType struct).
The UDAF class has an initialize() method that receives the aggregation step, the types, and assigns values to the data members in the UDAF class.
The AccumulatorType struct has a data member that is a pointer to the UDAF class. This would allow member methods inside AccumulatorType to access data members in the UDAF class. This UDAF-pointer is set in SimpleAggregateAdapter::initializeNewGroupsInternal().

netlify · 2025-01-12T11:12:00Z

✅ Deploy Preview for meta-velox canceled.

Name	Link
🔨 Latest commit	`203160a`
🔍 Latest deploy log	https://app.netlify.com/sites/meta-velox/deploys/67a2ae2b6d64e60008078291

liujiayi771 · 2025-01-12T11:19:31Z

@kagamiori Some time ago, we attempted to add function-level states to the simple UDAF interface in #9167 and we have discussed some details in #8711, but @mbasmanova suggested that we implement it in a manner similar to the simple scalar functions interface. You had previously provided a specific implementation plan, and this PR is implemented according to the plan you proposed at that time. However, I did not add support for constantInput, as we currently only need to use rawInputType and resultType. Adding constantInput would require introducing new variables in the AggregateFunctionFactory, and I believe we can expand it when we actually need it.

We currently need this capability to implement the decimal average aggregate function in Spark. Could you help to take a look. cc @rui-mo.

kagamiori · 2025-01-16T03:45:13Z

@kagamiori Some time ago, we attempted to add function-level states to the simple UDAF interface in #9167 and we have discussed some details in #8711, but @mbasmanova suggested that we implement it in a manner similar to the simple scalar functions interface. You had previously provided a specific implementation plan, and this PR is implemented according to the plan you proposed at that time. However, I did not add support for constantInput, as we currently only need to use rawInputType and resultType. Adding constantInput would require introducing new variables in the AggregateFunctionFactory, and I believe we can expand it when we actually need it.

We currently need this capability to implement the decimal average aggregate function in Spark. Could you help to take a look. cc @rui-mo.

Sure, let me take a look.

This was the second prototype after getting @mbasmanova's suggestions: #9775. It looks like this PR has the same design 👍.

kagamiori

Hi @liujiayi771, this looks mostly good to me except a few nits. Thank you for extending the simple UDAF interface to improve its usability!

kagamiori · 2025-01-18T00:44:41Z

velox/exec/SimpleAggregateAdapter.h

+      TypePtr resultType)
+      : Aggregate(std::move(resultType)), fn_{std::make_unique<FUNC>()} {
+    if constexpr (support_initialize_) {
+      FUNC::initialize(fn_.get(), step, argTypes, resultType_);


nit: Just fn_->initialize(step, argTypes, resultType_);

kagamiori · 2025-01-18T00:45:51Z

velox/exec/SimpleAggregateAdapter.h

@@ -103,6 +110,7 @@ class SimpleAggregateAdapter : public Aggregate {
  // These functions are called on groups of both non-null and null
  // accumulators. These functions also return a bool indicating whether the
  // current group should be a NULL in the result vector.
+  std::unique_ptr<FUNC> fn_;


nit: Move the member variable to the bottom of the class.

kagamiori · 2025-01-18T01:18:00Z

velox/exec/tests/SimpleAggregateAdapterTest.cpp

+  using IntermediateType = Row<int64_t, double>;
+  using OutputType = double;
+
+  TypePtr inputType_;


nit: Add a comment above to say explicitly that these are the function level variable we're going to test, and explain that they are set during the creation of the aggregation function and checked in addInput().

kagamiori · 2025-01-18T01:30:23Z

velox/functions/prestosql/aggregates/ApproxMostFrequentAggregate.cpp

            resultType,
            name,
-            valueType);
+            valueType,
+            step,
+            argTypes);


nit: reorder the arguments to be name, step, argType, resultType, valueType to be more consistent with the lambda's parameter list.

kagamiori · 2025-01-18T01:32:54Z

velox/functions/sparksql/aggregates/SumAggregate.cpp

    const TypePtr& inputType,
    const TypePtr& sumType,
-    const TypePtr& resultType) {
+    const TypePtr& resultType,
+    core::AggregationNode::Step step,
+    const std::vector<TypePtr>& argTypes) {


nit: Reorder the parameter list to be step, argTypes, resultType, inputType, sumType.

kagamiori · 2025-01-18T01:34:29Z

velox/docs/develop/aggregate-functions.rst

+    TypePtr inputType;
+    TypePtr resultType;


nit: inputType -> inputType_, resultType -> resultType_

liujiayi771 · 2025-01-22T01:17:39Z

@kagamiori Thank you for your review. All the comments have been addressed.

kagamiori

LGTM. Thank you for making the simple UDAF interface support function-level states! I left a few nits.

kagamiori · 2025-01-29T19:33:20Z

velox/docs/develop/aggregate-functions.rst

@@ -248,6 +271,9 @@ For aggregaiton functions of default-null behavior, the author defines an
    // Author defines data members
    ...

+    // Define a pointer to the UDAF class.


nit: Add in the comment that this is only needed if the aggregation function uses function-level states.

kagamiori · 2025-01-29T19:37:20Z

velox/docs/develop/aggregate-functions.rst

+Some function-level variables needs to be declared in the simple aggregation
+function class. These variables are initialized once when the aggregation


nit: The author can optionally define function-level variables in the simple aggregation function class. These variables...

kagamiori · 2025-01-29T19:39:02Z

velox/docs/develop/aggregate-functions.rst

+function, the author can hold them in the aggregate class variables, and
+initialize them in the initialize() method.


nit: the author can hold them in the aggregate class variables, and
initialize them in the initialize() method. --> these types can be defined as member variables in the aggregate class and initialized in the initialize() method.

facebook-github-bot · 2025-01-29T19:41:58Z

@kagamiori has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

facebook-github-bot · 2025-02-04T18:09:31Z

@kagamiori has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

kagamiori · 2025-02-04T23:26:12Z

Hi @liujiayi771, could you please rebase the PR onto the latest main? Thanks!

liujiayi771 · 2025-02-05T01:09:26Z

Hi @liujiayi771, could you please rebase the PR onto the latest main? Thanks!

Done.

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jan 12, 2025

liujiayi771 force-pushed the func-level-variables branch from 9013ee9 to 9775661 Compare January 12, 2025 11:19

liujiayi771 force-pushed the func-level-variables branch 4 times, most recently from 94af18b to 8d84155 Compare January 13, 2025 05:49

kagamiori reviewed Jan 18, 2025

View reviewed changes

kagamiori mentioned this pull request Jan 21, 2025

Prototype of simple UDAF interface with function-level variables #9775

Closed

kagamiori approved these changes Jan 29, 2025

View reviewed changes

liujiayi771 added 3 commits February 5, 2025 08:17

extend the simple UDAF interface with function-level variables

056d639

address comments

c466df6

Address comments

203160a

liujiayi771 force-pushed the func-level-variables branch from 7a37b60 to 203160a Compare February 5, 2025 00:17

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Extend the simple UDAF interface with function-level variables #12067

feat: Extend the simple UDAF interface with function-level variables #12067

liujiayi771 commented Jan 12, 2025 •

edited

Loading

netlify bot commented Jan 12, 2025 •

edited

Loading

liujiayi771 commented Jan 12, 2025 •

edited

Loading

kagamiori commented Jan 16, 2025 •

edited

Loading

kagamiori left a comment

kagamiori Jan 18, 2025

kagamiori Jan 18, 2025

kagamiori Jan 18, 2025

kagamiori Jan 18, 2025

kagamiori Jan 18, 2025

kagamiori Jan 18, 2025

liujiayi771 commented Jan 22, 2025

kagamiori left a comment

kagamiori Jan 29, 2025

kagamiori Jan 29, 2025

kagamiori Jan 29, 2025

facebook-github-bot commented Jan 29, 2025

facebook-github-bot commented Feb 4, 2025

kagamiori commented Feb 4, 2025

liujiayi771 commented Feb 5, 2025

		Some function-level variables needs to be declared in the simple aggregation
		function class. These variables are initialized once when the aggregation

		function, the author can hold them in the aggregate class variables, and
		initialize them in the initialize() method.

feat: Extend the simple UDAF interface with function-level variables #12067

Are you sure you want to change the base?

feat: Extend the simple UDAF interface with function-level variables #12067

Conversation

liujiayi771 commented Jan 12, 2025 • edited Loading

netlify bot commented Jan 12, 2025 • edited Loading

✅ Deploy Preview for meta-velox canceled.

liujiayi771 commented Jan 12, 2025 • edited Loading

kagamiori commented Jan 16, 2025 • edited Loading

kagamiori left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

liujiayi771 commented Jan 22, 2025

kagamiori left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

facebook-github-bot commented Jan 29, 2025

facebook-github-bot commented Feb 4, 2025

kagamiori commented Feb 4, 2025

liujiayi771 commented Feb 5, 2025

liujiayi771 commented Jan 12, 2025 •

edited

Loading

netlify bot commented Jan 12, 2025 •

edited

Loading

liujiayi771 commented Jan 12, 2025 •

edited

Loading

kagamiori commented Jan 16, 2025 •

edited

Loading