[Feature](materialized-view) support visitBitmapUnion mv rewrite (#29200

) * support visitBitmapUnion rewrite * add case
apache · Dec 28, 2023 · 276440a · 276440a
1 parent 388a158
commit 276440a
Show file tree

Hide file tree

Showing 3 changed files with 86 additions and 24 deletions.
diff --git a/.../java/org/apache/doris/nereids/rules/rewrite/mv/SelectMaterializedIndexWithAggregate.java b/.../java/org/apache/doris/nereids/rules/rewrite/mv/SelectMaterializedIndexWithAggregate.java
@@ -53,7 +53,6 @@
 import org.apache.doris.nereids.trees.expressions.functions.agg.Sum;
 import org.apache.doris.nereids.trees.expressions.functions.combinator.MergeCombinator;
 import org.apache.doris.nereids.trees.expressions.functions.combinator.StateCombinator;
-import org.apache.doris.nereids.trees.expressions.functions.scalar.BitmapHash;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.HllHash;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.ToBitmap;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.ToBitmapWithCheck;
@@ -1264,6 +1263,66 @@ public Expression visitCount(Count count, RewriteContext context) {
             return count;
         }
 
+        /**
+         * bitmap_union(to_bitmap(col)) ->
+         * bitmap_union(mva_BITMAP_UNION__to_bitmap_with_check(col))
+         */
+        @Override
+        public Expression visitBitmapUnion(BitmapUnion bitmapUnion, RewriteContext context) {
+            Expression result = visitAggregateFunction(bitmapUnion, context);
+            if (result != bitmapUnion) {
+                return result;
+            }
+            if (bitmapUnion.child() instanceof ToBitmap) {
+                ToBitmap toBitmap = (ToBitmap) bitmapUnion.child();
+                Optional<Slot> slotOpt = ExpressionUtils.extractSlotOrCastOnSlot(toBitmap.child());
+                if (slotOpt.isPresent()) {
+                    String bitmapUnionColumn = normalizeName(CreateMaterializedViewStmt
+                            .mvColumnBuilder(AggregateType.BITMAP_UNION, CreateMaterializedViewStmt
+                                    .mvColumnBuilder(new ToBitmapWithCheck(toBitmap.child()).toSql())));
+
+                    Column mvColumn = context.checkContext.getColumn(bitmapUnionColumn);
+                    // has bitmap_union column
+                    if (mvColumn != null && context.checkContext.valueNameToColumn.containsValue(mvColumn)) {
+
+                        Slot bitmapUnionSlot = context.checkContext.scan.getOutputByIndex(context.checkContext.index)
+                                .stream().filter(s -> bitmapUnionColumn.equalsIgnoreCase(normalizeName(s.getName())))
+                                .findFirst().orElseThrow(
+                                        () -> new AnalysisException("cannot find bitmap union slot when select mv"));
+
+                        context.exprRewriteMap.slotMap.put(slotOpt.get(), bitmapUnionSlot);
+                        context.exprRewriteMap.projectExprMap.put(toBitmap, bitmapUnionSlot);
+                        BitmapUnion newBitmapUnion = new BitmapUnion(bitmapUnionSlot);
+                        context.exprRewriteMap.aggFuncMap.put(bitmapUnion, newBitmapUnion);
+                        return newBitmapUnion;
+                    }
+                }
+            } else {
+                Expression child = bitmapUnion.child();
+                String bitmapUnionColumn = normalizeName(CreateMaterializedViewStmt.mvColumnBuilder(
+                        AggregateType.BITMAP_UNION, CreateMaterializedViewStmt.mvColumnBuilder(child.toSql())));
+
+                Column mvColumn = context.checkContext.getColumn(bitmapUnionColumn);
+                // has bitmap_union column
+                if (mvColumn != null && context.checkContext.valueNameToColumn.containsValue(mvColumn)) {
+
+                    Slot bitmapUnionSlot = context.checkContext.scan.getOutputByIndex(context.checkContext.index)
+                            .stream().filter(s -> bitmapUnionColumn.equalsIgnoreCase(normalizeName(s.getName())))
+                            .findFirst()
+                            .orElseThrow(() -> new AnalysisException("cannot find bitmap union slot when select mv"));
+                    if (child instanceof Slot) {
+                        context.exprRewriteMap.slotMap.put((Slot) child, bitmapUnionSlot);
+                    }
+                    context.exprRewriteMap.projectExprMap.put(child, bitmapUnionSlot);
+                    BitmapUnion newBitmapUnion = new BitmapUnion(bitmapUnionSlot);
+                    context.exprRewriteMap.aggFuncMap.put(bitmapUnion, newBitmapUnion);
+                    return newBitmapUnion;
+                }
+            }
+
+            return bitmapUnion;
+        }
+
         /**
          * bitmap_union_count(to_bitmap(col)) -> bitmap_union_count(mva_BITMAP_UNION__to_bitmap_with_check(col))
          */
