add approx distinct count

Signed-off-by: xinyual <xinyual@amazon.com>

Author: xinyual

core/build.gradle

@@ -63,6 +63,7 @@ dependencies {
     api 'org.apache.calcite:calcite-linq4j:1.38.0'
     api project(':common')
     implementation "com.github.seancfoley:ipaddress:5.4.2"
+    implementation "com.google.zetasketch:zetasketch:0.1.0"
 
     annotationProcessor('org.immutables:value:2.8.8')
     compileOnly('org.immutables:value-annotations:2.8.8')

core/src/main/java/org/opensearch/sql/calcite/udf/udaf/DistinctCountApproxAggFunction.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.calcite.udf.udaf;
+
+import com.google.zetasketch.HyperLogLogPlusPlus;
+import org.opensearch.sql.calcite.udf.UserDefinedAggFunction;
+
+public class DistinctCountApproxAggFunction
+    implements UserDefinedAggFunction<DistinctCountApproxAggFunction.HLLAccumulator> {
+
+  @Override
+  public HLLAccumulator init() {
+    return new HLLAccumulator();
+  }
+
+  @Override
+  public Object result(HLLAccumulator accumulator) {
+    return accumulator.value();
+  }
+
+  @Override
+  public HLLAccumulator add(HLLAccumulator acc, Object... values) {
+    for (Object value : values) {
+      if (value != null) {
+        acc.add(value.toString());
+      }
+    }
+    return acc;
+  }
+
+  public static class HLLAccumulator implements UserDefinedAggFunction.Accumulator {
+    private final HyperLogLogPlusPlus<String> hll;
+
+    public HLLAccumulator() {
+      this.hll = new HyperLogLogPlusPlus.Builder().buildForStrings();
+    }
+
+    public void add(String value) {
+      hll.add(value);
+    }
+
+    @Override
+    public Object value(Object... args) {
+        return hll.result();
+    }
+  }
+}

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

@@ -33,6 +33,7 @@
 import org.opensearch.sql.ast.expression.WindowBound;
 import org.opensearch.sql.ast.expression.WindowFrame;
 import org.opensearch.sql.calcite.CalcitePlanContext;
+import org.opensearch.sql.calcite.udf.udaf.DistinctCountApproxAggFunction;
 import org.opensearch.sql.calcite.udf.udaf.PercentileApproxFunction;
 import org.opensearch.sql.calcite.udf.udaf.TakeAggFunction;
 import org.opensearch.sql.expression.function.BuiltinFunctionName;
@@ -218,6 +219,8 @@ static RelBuilder.AggCall makeAggCall(
       case AVG:
         return context.relBuilder.avg(distinct, null, field);
       case COUNT:
+        // return context.relBuilder.aggregateCall(SqlStdOperatorTable.APPROX_COUNT_DISTINCT,
+        // field);
         return context.relBuilder.count(
             distinct, null, field == null ? ImmutableList.of() : ImmutableList.of(field));
       case SUM:
@@ -238,6 +241,14 @@ static RelBuilder.AggCall makeAggCall(
         //            case PERCENTILE_APPROX:
         //                return
         // context.relBuilder.aggregateCall(SqlStdOperatorTable.PERCENTILE_CONT, field);
+      case DISTINCT_COUNT_APPROX:
+        return TransferUserDefinedAggFunction(
+            DistinctCountApproxAggFunction.class,
+            "APPROX_DISTINCT_COUNT",
+            ReturnTypes.BIGINT_FORCE_NULLABLE,
+            List.of(field),
+            argList,
+            context.relBuilder);
       case TAKE:
         return TransferUserDefinedAggFunction(
             TakeAggFunction.class,

core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java

@@ -188,6 +188,7 @@ public enum BuiltinFunctionName {
   TAKE(FunctionName.of("take")),
   // t-digest percentile which is used in OpenSearch core by default.
   PERCENTILE_APPROX(FunctionName.of("percentile_approx")),
+  DISTINCT_COUNT_APPROX(FunctionName.of("distinct_count_approx")),
   // Not always an aggregation query
   NESTED(FunctionName.of("nested")),
 
@@ -310,6 +311,7 @@ public enum BuiltinFunctionName {
           .put("take", BuiltinFunctionName.TAKE)
           .put("percentile", BuiltinFunctionName.PERCENTILE_APPROX)
           .put("percentile_approx", BuiltinFunctionName.PERCENTILE_APPROX)
+          .put("distinct_count_approx", BuiltinFunctionName.DISTINCT_COUNT_APPROX)
           .build();
 
   private static final Map<String, BuiltinFunctionName> WINDOW_FUNC_MAPPING =

integ-test/src/test/java/org/opensearch/sql/calcite/standalone/CalcitePPLAggregationIT.java

@@ -514,6 +514,17 @@ public void testCountDistinct() {
     verifyDataRows(actual, rows(3, "F"), rows(4, "M"));
   }
 
+  @Test
+  public void testCountDistinctApprox() {
+    JSONObject actual =
+        executeQuery(
+            String.format(
+                "source=%s | stats distinct_count_approx(state) by gender", TEST_INDEX_BANK));
+    verifySchema(
+        actual, schema("gender", "string"), schema("distinct_count_approx(state)", "long"));
+    verifyDataRows(actual, rows(3, "F"), rows(4, "M"));
+  }
+
   @Test
   public void testCountDistinctWithAlias() {
     JSONObject actual =

ppl/src/main/antlr/OpenSearchPPLLexer.g4

@@ -203,6 +203,7 @@ BIT_XOR_OP:                         '^';
 AVG:                                'AVG';
 COUNT:                              'COUNT';
 DISTINCT_COUNT:                     'DISTINCT_COUNT';
+DISTINCT_COUNT_APPROX:              'DISTINCT_COUNT_APPROX';
 ESTDC:                              'ESTDC';
 ESTDC_ERROR:                        'ESTDC_ERROR';
 MAX:                                'MAX';

ppl/src/main/antlr/OpenSearchPPLParser.g4

@@ -389,6 +389,7 @@ statsFunction
    : statsFunctionName LT_PRTHS valueExpression RT_PRTHS        # statsFunctionCall
    | COUNT LT_PRTHS RT_PRTHS                                    # countAllFunctionCall
    | (DISTINCT_COUNT | DC) LT_PRTHS valueExpression RT_PRTHS    # distinctCountFunctionCall
+   | DISTINCT_COUNT_APPROX LT_PRTHS valueExpression RT_PRTHS    # distinctCountApproxFunctionCall
    | takeAggFunction                                            # takeAggFunctionCall
    | percentileApproxFunction                                   # percentileApproxFunctionCall
    ;
@@ -1096,6 +1097,7 @@ keywordsCanBeId
    | statsFunctionName
    | windowFunctionName
    | DISTINCT_COUNT
+   | DISTINCT_COUNT_APPROX
    | ESTDC
    | ESTDC_ERROR
    | MEAN

ppl/src/main/java/org/opensearch/sql/ppl/parser/AstExpressionBuilder.java

@@ -203,6 +203,12 @@ public UnresolvedExpression visitDistinctCountFunctionCall(DistinctCountFunction
     return new AggregateFunction("count", visit(ctx.valueExpression()), true);
   }
 
+  @Override
+  public UnresolvedExpression visitDistinctCountApproxFunctionCall(
+      OpenSearchPPLParser.DistinctCountApproxFunctionCallContext ctx) {
+    return new AggregateFunction("distinct_count_approx", visit(ctx.valueExpression()), true);
+  }
+
   @Override
   public UnresolvedExpression visitPercentileApproxFunctionCall(
       OpenSearchPPLParser.PercentileApproxFunctionCallContext ctx) {