Optimize count aggregation performance by utilizing native doc_count in v3

core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java

@@ -17,7 +17,9 @@
 import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_NAME;
 import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_NAME_MAIN;
 import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_NAME_SUBSEARCH;
+import static org.opensearch.sql.calcite.utils.PlanUtils.getInputRefs;
 import static org.opensearch.sql.calcite.utils.PlanUtils.getRelation;
+import static org.opensearch.sql.calcite.utils.PlanUtils.getRexCall;
 import static org.opensearch.sql.calcite.utils.PlanUtils.transformPlanToAttachChild;
 
 import com.google.common.base.Strings;
@@ -53,6 +55,7 @@
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.rex.RexVisitorImpl;
 import org.apache.calcite.rex.RexWindowBounds;
+import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
 import org.apache.calcite.sql.type.SqlTypeFamily;
 import org.apache.calcite.sql.type.SqlTypeName;
@@ -813,6 +816,22 @@ private void projectPlusOverriding(
     context.relBuilder.rename(expectedRenameFields);
   }
 
+  private List<RexInputRef> extractInputRefFromCountCalls(List<RelBuilder.AggCall> aggCalls) {
+    return aggCalls.stream()
+        .map(RelBuilder.AggCall::over)
+        .map(RelBuilder.OverCall::toRex)
+        .flatMap(node -> getRexCall(node, this::isCountField).stream())
+        .flatMap(rex -> getInputRefs(rex).stream())
+        .toList();
+  }
+
+  /** Is count(FIELD) */
+  private boolean isCountField(RexCall call) {
+    return call.isA(SqlKind.COUNT)
+        && call.getOperands().size() == 1 // count(FIELD)
+        && call.getOperands().get(0) instanceof RexInputRef;
+  }
+
   /**
    * Resolve the aggregation with trimming unused fields to avoid bugs in {@link
    * org.apache.calcite.sql2rel.RelDecorrelator#decorrelateRel(Aggregate, boolean)}
@@ -826,6 +845,48 @@ private Pair<List<RexNode>, List<AggCall>> aggregateWithTrimming(
       List<UnresolvedExpression> groupExprList,
       List<UnresolvedExpression> aggExprList,
       CalcitePlanContext context) {
+    Pair<List<RexNode>, List<AggCall>> resolved =
+        resolveAttributesForAggregation(groupExprList, aggExprList, context);
+    // Adding filter(isNotNull(RexInputRef)) before count(FIELD) aggregation
+    // could help on doc_count pushdown:
+    //
+    // Example 1: source=t | stats count(a)
+    // Before: Aggregate(count(a))
+    //         \- Scan t
+    // After: Aggregate(count(a))
+    //        \- Filter(isNotNull(a))
+    //              \- Scan t
+    //
+    // Example 2: source=t | stats count(): no change for count()
+    // Before: Aggregate(count())
+    //           \- Scan t
+    // After: Aggregate(count())
+    //           \- Scan t
+    //
+    // Example 3: source=t | stats count(), count(a): no change for multiple fields
+    // Before: Aggregate(count(), count(a))
+    //           \- Scan t
+    // After: Aggregate(count(), count(a))
+    //           \- Scan t
+    //
+    // Example 4: source=t | stats count(a), count(b): no change for multiple fields
+    // Before: Aggregate(count(a), count(b))
+    //           \- Scan t
+    // After: Aggregate(count(a), count(b))
+    //           \- Scan t
+    //
+    // Example 5: source=t | stats count(a+1): no change for expression
+    // Before: Aggregate(count(a+1))
+    //           \- Scan t
+    // After: Aggregate(count(a+1))
+    //           \- Scan t
+    List<RexInputRef> refsInCountCall = extractInputRefFromCountCalls(resolved.getRight());
+    if (refsInCountCall.size() == 1 && resolved.getRight().size() == refsInCountCall.size()) {
+      context.relBuilder.filter(context.relBuilder.isNotNull(refsInCountCall.getFirst()));
+    }
+
+    // Add project before aggregate:
+    //
     // Example 1: source=t | where a > 1 | stats avg(b + 1) by c
     // Before: Aggregate(avg(b + 1))
     //         \- Filter(a > 1)
@@ -836,20 +897,19 @@ private Pair<List<RexNode>, List<AggCall>> aggregateWithTrimming(
     //              \- Scan t
     //
     // Example 2: source=t | where a > 1 | top b by c
-    // Before: Aggregate(count)
-    //         \-Filter(a > 1)
+    // Before: Aggregate(count(b) by c)
+    //         \-Filter(a > 1 && isNotNull(b))
     //           \- Scan t
-    // After: Aggregate(count)
+    // After: Aggregate(count(b) by c)
     //        \- Project([c, b])
-    //           \- Filter(a > 1)
+    //           \- Filter(a > 1 && isNotNull(b))
     //              \- Scan t
-    // Example 3: source=t | stats count(): no project added for count()
-    // Before: Aggregate(count)
+    //
+    // Example 3: source=t | stats count(): no change for count()
+    // Before: Aggregate(count())
     //           \- Scan t
-    // After: Aggregate(count)
+    // After: Aggregate(count())
     //           \- Scan t
-    Pair<List<RexNode>, List<AggCall>> resolved =
-        resolveAttributesForAggregation(groupExprList, aggExprList, context);
     List<RexInputRef> trimmedRefs = new ArrayList<>();
     trimmedRefs.addAll(PlanUtils.getInputRefs(resolved.getLeft())); // group-by keys first
     trimmedRefs.addAll(PlanUtils.getInputRefsFromAggCall(resolved.getRight()));

core/src/main/java/org/opensearch/sql/calcite/utils/PlanUtils.java

@@ -15,6 +15,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import org.apache.calcite.plan.RelOptTable;
@@ -274,6 +275,26 @@ static List<RexInputRef> getInputRefs(List<RexNode> nodes) {
     return nodes.stream().flatMap(node -> getInputRefs(node).stream()).toList();
   }
 
+  /** Get all uniq RexCall from RexNode with a predicate */
+  static List<RexCall> getRexCall(RexNode node, Predicate<RexCall> predicate) {
+    List<RexCall> list = new ArrayList<>();
+    node.accept(
+        new RexVisitorImpl<Void>(true) {
+          @Override
+          public Void visitCall(RexCall inputCall) {
+            if (predicate.test(inputCall)) {
+              if (!list.contains(inputCall)) {
+                list.add(inputCall);
+              }
+            } else {
+              inputCall.getOperands().forEach(call -> call.accept(this));
+            }
+            return null;
+          }
+        });
+    return list;
+  }
+
   /** Get all uniq input references from a list of agg calls. */
   static List<RexInputRef> getInputRefsFromAggCall(List<RelBuilder.AggCall> aggCalls) {
     return aggCalls.stream()

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java

@@ -5,6 +5,7 @@
 
 package org.opensearch.sql.calcite.remote;
 
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_ACCOUNT;
 import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK;
 import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_LOGS;
 import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_NESTED_SIMPLE;
@@ -644,6 +645,64 @@ public void testExplainMinOnStringField() throws IOException {
         explainQueryToString("source=opensearch-sql_test_index_account | stats min(firstname)"));
   }
 
