refactor: unify aggregate and window functions in window handling #170

vbarua · 2023-08-28T20:18:06Z

feat: allow window function operands beyond just column references
fix: invocation was not set when building WindowFunction proto

BREAKING CHANGES:

Window and WindowFunction have been merged into WindowFunctionInvocation
WindowFunctionInvocation is now an inner class of expression
WindowFunctionInvocation now implements Expression.
Expression visitor now visit WindowFunctionInvocation
WindowFunctionConvert constructor requires fewer arguments
windowCreator has additional fields
POJO WindowBound offset must now be a long, better reflecting spec

vbarua · 2023-08-28T20:21:25Z

core/src/main/java/io/substrait/expression/Expression.java

-  public abstract static class WindowFunction {
-    public abstract WindowFunctionInvocation getFunction();
-
-    public abstract Optional<Expression> getPreMeasureFilter();


The preMeasureFilter field was not used anywhere, and the WindowFunction message cannot actually encode a pre-measure filter (for now).

I opted to remove it entirely. Support can be added in once the spec supports it.

vbarua · 2023-08-28T20:23:06Z

core/src/main/java/io/substrait/expression/WindowBound.java

@@ -45,7 +45,7 @@ public BoundedKind boundedKind() {

    public abstract Direction direction();

-    public abstract Expression offset();
+    public abstract long offset();


In the spec, offset must be int64.

The POJO layer should reflect this restriction.

vbarua · 2023-08-28T20:25:08Z

core/src/main/java/io/substrait/expression/WindowFunctionInvocation.java

-import org.immutables.value.Value;
-
-@Value.Immutable
-public abstract class WindowFunctionInvocation {


WindowFunctionInvocation is a / implements Expression, much like ScalarFunctionInvocation.

Moved this to an inner class of Expression for consistency.

vbarua · 2023-08-28T20:37:43Z

core/src/main/java/io/substrait/extension/SimpleExtension.java

+
+    List<WindowFunctionVariant> allWindowFunctionVariants =
+        Stream.concat(windowFunctionVariants, windowAggFunctionVariants)
+            .collect(Collectors.toList());


The spec is being loosed to allow usage of aggregate functions as window function.
substrait-io/substrait#540

substrait-java already supported this behaviour, but the implementation is simpler/cleaner if aggregate functions are added as valid window functions here.

vbarua · 2023-08-28T20:51:21Z

core/src/main/java/io/substrait/expression/Expression.java

-    public abstract Aggregate.Measure aggregateFunction();
-
-    @Nullable
-    public abstract WindowFunction windowFunction();


Instead of having separate handling for window calls with window functions, and window calls with aggregate functions, this PR updates the extension handling to treat all aggregate functions as valid window functions and a such removes the need for this distinction.

vbarua · 2023-08-28T20:52:02Z

isthmus/src/main/java/io/substrait/isthmus/expression/RexExpressionConverter.java

 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.stream.Collectors;
-import org.apache.calcite.rel.RelNode;
-import org.apache.calcite.rex.*;


Avoiding * imports

vbarua · 2023-08-28T20:54:37Z

isthmus/src/main/java/io/substrait/isthmus/expression/WindowFunctionConverter.java

-  public Expression.Window convert(
-      RelNode input,
-      Type.Struct inputType,
+  public Optional<Expression.WindowFunctionInvocation> convert(


public Optional<Expression.WindowFunctionInvocation> convert

This output types brings convert inline with both ScalarFunctionConverter and AggregateFunctionConverter.

vbarua · 2023-08-28T21:00:02Z

isthmus/src/test/java/io/substrait/isthmus/PlanTestBase.java

+   * </ul>
+   */
+  @Beta
+  protected void assertFullRoundTrip(String sqlQuery, List<String> createStatements)


assertProtoPlanRoundrip first converts from

SQL -> Calcite -> POJO -> Proto

and then starting from the proto does

Proto -> Pojo -> Proto

but this means that if the initial Proto -> Pojo translation is missing fields, the test doesn't detect them.

This assert helped identify some issues when processing window functions. My intent is:

make this assert strong by checking more than just Pojo -> Proto -> Pojo

replace all calls of assertProtoPlanRoundrip with

Marked as BETA for now as 1 is not done yet. I thought about doing it in this PR, which would requires adding support for mapping window functions from substrait into Calcite, but this is already big enough as is.

vbarua · 2023-08-28T21:02:26Z

isthmus/src/main/java/io/substrait/isthmus/expression/WindowFunctionConverter.java

-            null,
-            RelCollations.EMPTY,
-            over.getType(),
-            sqlAggFunction.getName());


The usage of this API meant that calls like ntile(4) were not supported, as 4 would be encoded as a RexLiteral which is not a RexSlot. Switching away from this handling enables those queries, and in fact I was able to enable the test for the ntile query in WindowFunctionTest.

EpsilonPrime

Good set of changes. I especially like how bounds and operands are much cleaner (as opposed to invoking a potentially unrelated builder in place).

BREAKING CHANGE * WindowFunctionConverter no longer takes an AggregateFunctionConverter

feat: allow window function operands beyond just column references fix: invocation was not set when building WindowFunction proto message BREAKING CHANGE: * Window and WindowFunction have been merged into WindowFunctionInvocation * WindowFunctionInvocation is now an inner class of expression * WindowFunctionInvocation now implements Expression. * Expression visitor now visits WindowFunctionInvocation * WindowFunctionConvert constructor requires fewer arguments * windowCreator has additional fields * POJO WindowBound offset must now be a long, better reflecting spec

…bstrait-io#170) feat: allow window function operands beyond just column references fix: invocation was not set when building WindowFunction proto message BREAKING CHANGES: * Window and WindowFunction have been merged into WindowFunctionInvocation * WindowFunctionInvocation is now an inner class of expression * WindowFunctionInvocation now implements Expression. * Expression visitor now visits WindowFunctionInvocation * WindowFunctionConvert constructor requires fewer arguments * windowCreator has additional fields * POJO WindowBound offset must now be a long, better reflecting spec

vbarua commented Aug 28, 2023

View reviewed changes

vbarua force-pushed the vbarua/agg-fns-in-windows branch 2 times, most recently from 03a9e26 to 3c733be Compare August 28, 2023 20:48

vbarua commented Aug 28, 2023

View reviewed changes

EpsilonPrime approved these changes Aug 28, 2023

View reviewed changes

vbarua marked this pull request as ready for review August 28, 2023 21:19

vbarua added 8 commits August 30, 2023 10:07

feat: enable usage of aggregate functions in windows

b9d3c95

refactor: simplify window function lookup

c5149a2

BREAKING CHANGE * WindowFunctionConverter no longer takes an AggregateFunctionConverter

refactor: merge Window into WindowFunctionInvocation

f5e33f5

test: introduce assertFullRoundTrip check

4b7886d

test: stronger WindowFunctionAssert

6f8924b

fix: invocation was not set on WindowFunction

38c97a2

fix: substrait only supports integer window offsets

00d52fc

refactor: move WindowFunctionInvocation into Expression

082139c

vbarua force-pushed the vbarua/agg-fns-in-windows branch from 3c733be to 082139c Compare August 30, 2023 17:09

vbarua merged commit afa1213 into main Aug 30, 2023
8 checks passed

vbarua deleted the vbarua/agg-fns-in-windows branch August 30, 2023 18:42

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

refactor: unify aggregate and window functions in window handling #170

refactor: unify aggregate and window functions in window handling #170

vbarua commented Aug 28, 2023 •

edited

Loading

vbarua Aug 28, 2023

vbarua Aug 28, 2023

vbarua Aug 28, 2023

vbarua Aug 28, 2023

vbarua Aug 28, 2023

vbarua Aug 28, 2023

vbarua Aug 28, 2023

vbarua Aug 28, 2023 •

edited

Loading

vbarua Aug 28, 2023 •

edited

Loading

EpsilonPrime left a comment

refactor: unify aggregate and window functions in window handling #170

refactor: unify aggregate and window functions in window handling #170

Conversation

vbarua commented Aug 28, 2023 • edited Loading

vbarua Aug 28, 2023

Choose a reason for hiding this comment

vbarua Aug 28, 2023

Choose a reason for hiding this comment

vbarua Aug 28, 2023

Choose a reason for hiding this comment

vbarua Aug 28, 2023

Choose a reason for hiding this comment

vbarua Aug 28, 2023

Choose a reason for hiding this comment

vbarua Aug 28, 2023

Choose a reason for hiding this comment

vbarua Aug 28, 2023

Choose a reason for hiding this comment

vbarua Aug 28, 2023 • edited Loading

Choose a reason for hiding this comment

vbarua Aug 28, 2023 • edited Loading

Choose a reason for hiding this comment

EpsilonPrime left a comment

Choose a reason for hiding this comment

vbarua commented Aug 28, 2023 •

edited

Loading

vbarua Aug 28, 2023 •

edited

Loading

vbarua Aug 28, 2023 •

edited

Loading