@@ -1300,32 +1359,26 @@ public Expression visitBitmapUnionCount(BitmapUnionCount bitmapUnionCount, Rewri
                         return newBitmapUnionCount;
                     }
                 }
-            } else if (bitmapUnionCount.child() instanceof BitmapHash) {
-                BitmapHash bitmapHash = (BitmapHash) bitmapUnionCount.child();
-                Optional<Slot> slotOpt = ExpressionUtils.extractSlotOrCastOnSlot(bitmapHash.child());
-                if (slotOpt.isPresent()) {
-                    String bitmapUnionCountColumn = normalizeName(
-                            CreateMaterializedViewStmt.mvColumnBuilder(AggregateType.BITMAP_UNION,
-                                    CreateMaterializedViewStmt.mvColumnBuilder(bitmapHash.toSql())));
-
-                    Column mvColumn = context.checkContext.getColumn(bitmapUnionCountColumn);
-                    // has bitmap_union_count column
-                    if (mvColumn != null && context.checkContext.valueNameToColumn.containsValue(mvColumn)) {
+            } else {
+                Expression child = bitmapUnionCount.child();
+                String bitmapUnionCountColumn = normalizeName(CreateMaterializedViewStmt.mvColumnBuilder(
+                        AggregateType.BITMAP_UNION, CreateMaterializedViewStmt.mvColumnBuilder(child.toSql())));
 
-                        Slot bitmapUnionCountSlot = context.checkContext.scan
-                                .getOutputByIndex(context.checkContext.index)
-                                .stream()
-                                .filter(s -> bitmapUnionCountColumn.equalsIgnoreCase(normalizeName(s.getName())))
-                                .findFirst()
-                                .orElseThrow(() -> new AnalysisException(
-                                        "cannot find bitmap union count slot when select mv"));
+                Column mvColumn = context.checkContext.getColumn(bitmapUnionCountColumn);
+                // has bitmap_union_count column
+                if (mvColumn != null && context.checkContext.valueNameToColumn.containsValue(mvColumn)) {
 
-                        context.exprRewriteMap.slotMap.put(slotOpt.get(), bitmapUnionCountSlot);
-                        context.exprRewriteMap.projectExprMap.put(bitmapHash, bitmapUnionCountSlot);
-                        BitmapUnionCount newBitmapUnionCount = new BitmapUnionCount(bitmapUnionCountSlot);
-                        context.exprRewriteMap.aggFuncMap.put(bitmapUnionCount, newBitmapUnionCount);
-                        return newBitmapUnionCount;
+                    Slot bitmapUnionCountSlot = context.checkContext.scan.getOutputByIndex(context.checkContext.index)
+                            .stream().filter(s -> bitmapUnionCountColumn.equalsIgnoreCase(normalizeName(s.getName())))
+                            .findFirst().orElseThrow(
+                                    () -> new AnalysisException("cannot find bitmap union count slot when select mv"));
+                    if (child instanceof Slot) {
+                        context.exprRewriteMap.slotMap.put((Slot) child, bitmapUnionCountSlot);
                     }
+                    context.exprRewriteMap.projectExprMap.put(child, bitmapUnionCountSlot);
+                    BitmapUnionCount newBitmapUnionCount = new BitmapUnionCount(bitmapUnionCountSlot);
+                    context.exprRewriteMap.aggFuncMap.put(bitmapUnionCount, newBitmapUnionCount);
+                    return newBitmapUnionCount;
                 }
             }
 

diff --git a/regression-test/data/mv_p0/ut/testBitmapUnionInQuery/testBitmapUnionInQuery.out b/regression-test/data/mv_p0/ut/testBitmapUnionInQuery/testBitmapUnionInQuery.out
@@ -7,3 +7,6 @@
 -- !select_mv --
 1	2
 
+-- !select_mv --
+1	2
+
diff --git a/regression-test/suites/mv_p0/ut/testBitmapUnionInQuery/testBitmapUnionInQuery.groovy b/regression-test/suites/mv_p0/ut/testBitmapUnionInQuery/testBitmapUnionInQuery.groovy
@@ -47,4 +47,10 @@ suite ("testBitmapUnionInQuery") {
         contains "(user_tags_mv)"
     }
     qt_select_mv "select user_id, bitmap_union_count(to_bitmap(tag_id)) a from user_tags group by user_id having a>1 order by a;"
+
+    explain {
+        sql("select user_id, bitmap_count(bitmap_union(to_bitmap(tag_id))) a from user_tags group by user_id having a>1 order by a;")
+        contains "(user_tags_mv)"
+    }
+    qt_select_mv "select user_id, bitmap_count(bitmap_union(to_bitmap(tag_id))) a from user_tags group by user_id having a>1 order by a;"
 }
-Original file line number
+Diff line change
@@ Expand Up / @@ -7,3 +7,6 @@ @@
     -- !select_mv --
 	2
+    -- !select_mv --
+	2