pushdown approx distinct count

xinyual · xinyual · commit 45b49f215739 · 2025-10-22T10:14:34.000+08:00
Signed-off-by: xinyual &lt;xinyual@amazon.com&gt;
diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java
@@ -620,6 +620,19 @@ public void testExplainOnAggregationWithSumEnhancement() throws IOException {
                 TEST_INDEX_BANK)));
   }
 
+  @Test
+  public void testStatsDistinctCountApproxFunctionExplainWithPushDown() throws IOException {
+    enabledOnlyWhenPushdownIsEnabled();
+    String query =
+        "source=opensearch-sql_test_index_account | stats distinct_count_approx(state) as"
+            + " distinct_states by gender";
+    var result = explainQueryToString(query);
+    String expected =
+        loadFromFile(
+            "expectedOutput/calcite/explain_agg_with_distinct_count_approx_enhancement.json");
+    assertJsonEqualsIgnoreId(expected, result);
+  }
+
   @Test
   public void testExplainRegexMatchInWhereWithScriptPushdown() throws IOException {
     enabledOnlyWhenPushdownIsEnabled();
diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_agg_with_distinct_count_approx_enhancement.json b/integ-test/src/test/resources/expectedOutput/calcite/explain_agg_with_distinct_count_approx_enhancement.json
@@ -0,0 +1,6 @@
+{
+  "calcite":{
+    "logical":"LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])\n  LogicalProject(distinct_states=[$1], gender=[$0])\n    LogicalAggregate(group=[{0}], distinct_states=[DISTINCT_COUNT_APPROX($1)])\n      LogicalProject(gender=[$4], state=[$7])\n        CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",
+    "physical":"CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},distinct_states=DISTINCT_COUNT_APPROX($1)), PROJECT->[distinct_states, 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\":{\"distinct_states\":{\"cardinality\":{\"field\":\"state.keyword\"}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])\n"
+  }
+}
diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_eventstats_dc.json b/integ-test/src/test/resources/expectedOutput/calcite/explain_eventstats_dc.json
@@ -1,6 +1,6 @@
 {
   "calcite": {
-    "logical": "LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])\n  LogicalProject(account_number=[$0], firstname=[$1], address=[$2], balance=[$3], gender=[$4], city=[$5], employer=[$6], state=[$7], age=[$8], email=[$9], lastname=[$10], distinct_states=[APPROX_DISTINCT_COUNT($7) OVER ()])\n    CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",
-    "physical": "EnumerableLimit(fetch=[10000])\n  EnumerableWindow(window#0=[window(aggs [APPROX_DISTINCT_COUNT($7)])])\n    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[account_number, firstname, address, balance, gender, city, employer, state, age, email, lastname]], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"account_number\",\"firstname\",\"address\",\"balance\",\"gender\",\"city\",\"employer\",\"state\",\"age\",\"email\",\"lastname\"],\"excludes\":[]}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])\n"
+    "logical": "LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])\n  LogicalProject(account_number=[$0], firstname=[$1], address=[$2], balance=[$3], gender=[$4], city=[$5], employer=[$6], state=[$7], age=[$8], email=[$9], lastname=[$10], distinct_states=[DISTINCT_COUNT_APPROX($7) OVER ()])\n    CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",
+    "physical": "EnumerableLimit(fetch=[10000])\n  EnumerableWindow(window#0=[window(aggs [DISTINCT_COUNT_APPROX($7)])])\n    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[account_number, firstname, address, balance, gender, city, employer, state, age, email, lastname]], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"account_number\",\"firstname\",\"address\",\"balance\",\"gender\",\"city\",\"employer\",\"state\",\"age\",\"email\",\"lastname\"],\"excludes\":[]}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])\n"
   }
 }
diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_eventstats_distinct_count.json b/integ-test/src/test/resources/expectedOutput/calcite/explain_eventstats_distinct_count.json
@@ -1,6 +1,6 @@
 {
   "calcite": {
-    "logical": "LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])\n  LogicalProject(account_number=[$0], firstname=[$1], address=[$2], balance=[$3], gender=[$4], city=[$5], employer=[$6], state=[$7], age=[$8], email=[$9], lastname=[$10], distinct_states=[APPROX_DISTINCT_COUNT($7) OVER (PARTITION BY $4)])\n    CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",
-    "physical": "EnumerableLimit(fetch=[10000])\n  EnumerableWindow(window#0=[window(partition {4} aggs [APPROX_DISTINCT_COUNT($7)])])\n    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[account_number, firstname, address, balance, gender, city, employer, state, age, email, lastname]], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"account_number\",\"firstname\",\"address\",\"balance\",\"gender\",\"city\",\"employer\",\"state\",\"age\",\"email\",\"lastname\"],\"excludes\":[]}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])\n"
+    "logical": "LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])\n  LogicalProject(account_number=[$0], firstname=[$1], address=[$2], balance=[$3], gender=[$4], city=[$5], employer=[$6], state=[$7], age=[$8], email=[$9], lastname=[$10], distinct_states=[DISTINCT_COUNT_APPROX($7) OVER (PARTITION BY $4)])\n    CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",
+    "physical": "EnumerableLimit(fetch=[10000])\n  EnumerableWindow(window#0=[window(partition {4} aggs [DISTINCT_COUNT_APPROX($7)])])\n    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[account_number, firstname, address, balance, gender, city, employer, state, age, email, lastname]], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"account_number\",\"firstname\",\"address\",\"balance\",\"gender\",\"city\",\"employer\",\"state\",\"age\",\"email\",\"lastname\"],\"excludes\":[]}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])\n"
   }
 }
diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java
@@ -5,6 +5,8 @@
 
 package org.opensearch.sql.opensearch.executor;
 
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.DISTINCT_COUNT_APPROX;
+
 import java.security.AccessController;
 import java.security.PrivilegedAction;
 import java.sql.PreparedStatement;
@@ -292,10 +294,10 @@ private void registerOpenSearchFunctions() {
     SqlUserDefinedAggFunction approxDistinctCountFunction =
         UserDefinedFunctionUtils.createUserDefinedAggFunction(
             DistinctCountApproxAggFunction.class,
-            "APPROX_DISTINCT_COUNT",
+            DISTINCT_COUNT_APPROX.toString(),
             ReturnTypes.BIGINT_FORCE_NULLABLE,
             null);
     PPLFuncImpTable.INSTANCE.registerExternalAggOperator(
-        BuiltinFunctionName.DISTINCT_COUNT_APPROX, approxDistinctCountFunction);
+        DISTINCT_COUNT_APPROX, approxDistinctCountFunction);
   }
 }
diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/AggregateAnalyzer.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/AggregateAnalyzer.java
@@ -382,141 +382,148 @@ private static Pair<AggregationBuilder, MetricParser> createRegularAggregation(
     switch (aggCall.getAggregation().kind) {
       case AVG:
         return Pair.of(
-            helper.build(args.get(0), AggregationBuilders.avg(aggFieldName)),
-            new SingleValueParser(aggFieldName));
+                helper.build(args.get(0), AggregationBuilders.avg(aggFieldName)),
+                new SingleValueParser(aggFieldName));
       case SUM:
         // 1. Only case SUM, skip SUM0 / COUNT since calling avg() in DSL should be faster.
         // 2. To align with databases, SUM0 is not preferred now.
         return Pair.of(
-            helper.build(args.get(0), AggregationBuilders.sum(aggFieldName)),
-            new SingleValueParser(aggFieldName));
+                helper.build(args.get(0), AggregationBuilders.sum(aggFieldName)),
+                new SingleValueParser(aggFieldName));
       case COUNT:
         return Pair.of(
-            helper.build(
-                !args.isEmpty() ? args.get(0) : null, AggregationBuilders.count(aggFieldName)),
-            new SingleValueParser(aggFieldName));
+                helper.build(
+                        !args.isEmpty() ? args.get(0) : null, AggregationBuilders.count(aggFieldName)),
+                new SingleValueParser(aggFieldName));
       case MIN: {
         ExprType fieldType =
-            OpenSearchTypeFactory.convertRelDataTypeToExprType(args.get(0).getType());
+                OpenSearchTypeFactory.convertRelDataTypeToExprType(args.get(0).getType());
         if (supportsMaxMinAggregation(fieldType)) {
           return Pair.of(
-              helper.build(args.get(0), AggregationBuilders.min(aggFieldName)),
-              new SingleValueParser(aggFieldName));
+                  helper.build(args.get(0), AggregationBuilders.min(aggFieldName)),
+                  new SingleValueParser(aggFieldName));
         } else {
           return Pair.of(
-              AggregationBuilders.topHits(aggFieldName)
-                  .fetchSource(helper.inferNamedField(args.get(0)).getRootName(), null)
-                  .size(1)
-                  .from(0)
-                  .sort(
-                      helper.inferNamedField(args.get(0)).getReferenceForTermQuery(),
-                      SortOrder.ASC),
-              new TopHitsParser(aggFieldName, true));
+                  AggregationBuilders.topHits(aggFieldName)
+                          .fetchSource(helper.inferNamedField(args.get(0)).getRootName(), null)
+                          .size(1)
+                          .from(0)
+                          .sort(
+                                  helper.inferNamedField(args.get(0)).getReferenceForTermQuery(),
+                                  SortOrder.ASC),
+                  new TopHitsParser(aggFieldName, true));
         }
       }
       case MAX: {
         ExprType fieldType =
-            OpenSearchTypeFactory.convertRelDataTypeToExprType(args.get(0).getType());
+                OpenSearchTypeFactory.convertRelDataTypeToExprType(args.get(0).getType());
         if (supportsMaxMinAggregation(fieldType)) {
           return Pair.of(
-              helper.build(args.get(0), AggregationBuilders.max(aggFieldName)),
-              new SingleValueParser(aggFieldName));
+                  helper.build(args.get(0), AggregationBuilders.max(aggFieldName)),
+                  new SingleValueParser(aggFieldName));
         } else {
           return Pair.of(
-              AggregationBuilders.topHits(aggFieldName)
-                  .fetchSource(helper.inferNamedField(args.get(0)).getRootName(), null)
-                  .size(1)
-                  .from(0)
-                  .sort(
-                      helper.inferNamedField(args.get(0)).getReferenceForTermQuery(),
-                      SortOrder.DESC),
-              new TopHitsParser(aggFieldName, true));
+                  AggregationBuilders.topHits(aggFieldName)
+                          .fetchSource(helper.inferNamedField(args.get(0)).getRootName(), null)
+                          .size(1)
+                          .from(0)
+                          .sort(
+                                  helper.inferNamedField(args.get(0)).getReferenceForTermQuery(),
+                                  SortOrder.DESC),
+                  new TopHitsParser(aggFieldName, true));
         }
       }
       case VAR_SAMP:
         return Pair.of(
-            helper.build(args.get(0), AggregationBuilders.extendedStats(aggFieldName)),
-            new StatsParser(ExtendedStats::getVarianceSampling, aggFieldName));
+                helper.build(args.get(0), AggregationBuilders.extendedStats(aggFieldName)),
+                new StatsParser(ExtendedStats::getVarianceSampling, aggFieldName));
       case VAR_POP:
         return Pair.of(
-            helper.build(args.get(0), AggregationBuilders.extendedStats(aggFieldName)),
-            new StatsParser(ExtendedStats::getVariancePopulation, aggFieldName));
+                helper.build(args.get(0), AggregationBuilders.extendedStats(aggFieldName)),
+                new StatsParser(ExtendedStats::getVariancePopulation, aggFieldName));
       case STDDEV_SAMP:
         return Pair.of(
-            helper.build(args.get(0), AggregationBuilders.extendedStats(aggFieldName)),
-            new StatsParser(ExtendedStats::getStdDeviationSampling, aggFieldName));
+                helper.build(args.get(0), AggregationBuilders.extendedStats(aggFieldName)),
+                new StatsParser(ExtendedStats::getStdDeviationSampling, aggFieldName));
       case STDDEV_POP:
         return Pair.of(
-            helper.build(args.get(0), AggregationBuilders.extendedStats(aggFieldName)),
-            new StatsParser(ExtendedStats::getStdDeviationPopulation, aggFieldName));
+                helper.build(args.get(0), AggregationBuilders.extendedStats(aggFieldName)),
+                new StatsParser(ExtendedStats::getStdDeviationPopulation, aggFieldName));
       case ARG_MAX:
         return Pair.of(
-          AggregationBuilders.topHits(aggFieldName)
-              .fetchSource(helper.inferNamedField(args.get(0)).getRootName(), null)
-              .size(1)
-              .from(0)
-              .sort(
-                  helper.inferNamedField(args.get(1)).getRootName(),
-                  org.opensearch.search.sort.SortOrder.DESC),
-          new ArgMaxMinParser(aggFieldName));
+                AggregationBuilders.topHits(aggFieldName)
+                        .fetchSource(helper.inferNamedField(args.get(0)).getRootName(), null)
+                        .size(1)
+                        .from(0)
+                        .sort(
+                                helper.inferNamedField(args.get(1)).getRootName(),
+                                org.opensearch.search.sort.SortOrder.DESC),
+                new ArgMaxMinParser(aggFieldName));
       case ARG_MIN:
         return Pair.of(
-          AggregationBuilders.topHits(aggFieldName)
-              .fetchSource(helper.inferNamedField(args.get(0)).getRootName(), null)
-              .size(1)
-              .from(0)
-              .sort(
-                  helper.inferNamedField(args.get(1)).getRootName(),
-                  org.opensearch.search.sort.SortOrder.ASC),
-          new ArgMaxMinParser(aggFieldName));
+                AggregationBuilders.topHits(aggFieldName)
+                        .fetchSource(helper.inferNamedField(args.get(0)).getRootName(), null)
+                        .size(1)
+                        .from(0)
+                        .sort(
+                                helper.inferNamedField(args.get(1)).getRootName(),
+                                org.opensearch.search.sort.SortOrder.ASC),
+                new ArgMaxMinParser(aggFieldName));
       case OTHER_FUNCTION:
         BuiltinFunctionName functionName =
