[multistage] Add support for the ranking ROW_NUMBER() window function

pinot-query-planner/src/main/java/org/apache/calcite/rel/rules/PinotWindowExchangeNodeInsertRule.java

@@ -19,6 +19,7 @@
 package org.apache.calcite.rel.rules;
 
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import java.util.Collections;
 import java.util.HashSet;
@@ -49,7 +50,7 @@ public class PinotWindowExchangeNodeInsertRule extends RelOptRule {
   // Supported window functions
   // OTHER_FUNCTION supported are: BOOL_AND, BOOL_OR
   private static final Set<SqlKind> SUPPORTED_WINDOW_FUNCTION_KIND = ImmutableSet.of(SqlKind.SUM, SqlKind.SUM0,
-      SqlKind.MIN, SqlKind.MAX, SqlKind.COUNT, SqlKind.OTHER_FUNCTION);
+      SqlKind.MIN, SqlKind.MAX, SqlKind.COUNT, SqlKind.ROW_NUMBER, SqlKind.OTHER_FUNCTION);
 
   public PinotWindowExchangeNodeInsertRule(RelBuilderFactory factory) {
     super(operand(LogicalWindow.class, any()), factory, null);
@@ -145,19 +146,26 @@ private void validateWindowAggCallsSupported(Window.Group windowGroup) {
   }
 
   private void validateWindowFrames(Window.Group windowGroup) {
+    // Has ROWS only aggregation call kind (e.g. ROW_NUMBER)?
+    boolean isRowsOnlyTypeAggregateCall = isRowsOnlyAggregationCallType(windowGroup.aggCalls);
     // For Phase 1 only the default frame is supported
-    Preconditions.checkState(!windowGroup.isRows, "Default frame must be of type RANGE and not ROWS");
+    Preconditions.checkState(!windowGroup.isRows || isRowsOnlyTypeAggregateCall,
+        "Default frame must be of type RANGE and not ROWS unless this is a ROWS only aggregation function");
     Preconditions.checkState(windowGroup.lowerBound.isPreceding() && windowGroup.lowerBound.isUnbounded(),
         String.format("Lower bound must be UNBOUNDED PRECEDING but it is: %s", windowGroup.lowerBound));
-    if (windowGroup.orderKeys.getKeys().isEmpty()) {
+    if (windowGroup.orderKeys.getKeys().isEmpty() && !isRowsOnlyTypeAggregateCall) {
       Preconditions.checkState(windowGroup.upperBound.isFollowing() && windowGroup.upperBound.isUnbounded(),
-          String.format("Upper bound must be UNBOUNDED PRECEDING but it is: %s", windowGroup.upperBound));
+          String.format("Upper bound must be UNBOUNDED FOLLOWING but it is: %s", windowGroup.upperBound));
     } else {
       Preconditions.checkState(windowGroup.upperBound.isCurrentRow(),
           String.format("Upper bound must be CURRENT ROW but it is: %s", windowGroup.upperBound));
     }
   }
 
+  private boolean isRowsOnlyAggregationCallType(ImmutableList<Window.RexWinAggCall> aggCalls) {
+    return aggCalls.stream().anyMatch(aggCall -> aggCall.getKind().equals(SqlKind.ROW_NUMBER));
+  }
+
   private boolean isPartitionByOnlyQuery(Window.Group windowGroup) {
     boolean isPartitionByOnly = false;
     if (windowGroup.orderKeys.getKeys().isEmpty()) {

pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java

@@ -109,6 +109,7 @@ protected Object[][] provideQueries() {
         new Object[]{"SELECT a.col1, SUM(a.col3) OVER (PARTITION BY a.col2, a.col1) FROM a"},
         new Object[]{"SELECT a.col1, SUM(a.col3) OVER (ORDER BY a.col2, a.col1), MIN(a.col3) OVER (ORDER BY a.col2, "
             + "a.col1) FROM a"},
+        new Object[]{"SELECT a.col1, ROW_NUMBER() OVER(PARTITION BY a.col2 ORDER BY a.col3) FROM a"},
         new Object[]{"SELECT a.col1, SUM(a.col3) OVER (ORDER BY a.col2), MIN(a.col3) OVER (ORDER BY a.col2) FROM a"},
         new Object[]{"SELECT /*+ skipLeafStageGroupByAggregation */ a.col1, SUM(a.col3) FROM a WHERE a.col3 >= 0"
             + " AND a.col2 = 'a' GROUP BY a.col1"},