revert and refactor

apache · Aug 9, 2023 · 8e800ea · 8e800ea
1 parent 9099aeb
commit 8e800ea
Show file tree

Hide file tree

Showing 9 changed files with 100 additions and 182 deletions.
diff --git a/...e-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java b/...e-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java
@@ -567,9 +567,6 @@ public PlanFragment visitPhysicalOlapScan(PhysicalOlapScan olapScan, PlanTransla
         );
         // TODO: we need to remove all finalizeForNereids
         olapScanNode.finalizeForNereids();
-        if (olapScan.getPushDownCountOnIndex()) {
-            olapScanNode.setPushDownAggNoGrouping(TPushAggOp.COUNT_ON_INDEX);
-        }
         // Create PlanFragment
         // TODO: use a util function to convert distribution to DataPartition
         DataPartition dataPartition = DataPartition.RANDOM;
@@ -799,6 +796,9 @@ public PlanFragment visitPhysicalStorageLayerAggregate(
             case COUNT:
                 pushAggOp = TPushAggOp.COUNT;
                 break;
+            case COUNT_ON_MATCH:
+                pushAggOp = TPushAggOp.COUNT_ON_INDEX;
+                break;
             case MIN_MAX:
                 pushAggOp = TPushAggOp.MINMAX;
                 break;

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java
@@ -82,7 +82,6 @@
 import org.apache.doris.nereids.rules.rewrite.PushFilterInsideJoin;
 import org.apache.doris.nereids.rules.rewrite.PushProjectIntoOneRowRelation;
 import org.apache.doris.nereids.rules.rewrite.PushProjectThroughUnion;
-import org.apache.doris.nereids.rules.rewrite.PushdownCountOnIndex;
 import org.apache.doris.nereids.rules.rewrite.PushdownFilterThroughProject;
 import org.apache.doris.nereids.rules.rewrite.PushdownFilterThroughWindow;
 import org.apache.doris.nereids.rules.rewrite.PushdownLimit;
@@ -291,11 +290,6 @@ public class Rewriter extends AbstractBatchJobExecutor {
             topic("topn optimize",
                     topDown(new DeferMaterializeTopNResult())
             ),
-            topic("Count agg push down",
-                    topDown(
-                            new PushdownCountOnIndex()
-                    )
-            ),
             // this rule batch must keep at the end of rewrite to do some plan check
             topic("Final rewrite and check",
                     custom(RuleType.ENSURE_PROJECT_ON_TOP_JOIN, EnsureProjectOnTopJoin::new),

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java
@@ -221,7 +221,6 @@ public enum RuleType {
     PUSH_CONJUNCTS_INTO_JDBC_SCAN(RuleTypeClass.REWRITE),
     OLAP_SCAN_TABLET_PRUNE(RuleTypeClass.REWRITE),
     PUSH_AGGREGATE_TO_OLAP_SCAN(RuleTypeClass.REWRITE),
-    PUSH_COUNT_ON_INDEX_TO_OLAP_SCAN(RuleTypeClass.REWRITE),
     EXTRACT_SINGLE_TABLE_EXPRESSION_FROM_DISJUNCTION(RuleTypeClass.REWRITE),
     HIDE_ONE_ROW_RELATION_UNDER_UNION(RuleTypeClass.REWRITE),
     PUSH_PROJECT_THROUGH_UNION(RuleTypeClass.REWRITE),
@@ -323,6 +322,7 @@ public enum RuleType {
     LOGICAL_ASSERT_NUM_ROWS_TO_PHYSICAL_ASSERT_NUM_ROWS(RuleTypeClass.IMPLEMENTATION),
     STORAGE_LAYER_AGGREGATE_WITHOUT_PROJECT(RuleTypeClass.IMPLEMENTATION),
     STORAGE_LAYER_AGGREGATE_WITH_PROJECT(RuleTypeClass.IMPLEMENTATION),
+    COUNT_ON_INDEX(RuleTypeClass.IMPLEMENTATION),
     ONE_PHASE_AGGREGATE_WITHOUT_DISTINCT(RuleTypeClass.IMPLEMENTATION),
     TWO_PHASE_AGGREGATE_WITHOUT_DISTINCT(RuleTypeClass.IMPLEMENTATION),
     TWO_PHASE_AGGREGATE_WITH_COUNT_DISTINCT_MULTI(RuleTypeClass.IMPLEMENTATION),

diff --git a/...core/src/main/java/org/apache/doris/nereids/rules/implementation/AggregateStrategies.java b/...core/src/main/java/org/apache/doris/nereids/rules/implementation/AggregateStrategies.java
@@ -36,6 +36,7 @@
 import org.apache.doris.nereids.trees.expressions.Cast;
 import org.apache.doris.nereids.trees.expressions.Expression;
 import org.apache.doris.nereids.trees.expressions.IsNull;
+import org.apache.doris.nereids.trees.expressions.Match;
 import org.apache.doris.nereids.trees.expressions.NamedExpression;
 import org.apache.doris.nereids.trees.expressions.SlotReference;
 import org.apache.doris.nereids.trees.expressions.functions.ExpressionTrait;
@@ -56,6 +57,7 @@
 import org.apache.doris.nereids.trees.plans.algebra.Project;
 import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate;
 import org.apache.doris.nereids.trees.plans.logical.LogicalFileScan;
+import org.apache.doris.nereids.trees.plans.logical.LogicalFilter;
 import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan;
 import org.apache.doris.nereids.trees.plans.logical.LogicalProject;
 import org.apache.doris.nereids.trees.plans.logical.LogicalRelation;
@@ -85,10 +87,12 @@
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 
-/** AggregateStrategies */
+/**
+ * AggregateStrategies
+ */
 @DependsRules({
-    NormalizeAggregate.class,
-    FoldConstantRuleOnFE.class
+        NormalizeAggregate.class,
+        FoldConstantRuleOnFE.class
 })
 public class AggregateStrategies implements ImplementationRuleFactory {
 
@@ -97,6 +101,28 @@ public List<Rule> buildRules() {
         PatternDescriptor<LogicalAggregate<GroupPlan>> basePattern = logicalAggregate();
 
         return ImmutableList.of(
+            RuleType.COUNT_ON_INDEX.build(
+                logicalAggregate(
+                    logicalProject(
+                        logicalFilter(
+                            logicalOlapScan()
+                        ).when(filter -> containsMatchExpression(filter.getExpressions())
+                                && filter.getExpressions().size() == 1)
+                    ))
+                    .when(agg -> enablePushDownCountOnIndex())
+                    .when(agg -> agg.getGroupByExpressions().size() == 0)
+                    .when(agg -> {
+                        Set<AggregateFunction> funcs = agg.getAggregateFunctions();
+                        return !funcs.isEmpty() && funcs.stream().allMatch(f -> f instanceof Count && !f.isDistinct());
+                    })
+                    .thenApply(ctx -> {
+                        LogicalAggregate<LogicalProject<LogicalFilter<LogicalOlapScan>>> agg = ctx.root;
+                        LogicalProject<LogicalFilter<LogicalOlapScan>> project = agg.child();
+                        LogicalFilter<LogicalOlapScan> filter = project.child();
+                        LogicalOlapScan olapScan = filter.child();
+                        return pushdownCountOnIndex(agg, project, filter, olapScan, ctx.cascadesContext);
+                    })
+            ),
             RuleType.STORAGE_LAYER_AGGREGATE_WITHOUT_PROJECT.build(
                 logicalAggregate(
                     logicalOlapScan()
@@ -188,6 +214,55 @@ && couldConvertToMulti(agg))
         );
     }
 
+    private boolean containsMatchExpression(List<Expression> expressions) {
+        return expressions.stream().allMatch(expr -> expr instanceof Match);
+    }
+
+    private boolean enablePushDownCountOnIndex() {
+        ConnectContext connectContext = ConnectContext.get();
+        return connectContext == null || connectContext.getSessionVariable().isEnablePushDownCountOnIndex();
+    }
+
+    /**
+     * sql: select count(*) from tbl where column match 'token'
+     * <p>
+     * before:
+     * <p>
+     *               LogicalAggregate(groupBy=[], output=[count(*)])
+     *                                |
+     *                     LogicalFilter(column match 'token')
+     *                                |
+     *                       LogicalOlapScan(table=tbl)
+     * <p>
+     * after:
+     * <p>
+     *               LogicalAggregate(groupBy=[], output=[count(*)])
+     *                                |
+     *                    LogicalFilter(column match 'token')
+     *                                |
+     *        PhysicalStorageLayerAggregate(pushAggOp=COUNT_ON_INDEX, table=PhysicalOlapScan(table=tbl))
+     *
+     */
+    private LogicalAggregate<? extends Plan> pushdownCountOnIndex(
+            LogicalAggregate<? extends Plan> agg,
+            LogicalProject<? extends Plan> project,
+            LogicalFilter<? extends Plan> filter,
+            LogicalOlapScan olapScan,
+            CascadesContext cascadesContext) {
+        PhysicalOlapScan physicalOlapScan
+                = (PhysicalOlapScan) new LogicalOlapScanToPhysicalOlapScan()
+                .build()
+                .transform(olapScan, cascadesContext)
+                .get(0);
+        return agg.withChildren(ImmutableList.of(
+                project.withChildren(ImmutableList.of(
+                        filter.withChildren(ImmutableList.of(
+                                new PhysicalStorageLayerAggregate(
+                                        physicalOlapScan,
+                                        PushDownAggOp.COUNT_ON_MATCH)))))
+        ));
+    }
+
     /**
      * sql: select count(*) from tbl
      * <p>

diff --git a/...java/org/apache/doris/nereids/rules/implementation/LogicalOlapScanToPhysicalOlapScan.java b/...java/org/apache/doris/nereids/rules/implementation/LogicalOlapScanToPhysicalOlapScan.java
@@ -59,7 +59,6 @@ public Rule build() {
                         convertDistribution(olapScan),
                         olapScan.getPreAggStatus(),
                         olapScan.getOutputByIndex(olapScan.getTable().getBaseIndexId()),
-                        olapScan.isCountOnIndexPushedDown(),
                         Optional.empty(),
                         olapScan.getLogicalProperties())
         ).toRule(RuleType.LOGICAL_OLAP_SCAN_TO_PHYSICAL_OLAP_SCAN_RULE);

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushdownCountOnIndex.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushdownCountOnIndex.java