+  @Test
+  public void testExplainCountsByAgg() throws IOException {
+    enabledOnlyWhenPushdownIsEnabled();
+    String expected = loadExpectedPlan("explain_agg_counts_by1.yaml");
+    // case of only count(): doc_count works
+    assertYamlEqualsJsonIgnoreId(
+        expected,
+        explainQueryToString(
+            String.format(
+                "source=%s | stats count(), count() as c1 by gender", TEST_INDEX_ACCOUNT)));
+
+    // case of only one count(FIELD): doc_count works
+    expected = loadExpectedPlan("explain_agg_counts_by2.yaml");
+    assertYamlEqualsJsonIgnoreId(
+        expected,
+        explainQueryToString(
+            String.format(
+                "source=%s | stats count(balance) as c1, count(balance) as c2 by gender",
+                TEST_INDEX_ACCOUNT)));
+
+    // case of only one count(FIELD): doc_count works
+    expected = loadExpectedPlan("explain_agg_counts_by3.yaml");
+    assertYamlEqualsJsonIgnoreId(
+        expected,
+        explainQueryToString(
+            String.format(
+                "source=%s | eval account_number_alias = account_number"
+                    + " | stats count(account_number), count(account_number_alias) as c2 by gender",
+                TEST_INDEX_ACCOUNT)));
+
+    // case of count() + count(FIELD)): doc_count doesn't work
+    expected = loadExpectedPlan("explain_agg_counts_by4.yaml");
+    assertYamlEqualsJsonIgnoreId(
+        expected,
+        explainQueryToString(
+            String.format(
+                "source=%s | stats count(), count(account_number) by gender", TEST_INDEX_ACCOUNT)));
+
+    // case of count(FIELD1) + count(FIELD2)): doc_count doesn't work
+    expected = loadExpectedPlan("explain_agg_counts_by5.yaml");
+    assertYamlEqualsJsonIgnoreId(
+        expected,
+        explainQueryToString(
+            String.format(
+                "source=%s | stats count(balance), count(account_number) by gender",
+                TEST_INDEX_ACCOUNT)));
+
+    // case of count(EXPRESSION): doc_count doesn't work
+    expected = loadExpectedPlan("explain_agg_counts_by6.yaml");
+    assertYamlEqualsJsonIgnoreId(
+        expected,
+        explainQueryToString(
+            String.format(
+                "source=%s | eval b_1 = balance + 1"
+                    + " | stats count(b_1), count(pow(balance, 2)) as c3 by gender",
+                TEST_INDEX_ACCOUNT)));
+  }
+
   @Test
   public void testExplainSortOnMetricsNoBucketNullable() throws IOException {
     // TODO enhancement later: https://github.com/opensearch-project/sql/issues/4282

integ-test/src/test/java/org/opensearch/sql/ppl/StatsCommandIT.java

@@ -772,4 +772,29 @@ public void testStatsBySpanTimeWithNullBucket() throws IOException {
         rows(8213, "2025-07-31 00:00:00"),
         rows(8490, "2025-07-31 12:00:00"));
   }
