feat: UDAFs with multiple/variadic args #9361

reneesoika · 2022-07-28T17:13:25Z

Description

This is the main PR to add support for multiple UDAF arguments. One argument can be variadic, either through wrapping the last column argument type in VariadicArgs or making the UDAF factory method have a variadic argument.

There are other changes that will be made as separate, smaller PRs:

support for up to ~~ten~~ five column arguments. This PR only adds support for two and three column arguments (though adding the rest is trivial).
correlation
var args when there are initial arguments. Currently, this is disallowed because we do not support resolving functions with var args in the middle.
variadic TOPK
user documentation

This PR should not introduce breaking changes.

The most important changes are in ExpressionTypeManager, AggregateNode, UdafUtil, UdafAggregateFunctionFactory (resolving multi-param UDAFs), UdafTypes, UdafFactoryInvoker (handling multi-param UDAF signatures), and KudafAggregator (applying multiple columns). Most of the other changes are tweaks to use lists for column indices or arrays for param schemas.

Testing done

Added unit tests to existing classes that were tweaked. Added QTTs to cover typical and edge cases with multi-param UDAFs.

Reviewer checklist

Ensure docs are updated if necessary. (eg. if a user visible feature is being added or changed).
Ensure relevant issues are linked (description should include text like "Fixes #")

…column ref

This reverts commit 392f618.

…of per row

… for unused indices

…ove old single index methods

…type passed for variadic args

… agg functions

…ppression issue

…aram list)

jzaralim

first pass at reviewing. i havent reviewed the unit tests. this looks great so far!

ksqldb-engine/src/main/java/io/confluent/ksql/function/UdafTypes.java

ksqldb-engine/src/main/java/io/confluent/ksql/function/UdafFactoryInvoker.java

jzaralim · 2022-07-29T18:39:35Z

ksqldb-common/src/main/java/io/confluent/ksql/util/Triple.java

