|
19 | 19 | package org.apache.calcite.rel.rules; |
20 | 20 |
|
21 | 21 | import com.google.common.base.Preconditions; |
| 22 | +import com.google.common.collect.ImmutableList; |
22 | 23 | import com.google.common.collect.ImmutableSet; |
23 | 24 | import java.util.Collections; |
24 | 25 | import java.util.HashSet; |
@@ -49,7 +50,7 @@ public class PinotWindowExchangeNodeInsertRule extends RelOptRule { |
49 | 50 | // Supported window functions |
50 | 51 | // OTHER_FUNCTION supported are: BOOL_AND, BOOL_OR |
51 | 52 | private static final Set<SqlKind> SUPPORTED_WINDOW_FUNCTION_KIND = ImmutableSet.of(SqlKind.SUM, SqlKind.SUM0, |
52 | | - SqlKind.MIN, SqlKind.MAX, SqlKind.COUNT, SqlKind.OTHER_FUNCTION); |
| 53 | + SqlKind.MIN, SqlKind.MAX, SqlKind.COUNT, SqlKind.ROW_NUMBER, SqlKind.OTHER_FUNCTION); |
53 | 54 |
|
54 | 55 | public PinotWindowExchangeNodeInsertRule(RelBuilderFactory factory) { |
55 | 56 | super(operand(LogicalWindow.class, any()), factory, null); |
@@ -145,19 +146,26 @@ private void validateWindowAggCallsSupported(Window.Group windowGroup) { |
145 | 146 | } |
146 | 147 |
|
147 | 148 | private void validateWindowFrames(Window.Group windowGroup) { |
| 149 | + // Has ROWS only aggregation call kind (e.g. ROW_NUMBER)? |
| 150 | + boolean isRowsOnlyTypeAggregateCall = isRowsOnlyAggregationCallType(windowGroup.aggCalls); |
148 | 151 | // For Phase 1 only the default frame is supported |
149 | | - Preconditions.checkState(!windowGroup.isRows, "Default frame must be of type RANGE and not ROWS"); |
| 152 | + Preconditions.checkState(!windowGroup.isRows || isRowsOnlyTypeAggregateCall, |
| 153 | + "Default frame must be of type RANGE and not ROWS unless this is a ROWS only aggregation function"); |
150 | 154 | Preconditions.checkState(windowGroup.lowerBound.isPreceding() && windowGroup.lowerBound.isUnbounded(), |
151 | 155 | String.format("Lower bound must be UNBOUNDED PRECEDING but it is: %s", windowGroup.lowerBound)); |
152 | | - if (windowGroup.orderKeys.getKeys().isEmpty()) { |
| 156 | + if (windowGroup.orderKeys.getKeys().isEmpty() && !isRowsOnlyTypeAggregateCall) { |
153 | 157 | Preconditions.checkState(windowGroup.upperBound.isFollowing() && windowGroup.upperBound.isUnbounded(), |
154 | | - String.format("Upper bound must be UNBOUNDED PRECEDING but it is: %s", windowGroup.upperBound)); |
| 158 | + String.format("Upper bound must be UNBOUNDED FOLLOWING but it is: %s", windowGroup.upperBound)); |
155 | 159 | } else { |
156 | 160 | Preconditions.checkState(windowGroup.upperBound.isCurrentRow(), |
157 | 161 | String.format("Upper bound must be CURRENT ROW but it is: %s", windowGroup.upperBound)); |
158 | 162 | } |
159 | 163 | } |
160 | 164 |
|
| 165 | + private boolean isRowsOnlyAggregationCallType(ImmutableList<Window.RexWinAggCall> aggCalls) { |
| 166 | + return aggCalls.stream().anyMatch(aggCall -> aggCall.getKind().equals(SqlKind.ROW_NUMBER)); |
| 167 | + } |
| 168 | + |
161 | 169 | private boolean isPartitionByOnlyQuery(Window.Group windowGroup) { |
162 | 170 | boolean isPartitionByOnly = false; |
163 | 171 | if (windowGroup.orderKeys.getKeys().isEmpty()) { |
|
0 commit comments