-            BuiltinFunctionName.ofAggregation(aggCall.getAggregation().getName()).get();
+                BuiltinFunctionName.ofAggregation(aggCall.getAggregation().getName()).get();
         switch (functionName) {
           case TAKE:
             return Pair.of(
-                AggregationBuilders.topHits(aggFieldName)
-                    .fetchSource(helper.inferNamedField(args.get(0)).getRootName(), null)
-                    .size(helper.inferValue(args.get(1), Integer.class))
-                    .from(0),
-                new TopHitsParser(aggFieldName));
+                    AggregationBuilders.topHits(aggFieldName)
+                            .fetchSource(helper.inferNamedField(args.get(0)).getRootName(), null)
+                            .size(helper.inferValue(args.get(1), Integer.class))
+                            .from(0),
+                    new TopHitsParser(aggFieldName));
           case FIRST:
             TopHitsAggregationBuilder firstBuilder =
-                AggregationBuilders.topHits(aggFieldName).size(1).from(0);
+                    AggregationBuilders.topHits(aggFieldName).size(1).from(0);
             if (!args.isEmpty()) {
               firstBuilder.fetchSource(helper.inferNamedField(args.get(0)).getRootName(), null);
             }
             return Pair.of(firstBuilder, new TopHitsParser(aggFieldName, true));
           case LAST:
             TopHitsAggregationBuilder lastBuilder =
