[feat][dingo] Optimize count

dingodb · Feb 21, 2025 · 64bee41 · 64bee41
1 parent 7e3a7be
commit 64bee41
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/dingo-common/src/main/java/io/dingodb/common/CoprocessorV2.java b/dingo-common/src/main/java/io/dingodb/common/CoprocessorV2.java
@@ -47,6 +47,9 @@ public class CoprocessorV2 {
     @Setter
     private int limit;
 
+    @Setter
+    private boolean forAggCount;
+
     @Setter
     private int codecVersion;
 }
diff --git a/dingo-exec/src/main/java/io/dingodb/exec/operator/params/ScanWithRelOpParam.java b/dingo-exec/src/main/java/io/dingodb/exec/operator/params/ScanWithRelOpParam.java
@@ -37,6 +37,9 @@
 import io.dingodb.expr.rel.RelOp;
 import io.dingodb.expr.rel.json.RelOpDeserializer;
 import io.dingodb.expr.rel.json.RelOpSerializer;
+import io.dingodb.expr.rel.op.UngroupedAggregateOp;
+import io.dingodb.expr.runtime.expr.Expr;
+import io.dingodb.expr.runtime.expr.NullaryAggExpr;
 import lombok.Getter;
 import lombok.Setter;
 import org.checkerframework.checker.nullness.qual.NonNull;
@@ -147,12 +150,26 @@ public void init(Vertex vertex) {
                 List<Integer> selection = IntStream.range(0, schema.fieldCount())
                     .boxed()
                     .collect(Collectors.toList());
+
+                boolean forAggCount = false;
+                if(relOp instanceof UngroupedAggregateOp) {
+                    if(((UngroupedAggregateOp) relOp).getAggList().size() == 1) {
+                        Expr expr = ((UngroupedAggregateOp) relOp).getAggList().get(0);
+                        if(expr instanceof NullaryAggExpr) {
+                            if((((NullaryAggExpr)expr).getOp()).getName().equals("COUNT")) {
+                                forAggCount = true;
+                            }
+                        }
+                    }
+                }
+
                 TupleMapping outputKeyMapping = TupleMapping.of(new int[]{});
                 coprocessor = CoprocessorV2.builder()
                     .originalSchema(SchemaWrapperUtils.buildSchemaWrapper(schema, keyMapping, tableId.seq))
                     .resultSchema(SchemaWrapperUtils.buildSchemaWrapper(outputSchema, outputKeyMapping, tableId.seq))
                     .selection(selection)
                     .relExpr(os.toByteArray())
+                    .forAggCount(forAggCount)
                     .codecVersion(codecVersion)
                     .build();
                 if (limit > 0) {