[Enhancement] Support mv rewrite with function array_agg_distinct whe…

…n aggregate rollup (StarRocks#39344) Signed-off-by: Youngwb <[email protected]>
baibaichen · Jan 18, 2024 · f3b2bb4 · f3b2bb4
1 parent 0e18d6f
commit f3b2bb4
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/...sql/optimizer/rule/transformation/materialization/AggregatedMaterializedViewRewriter.java b/...sql/optimizer/rule/transformation/materialization/AggregatedMaterializedViewRewriter.java
@@ -76,6 +76,7 @@ public class AggregatedMaterializedViewRewriter extends MaterializedViewRewriter
     private static final Map<String, String> REWRITE_ROLLUP_FUNCTION_MAP = ImmutableMap.<String, String>builder()
             .put(FunctionSet.COUNT, FunctionSet.SUM)
             .put(FunctionSet.BITMAP_AGG, FunctionSet.BITMAP_UNION)
+            .put(FunctionSet.ARRAY_AGG_DISTINCT, FunctionSet.ARRAY_UNIQUE_AGG)
             .build();
 
     private static final Set<String> SUPPORTED_ROLLUP_FUNCTIONS = ImmutableSet.<String>builder()
@@ -88,6 +89,7 @@ public class AggregatedMaterializedViewRewriter extends MaterializedViewRewriter
             .add(FunctionSet.BITMAP_AGG)
             .add(FunctionSet.HLL_UNION)
             .add(FunctionSet.PERCENTILE_UNION)
+            .add(FunctionSet.ARRAY_AGG_DISTINCT)
             .build();
 
     public AggregatedMaterializedViewRewriter(MvRewriteContext mvRewriteContext) {

diff --git a/fe/fe-core/src/test/java/com/starrocks/planner/MaterializedViewTest.java b/fe/fe-core/src/test/java/com/starrocks/planner/MaterializedViewTest.java
@@ -2718,6 +2718,17 @@ public void testCountWithRollup() {
                 .contains("coalesce");
     }
 
+    @Test
+    public void testArrayAggDistinctWithRollup() {
+        String mv = "select user_id, array_agg_distinct(tag_id) from user_tags group by user_id, time;";
+        testRewriteOK(mv, "select user_id, array_agg_distinct(tag_id) from user_tags group by user_id, time;")
+                .notContain("array_unique_agg");
+        testRewriteOK(mv, "select user_id, array_agg_distinct(tag_id) from user_tags group by user_id")
+                .contains("array_unique_agg");
+        testRewriteOK(mv, "select array_agg_distinct(tag_id) from user_tags")
+                .contains("array_unique_agg");
+    }
+
     @Test
     public void testCountDistinctToBitmapCount1() {
         String mv = "select user_id, bitmap_union(to_bitmap(tag_id)) from user_tags group by user_id;";