-                AggregationBuilders.topHits(aggFieldName)
-                    .size(1)
-                    .from(0)
-                    .sort("_doc", org.opensearch.search.sort.SortOrder.DESC);
+                    AggregationBuilders.topHits(aggFieldName)
+                            .size(1)
+                            .from(0)
+                            .sort("_doc", org.opensearch.search.sort.SortOrder.DESC);
             if (!args.isEmpty()) {
               lastBuilder.fetchSource(helper.inferNamedField(args.get(0)).getRootName(), null);
             }
             return Pair.of(lastBuilder, new TopHitsParser(aggFieldName, true));
           case PERCENTILE_APPROX:
             PercentilesAggregationBuilder aggBuilder =
-                helper
-                    .build(args.get(0), AggregationBuilders.percentiles(aggFieldName))
-                    .percentiles(helper.inferValue(args.get(1), Double.class));
+                    helper
+                            .build(args.get(0), AggregationBuilders.percentiles(aggFieldName))
+                            .percentiles(helper.inferValue(args.get(1), Double.class));
             /* See {@link PercentileApproxFunction}, PERCENTILE_APPROX accepts args of [FIELD, PERCENTILE, TYPE, COMPRESSION(optional)] */
             if (args.size() > 3) {
               aggBuilder.compression(helper.inferValue(args.get(3), Double.class));
             }
             return Pair.of(aggBuilder, new SinglePercentileParser(aggFieldName));