+
+  @Test
+  public void testStatsByCounts() throws IOException {
+    JSONObject response =
+        executeQuery(
+            String.format(
+                "source=%s | eval b_1 = balance + 1 | stats count(), count() as c1,"
+                    + " count(account_number), count(lastname) as c2, count(balance/10),"
+                    + " count(pow(balance, 2)) as c3, count(b_1) by gender",
+                TEST_INDEX_ACCOUNT));
+    verifySchema(
+        response,
+        schema("count()", null, isCalciteEnabled() ? "bigint" : "int"),
+        schema("c1", null, isCalciteEnabled() ? "bigint" : "int"),
+        schema("count(account_number)", null, isCalciteEnabled() ? "bigint" : "int"),
+        schema("c2", null, isCalciteEnabled() ? "bigint" : "int"),
+        schema("count(balance/10)", null, isCalciteEnabled() ? "bigint" : "int"),
+        schema("c3", null, isCalciteEnabled() ? "bigint" : "int"),
+        schema("count(b_1)", null, isCalciteEnabled() ? "bigint" : "int"),
+        schema("gender", null, "string"));
+    verifyDataRows(
+        response,
+        rows(493, 493, 493, 493, 493, 493, 493, "F"),
+        rows(507, 507, 507, 507, 507, 507, 507, "M"));
+  }
 }

integ-test/src/test/resources/expectedOutput/calcite/explain_agg_counts_by1.yaml

@@ -0,0 +1,10 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalProject(count()=[$1], c1=[$1], gender=[$0])
+        LogicalAggregate(group=[{0}], count()=[COUNT()])
+          LogicalProject(gender=[$4])
+            CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
+  physical: |
+    EnumerableCalc(expr#0..1=[{inputs}], count()=[$t0], count()0=[$t0], gender=[$t1])
+      CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},count()=COUNT()), LIMIT->10000, PROJECT->[count(), gender]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"gender":{"terms":{"field":"gender.keyword","missing_bucket":true,"missing_order":"first","order":"asc"}}}]}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])

integ-test/src/test/resources/expectedOutput/calcite/explain_agg_counts_by2.yaml

@@ -0,0 +1,10 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalProject(c1=[$1], c2=[$1], gender=[$0])
+        LogicalAggregate(group=[{0}], c1=[COUNT($1)])
+          LogicalProject(gender=[$4], balance=[$3])
+            CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
+  physical: |
+    EnumerableCalc(expr#0..1=[{inputs}], c1=[$t0], c10=[$t0], gender=[$t1])
+      CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},c1=COUNT($1)), LIMIT->10000, PROJECT->[c1, gender]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"gender":{"terms":{"field":"gender.keyword","missing_bucket":true,"missing_order":"first","order":"asc"}}}]}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])

integ-test/src/test/resources/expectedOutput/calcite/explain_agg_counts_by3.yaml

@@ -0,0 +1,9 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalProject(count(account_number)=[$1], c2=[$2], gender=[$0])
+        LogicalAggregate(group=[{0}], count(account_number)=[COUNT($1)], c2=[COUNT($2)])
+          LogicalProject(gender=[$4], account_number=[$0], account_number_alias=[$0])
+            CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
+  physical: |
+    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},count(account_number)=COUNT($1),c2=COUNT($2)), PROJECT->[count(account_number), c2, gender], LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"gender":{"terms":{"field":"gender.keyword","missing_bucket":true,"missing_order":"first","order":"asc"}}}]}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])

integ-test/src/test/resources/expectedOutput/calcite/explain_agg_counts_by4.yaml

@@ -0,0 +1,9 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalProject(count()=[$1], count(account_number)=[$2], gender=[$0])
+        LogicalAggregate(group=[{0}], count()=[COUNT()], count(account_number)=[COUNT($1)])
+          LogicalProject(gender=[$4], account_number=[$0])
+            CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
+  physical: |
+    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},count()=COUNT(),count(account_number)=COUNT($1)), PROJECT->[count(), count(account_number), gender], LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"gender":{"terms":{"field":"gender.keyword","missing_bucket":true,"missing_order":"first","order":"asc"}}}]},"aggregations":{"count()":{"value_count":{"field":"_index"}},"count(account_number)":{"value_count":{"field":"account_number"}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])

integ-test/src/test/resources/expectedOutput/calcite/explain_agg_counts_by5.yaml

@@ -0,0 +1,9 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalProject(count(balance)=[$1], count(account_number)=[$2], gender=[$0])
+        LogicalAggregate(group=[{0}], count(balance)=[COUNT($1)], count(account_number)=[COUNT($2)])
+          LogicalProject(gender=[$4], balance=[$3], account_number=[$0])
+            CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
+  physical: |
+    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},count(balance)=COUNT($1),count(account_number)=COUNT($2)), PROJECT->[count(balance), count(account_number), gender], LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"gender":{"terms":{"field":"gender.keyword","missing_bucket":true,"missing_order":"first","order":"asc"}}}]},"aggregations":{"count(balance)":{"value_count":{"field":"balance"}},"count(account_number)":{"value_count":{"field":"account_number"}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])