From 64bee4199b8b9f7de951eba0233e335f858e9d31 Mon Sep 17 00:00:00 2001
From: guoxu <13910754971@163.com>
Date: Fri, 21 Feb 2025 15:27:37 +0800
Subject: [PATCH] [feat][dingo] Optimize count

---
 .../java/io/dingodb/common/CoprocessorV2.java   |  3 +++
 .../operator/params/ScanWithRelOpParam.java     | 17 +++++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/dingo-common/src/main/java/io/dingodb/common/CoprocessorV2.java b/dingo-common/src/main/java/io/dingodb/common/CoprocessorV2.java
index 2be4e9b824..0295f5f15b 100644
--- 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
index 2b7269d68c..d1feee3d91 100644
--- 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) {