+          case DISTINCT_COUNT_APPROX:
+            return Pair.of(
+                    helper.build(
+                            !args.isEmpty() ? args.getFirst() : null,
+                            AggregationBuilders.cardinality(aggFieldName)),
+                    new SingleValueParser(aggFieldName));
           default:
             throw new AggregateAnalyzer.AggregateAnalyzerException(
-                String.format("Unsupported push-down aggregator %s", aggCall.getAggregation()));
+                    String.format("Unsupported push-down aggregator %s", aggCall.getAggregation()));
         }
       default:
         throw new AggregateAnalyzer.AggregateAnalyzerException(
-            String.format("unsupported aggregator %s", aggCall.getAggregation()));
+                String.format("unsupported aggregator %s", aggCall.getAggregation()));
     }
   }
 
+
   private static boolean supportsMaxMinAggregation(ExprType fieldType) {
     ExprType coreType =
         (fieldType instanceof OpenSearchDataType)

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`{`
`2`	`2`	`"calcite": {`
`3`		`- "logical": "LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])\n LogicalProject(account_number=[$0], firstname=[$1], address=[$2], balance=[$3], gender=[$4], city=[$5], employer=[$6], state=[$7], age=[$8], email=[$9], lastname=[$10], distinct_states=[APPROX_DISTINCT_COUNT($7) OVER ()])\n CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",`
`4`		- "physical": "EnumerableLimit(fetch=[10000])\n EnumerableWindow(window#0=[window(aggs [APPROX_DISTINCT_COUNT($7)])])\n CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[account_number, firstname, address, balance, gender, city, employer, state, age, email, lastname]], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"account_number\",\"firstname\",\"address\",\"balance\",\"gender\",\"city\",\"employer\",\"state\",\"age\",\"email\",\"lastname\"],\"excludes\":[]}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])\n"
	`3`	`+ "logical": "LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])\n LogicalProject(account_number=[$0], firstname=[$1], address=[$2], balance=[$3], gender=[$4], city=[$5], employer=[$6], state=[$7], age=[$8], email=[$9], lastname=[$10], distinct_states=[DISTINCT_COUNT_APPROX($7) OVER ()])\n CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",`
	`4`	+ "physical": "EnumerableLimit(fetch=[10000])\n EnumerableWindow(window#0=[window(aggs [DISTINCT_COUNT_APPROX($7)])])\n CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[account_number, firstname, address, balance, gender, city, employer, state, age, email, lastname]], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"account_number\",\"firstname\",\"address\",\"balance\",\"gender\",\"city\",\"employer\",\"state\",\"age\",\"email\",\"lastname\"],\"excludes\":[]}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])\n"
`5`	`5`	`}`
`6`	`6`	`}`