[BugFix] Pushdown distinct agg across window

Signed-off-by: satanson <[email protected]>
StarRocks · Dec 4, 2023 · f469c95 · f469c95
1 parent 6b31c86
commit f469c95
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 3 deletions.
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/LogicalPlanPrinter.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/LogicalPlanPrinter.java
@@ -30,6 +30,7 @@
 import com.starrocks.sql.optimizer.operator.logical.LogicalScanOperator;
 import com.starrocks.sql.optimizer.operator.logical.LogicalTopNOperator;
 import com.starrocks.sql.optimizer.operator.logical.LogicalValuesOperator;
+import com.starrocks.sql.optimizer.operator.logical.LogicalWindowOperator;
 import com.starrocks.sql.optimizer.operator.physical.PhysicalAssertOneRowOperator;
 import com.starrocks.sql.optimizer.operator.physical.PhysicalCTEAnchorOperator;
 import com.starrocks.sql.optimizer.operator.physical.PhysicalCTEConsumeOperator;
@@ -251,6 +252,26 @@ public OperatorStr visitLogicalJoin(OptExpression optExpression, Integer step) {
             return new OperatorStr(sb.toString(), step, Arrays.asList(leftChild, rightChild));
         }
 
+        @Override
+        public OperatorStr visitLogicalWindow(OptExpression optExpression, Integer step) {
+            OperatorStr child = visit(optExpression.getInputs().get(0), step + 1);
+
+            LogicalWindowOperator window = optExpression.getOp().cast();
+            String windowCallStr = window.getWindowCall().entrySet().stream()
+                    .map(e -> String.format("%d: %s", e.getKey().getId(),
+                            scalarOperatorStringFunction.apply(e.getValue())))
+                    .collect(Collectors.joining(", "));
+            String windowDefStr = window.getAnalyticWindow() != null ? window.getAnalyticWindow().toSql() : "NONE";
+            String partitionByStr = window.getPartitionExpressions().stream()
+                    .map(scalarOperatorStringFunction).collect(Collectors.joining(", "));
+            String orderByStr = window.getOrderByElements().stream().map(Ordering::toString)
+                    .collect(Collectors.joining(", "));
+            return new OperatorStr("logical window( calls=[" +
+                    windowCallStr + "], window=" +
+                    windowDefStr + ", partitionBy=" +
+                    partitionByStr + ", orderBy=" + orderByStr + ")", step, Collections.singletonList(child));
+        }
+
         @Override
         public OperatorStr visitLogicalApply(OptExpression optExpression, Integer step) {
             OperatorStr leftChild = visit(optExpression.getInputs().get(0), step + 1);

diff --git a/...n/java/com/starrocks/sql/optimizer/rule/tree/pdagg/PushDownDistinctAggregateRewriter.java b/...n/java/com/starrocks/sql/optimizer/rule/tree/pdagg/PushDownDistinctAggregateRewriter.java
@@ -285,7 +285,10 @@ public AggregatePushDownContext visitLogicalProject(OptExpression optExpression,
             // rewrite
             ReplaceColumnRefRewriter rewriter = new ReplaceColumnRefRewriter(projectOp.getColumnRefMap());
             context.aggregations.replaceAll((k, v) -> (CallOperator) rewriter.rewrite(v));
-            context.groupBys.replaceAll((k, v) -> rewriter.rewrite(v));
+            Set<ColumnRefOperator> groupByUsedCols = context.groupBys.values().stream()
+                    .flatMap(v->rewriter.rewrite(v).getColumnRefs().stream()).collect(Collectors.toSet());
+            context.groupBys.clear();
+            groupByUsedCols.forEach(col->context.groupBys.put(col, col));
 
             if (projectOp.getColumnRefMap().values().stream().allMatch(ScalarOperator::isColumnRef)) {
                 return context;

diff --git a/fe/fe-core/src/test/java/com/starrocks/sql/plan/ReplayFromDumpTest.java b/fe/fe-core/src/test/java/com/starrocks/sql/plan/ReplayFromDumpTest.java
@@ -551,10 +551,10 @@ public void testPushDownDistinctAggBelowWindowRewrite() throws Exception {
         Pair<QueryDumpInfo, String> replayPair =
                 getPlanFragment(getDumpInfoFromFile("query_dump/pushdown_distinct_agg_below_window"), null,
                         TExplainLevel.COSTS);
-        Assert.assertTrue(replayPair.second.contains("  1:AGGREGATE (update finalize)\n" +
+        Assert.assertTrue(replayPair.second, replayPair.second.contains("  1:AGGREGATE (update finalize)\n" +
                 "  |  aggregate: sum[([3: gross, DECIMAL128(10,2), false]); args: DECIMAL128; " +
                 "result: DECIMAL128(38,2); args nullable: false; result nullable: true]\n" +
-                "  |  group by: [2: trans_date, DATE, false], [1: country, VARCHAR, true]\n" +
+                "  |  group by: [1: country, VARCHAR, true], [2: trans_date, DATE, false]\n" +
                 "  |  cardinality: 49070\n"));
     }
 
@@ -795,4 +795,24 @@ public void testTwoStageAgg() throws Exception {
         Assert.assertTrue(replayPair.second, replayPair.second.contains("0:OlapScanNode\n" +
                 "     table: lineorder_2, rollup: lineorder_2"));
     }
+
+    @Test
+    public void testPushDistinctAggDownWindow() throws Exception {
+        Pair<QueryDumpInfo, String> replayPair =
+                getPlanFragment(getDumpInfoFromFile("query_dump/pushdown_distinct_agg_below_window2"),
+                        null, TExplainLevel.NORMAL);
+        System.out.println(replayPair.second);
+        Assert.assertTrue(replayPair.second, replayPair.second.contains("  3:ANALYTIC\n" +
+                "  |  functions: [, sum(5: sum), ]\n" +
+                "  |  partition by: 1: TIME\n" +
+                "  |  \n" +
+                "  2:SORT\n" +
+                "  |  order by: <slot 1> 1: TIME ASC\n" +
+                "  |  offset: 0\n" +
+                "  |  \n" +
+                "  1:AGGREGATE (update finalize)\n" +
+                "  |  output: sum(2: NUM)\n" +
+                "  |  group by: 1: TIME"));
+    }
+
 }