+    this.right = right;
+  }
+
+  public T1 getLeft() {


The KLIP only mentions 10 because the javatuples library goes up to 10. If we're implementing our own classes, I don't think we really need to go up to 10. How were you planning to implement Quartet+?

I created classes for each tuple type (Quadruple, Quintuple, ...), so I was going to add these as cases in BaseAggregateFunction#determineInputConverter(). Ten column arguments does seem a bit excessive, though. Maybe stop at five?

Five is reasonable. My question was more, were you planning on following the same getLeft/getRight pattern?

For Quadruple and later, I've been using getFirst(), getSecond(), getThird(), and so on since getLeft() and similar didn't fit. I used getLeft() and getRight() for Triple since Pair uses these terms.

jzaralim · 2022-07-29T19:10:35Z

ksqldb-common/src/main/java/io/confluent/ksql/function/AggregateFunctionFactory.java

@@ -73,8 +73,7 @@ public AggregateFunctionFactory(final UdfMetadata metadata) {
    this.metadata = Objects.requireNonNull(metadata, "metadata can't be null");
  }

-  public abstract KsqlAggregateFunction<?, ?, ?> createAggregateFunction(
-      List<SqlArgument> argTypeList, AggregateFunctionInitArguments initArgs);
+  public abstract FunctionSource getFunction(List<SqlType> argTypeList);


This makes it impossible to write aggregate functions with time unit initial args.

Zara and I discussed this on Slack, and this is not a regression. In the old implementation and the current implementation, initial args must be Literals, but time units are not Literals (IntervalUnit extends Expression instead of Literal). This might be reported as a separate issue.

ksqldb-execution/src/main/java/io/confluent/ksql/execution/function/UdafUtil.java

...ecution/src/test/java/io/confluent/ksql/execution/function/udaf/KudafUndoAggregatorTest.java

jzaralim

reneesoika added 30 commits July 15, 2022 14:56

feat: dummy correlation function working

133ae55

refactor: remove unnecessary arg getter, divide initial args at last …

e26d819

…column ref

fix: incorrect # of init args for dummy function

02097e8

Revert "fix: incorrect # of init args for dummy function"

565a7e2

This reverts commit 392f618.

refactor: allow input conversion to be decided ahead of time instead …

ae04a50

…of per row

refactor: determine input converter ahead of time for correlation

183c164

refactor: move dummy init args creation to own method, use empty list…

bff5bea

… for unused indices

feat/refactor: add support for multi args to table agg functions, rem…

2858518

…ove old single index methods

feat: add pairs and triples for multi arg agg functions

b31a5da

feat: add initial support for variadic args

022a58d

feat: support multiple + variadic args for annotated UDAFs

0866080

fix: non-configurable constructor only supporting one column index

bc0023c

fix: checkstyle + exposed API errors

03b5386

fix: cols inside expressions disallowed by ExpressionTypeManager

0b73e51

fix: incorrect # of reg args in ExpressionTypeManager

cae8b62

fix: updated param name in test

9354926

fix: missing param info in test

377f615

fix: first param should always be considered a column ref

75a3467

refactor: migrate correlation test UDAF to annotations

1c2a8d6

fix: incorrect type retrieval for tuples and varargs

488403a

feat: make varargs iterable + add stream method

6df720b

fix: factory invokers not setting variadic correctly, improper param …

8504b67

…type passed for variadic args

fix: literals in col positions besides first now handled correctly by…

dcf9be8

… agg functions

fix: incorrect schema used to resolve function args in AggregateNode

9f8707f

fix: null type not handled correctly

c136d17

fix: test compile

c04ef00

fix: checkstyle

fceaa8d

fix: existing tests

ad4a8e4

fix: missing checks for variadic init args

696b109

refactor: inline extra variable

829b7e2

reneesoika added 9 commits July 27, 2022 16:13

cleanup: remove correlation test code

95586e9

test: add QTTs for generic variadic UDAF

4d8001b

fix: test not actually testing UDAF type checking

84775d0

fix: mishandled col args when init args are variadic

490aa42

test: multi arg UDAF QTT tests

2c5938c

test: non-generic var arg UDAF QTT tests

88a2548

refactor: use wrapper class instead of very long generic type

3775972

fix: checkstyle

900672b

test: add historical plans

21b5ca7

reneesoika requested a review from jzaralim July 28, 2022 17:13

reneesoika requested a review from a team as a code owner July 28, 2022 17:13

reneesoika added 5 commits July 28, 2022 13:42

fix: extraneous spotbugs warnings

33740dd

fix: suppress spotbugs at class level to try to work around lambda su…

e90e65c

…ppression issue

fix: incorrectly mocked variadic invoker (missing literal params in p…

663c176

…aram list)

fix: incorrect expected schemas after fxn signatures changed

0fc10f9

fix: outdated generated param name in test

8d416ed

jzaralim reviewed Jul 29, 2022

View reviewed changes

...ecution/src/test/java/io/confluent/ksql/execution/function/udaf/KudafUndoAggregatorTest.java Show resolved Hide resolved

reneesoika added 2 commits July 29, 2022 22:17

refactor: address non-test code review comments

045a486

test: add varargs tests to KudafAggregator and KudafUndoAggregator

2c6172c

jzaralim approved these changes Aug 2, 2022

View reviewed changes

reneesoika merged commit 02899ff into confluentinc:master Aug 2, 2022

reneesoika deleted the feat_multiple_udaf_args branch August 2, 2022 18:46

This was referenced Aug 2, 2022

feat: Add support for four and five column arguments to UDAFs #9378

Merged

feat: add CORRELATION UDAF #9383

Merged

reneesoika mentioned this pull request Aug 16, 2022

feat: Support variadic arguments in the middle of a function signature #9451

Merged

2 tasks

jzaralim mentioned this pull request Aug 18, 2022

UDAFs cannot be created with non-primitive initial arguments #9476

Open

This was referenced Aug 22, 2022

feat: Support Object as the type parameter for a UDAF variadic column argument #9481

Merged

feat: Variadic TopK that can select other columns #9493

Merged

docs: Add user docs for multi-arg and variadic UDAFs #9494

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: UDAFs with multiple/variadic args #9361

feat: UDAFs with multiple/variadic args #9361

reneesoika commented Jul 28, 2022 •

edited

jzaralim left a comment

jzaralim Jul 29, 2022

reneesoika Jul 29, 2022

jzaralim Jul 29, 2022

reneesoika Jul 29, 2022

jzaralim Jul 29, 2022

reneesoika Jul 29, 2022

jzaralim left a comment

feat: UDAFs with multiple/variadic args #9361

feat: UDAFs with multiple/variadic args #9361

Conversation

reneesoika commented Jul 28, 2022 • edited

Description

Testing done

Reviewer checklist

jzaralim left a comment

Choose a reason for hiding this comment

jzaralim Jul 29, 2022

Choose a reason for hiding this comment

reneesoika Jul 29, 2022

Choose a reason for hiding this comment

jzaralim Jul 29, 2022

Choose a reason for hiding this comment

reneesoika Jul 29, 2022

Choose a reason for hiding this comment

jzaralim Jul 29, 2022

Choose a reason for hiding this comment

reneesoika Jul 29, 2022

Choose a reason for hiding this comment

jzaralim left a comment

Choose a reason for hiding this comment

reneesoika commented Jul 28, 2022 •

edited