[Opt](Iceberg) handle count pushdown in fe side

apache · May 15, 2024 · d08ddcc · d08ddcc
1 parent fd749b8
commit d08ddcc
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 1 deletion.
diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/FileScanNode.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/FileScanNode.java
@@ -71,6 +71,7 @@ public abstract class FileScanNode extends ExternalScanNode {
     protected long totalPartitionNum = 0;
     protected long readPartitionNum = 0;
     protected long fileSplitSize;
+    public long rowCount = 0;
 
     public FileScanNode(PlanNodeId id, TupleDescriptor desc, String planNodeName, StatisticalType statisticalType,
             boolean needCheckColumnPriv) {

diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/source/IcebergScanNode.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/source/IcebergScanNode.java
@@ -212,6 +212,11 @@ private List<Split> doGetSplits() throws UserException {
         HashSet<String> partitionPathSet = new HashSet<>();
         boolean isPartitionedTable = icebergTable.spec().isPartitioned();
 
+        long rowCount = getCountFromSnapshot();
+        if (getPushDownAggNoGroupingOp().equals(TPushAggOp.COUNT) && rowCount > 0) {
+            this.rowCount = rowCount;
+            return new ArrayList<>();
+        }
         CloseableIterable<FileScanTask> fileScanTasks = TableScanUtil.splitFiles(scan.planFiles(), splitSize);
         try (CloseableIterable<CombinedScanTask> combinedScanTasks =
                 TableScanUtil.planTasks(fileScanTasks, splitSize, 1, 0)) {
@@ -264,6 +269,7 @@ private List<Split> doGetSplits() throws UserException {
             throw new UserException(e.getMessage(), e.getCause());
         }
 
+        // TODO: Need to delete this as we can handle count pushdown in fe side
         TPushAggOp aggOp = getPushDownAggNoGroupingOp();
         if (aggOp.equals(TPushAggOp.COUNT) && getCountFromSnapshot() > 0) {
             // we can create a special empty split and skip the plan process

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/NereidsPlanner.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/NereidsPlanner.java
@@ -27,6 +27,7 @@
 import org.apache.doris.common.Pair;
 import org.apache.doris.common.UserException;
 import org.apache.doris.common.profile.SummaryProfile;
+import org.apache.doris.datasource.iceberg.source.IcebergScanNode;
 import org.apache.doris.nereids.CascadesContext.Lock;
 import org.apache.doris.nereids.exceptions.AnalysisException;
 import org.apache.doris.nereids.glue.LogicalPlanAdapter;
@@ -54,6 +55,7 @@
 import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
 import org.apache.doris.nereids.trees.plans.logical.LogicalSqlCache;
 import org.apache.doris.nereids.trees.plans.physical.PhysicalEmptyRelation;
+import org.apache.doris.nereids.trees.plans.physical.PhysicalHashAggregate;
 import org.apache.doris.nereids.trees.plans.physical.PhysicalOneRowRelation;
 import org.apache.doris.nereids.trees.plans.physical.PhysicalPlan;
 import org.apache.doris.nereids.trees.plans.physical.PhysicalResultSink;
@@ -414,6 +416,7 @@ private PhysicalPlan chooseBestPlan(Group rootGroup, PhysicalProperties physical
 
     /**
      * getting hints explain string, which specified by enumerate and show in lists
+     *
      * @param hints hint map recorded in statement context
      * @return explain string shows using of hint
      */
@@ -594,6 +597,18 @@ public Optional<ResultSet> handleQueryInFe(StatementBase parsedStmt) {
                 );
             }
             return Optional.of(resultSet);
+        } else if (child instanceof PhysicalHashAggregate && getScanNodes().get(0) instanceof IcebergScanNode) {
+            List<Column> columns = Lists.newArrayList();
+            NamedExpression output = physicalPlan.getOutput().get(0);
+            columns.add(new Column(output.getName(), output.getDataType().toCatalogDataType()));
+            if (((IcebergScanNode) getScanNodes().get(0)).rowCount > 0) {
+                ResultSetMetaData metadata = new CommonResultSet.CommonResultSetMetaData(columns);
+                ResultSet resultSet = new CommonResultSet(metadata, Collections.singletonList(
+                        Lists.newArrayList(String.valueOf(((IcebergScanNode) getScanNodes().get(0)).rowCount))));
+                // only support one iceberg scan node and one count, e.g. select count(*) from icetbl;
+                return Optional.of(resultSet);
+            }
+            return Optional.empty();
         } else {
             return Optional.empty();
         }

diff --git a/fe/fe-core/src/main/java/org/apache/doris/planner/OriginalPlanner.java b/fe/fe-core/src/main/java/org/apache/doris/planner/OriginalPlanner.java
@@ -40,6 +40,7 @@
 import org.apache.doris.catalog.Type;
 import org.apache.doris.common.Config;
 import org.apache.doris.common.UserException;
+import org.apache.doris.datasource.iceberg.source.IcebergScanNode;
 import org.apache.doris.nereids.PlannerHook;
 import org.apache.doris.qe.CommonResultSet;
 import org.apache.doris.qe.ConnectContext;
@@ -636,13 +637,27 @@ public Optional<ResultSet> handleQueryInFe(StatementBase parsedStmt) {
             return Optional.empty();
         }
         SelectStmt parsedSelectStmt = (SelectStmt) parsedStmt;
-        if (!parsedSelectStmt.getTableRefs().isEmpty()) {
+        if (!(singleNodePlanner.getScanNodes().get(0) instanceof IcebergScanNode)
+                && !parsedSelectStmt.getTableRefs().isEmpty()) {
             return Optional.empty();
         }
+
         List<SelectListItem> selectItems = parsedSelectStmt.getSelectList().getItems();
         List<Column> columns = new ArrayList<>(selectItems.size());
         List<String> columnLabels = parsedSelectStmt.getColLabels();
         List<String> data = new ArrayList<>();
+        if (((IcebergScanNode) getScanNodes().get(0)).rowCount > 0) {
+            SelectListItem item = selectItems.get(0);
+            Expr expr = item.getExpr();
+            String columnName = columnLabels.get(0);
+            columns.add(new Column(columnName, expr.getType()));
+            data.add(String.valueOf(((IcebergScanNode) getScanNodes().get(0)).rowCount));
+            ResultSetMetaData metadata = new CommonResultSet.CommonResultSetMetaData(columns);
+            ResultSet resultSet = new CommonResultSet(metadata, Collections.singletonList(data));
+            // only support one iceberg scan node and one count, e.g. select count(*) from icetbl;
+            return Optional.of(resultSet);
+        }
+
         for (int i = 0; i < selectItems.size(); i++) {
             SelectListItem item = selectItems.get(i);
             Expr expr = item.getExpr();