Add DISTINCT_COUNT_APPROX function

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

@@ -10,11 +10,6 @@
 import static org.apache.calcite.rex.RexWindowBounds.UNBOUNDED_PRECEDING;
 import static org.apache.calcite.rex.RexWindowBounds.following;
 import static org.apache.calcite.rex.RexWindowBounds.preceding;
-import static org.opensearch.sql.calcite.utils.CalciteToolsHelper.STDDEV_POP_NULLABLE;
-import static org.opensearch.sql.calcite.utils.CalciteToolsHelper.STDDEV_SAMP_NULLABLE;
-import static org.opensearch.sql.calcite.utils.CalciteToolsHelper.VAR_POP_NULLABLE;
-import static org.opensearch.sql.calcite.utils.CalciteToolsHelper.VAR_SAMP_NULLABLE;
-import static org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils.TransferUserDefinedAggFunction;
 
 import com.google.common.collect.ImmutableList;
 import java.util.ArrayList;
@@ -25,7 +20,6 @@
 import org.apache.calcite.rex.RexVisitorImpl;
 import org.apache.calcite.rex.RexWindowBound;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
-import org.apache.calcite.sql.type.ReturnTypes;
 import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.calcite.tools.RelBuilder;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
@@ -37,9 +31,8 @@
 import org.opensearch.sql.ast.tree.Relation;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.calcite.CalcitePlanContext;
-import org.opensearch.sql.calcite.udf.udaf.PercentileApproxFunction;
-import org.opensearch.sql.calcite.udf.udaf.TakeAggFunction;
 import org.opensearch.sql.expression.function.BuiltinFunctionName;
+import org.opensearch.sql.expression.function.PPLFuncImpTable;
 
 public interface PlanUtils {
 
@@ -232,56 +225,7 @@ static RelBuilder.AggCall makeAggCall(
       boolean distinct,
       RexNode field,
       List<RexNode> argList) {
-    switch (functionName) {
-      case MAX:
-        return context.relBuilder.max(field);
-      case MIN:
-        return context.relBuilder.min(field);
-      case AVG:
-        return context.relBuilder.avg(distinct, null, field);
-      case COUNT:
-        return context.relBuilder.count(
-            distinct, null, field == null ? ImmutableList.of() : ImmutableList.of(field));
-      case SUM:
-        return context.relBuilder.sum(distinct, null, field);
-        //            case MEAN:
-        //                throw new UnsupportedOperationException("MEAN is not supported in PPL");
-        //            case STDDEV:
-        //                return context.relBuilder.aggregateCall(SqlStdOperatorTable.STDDEV,
-        // field);
-      case VARSAMP:
-        return context.relBuilder.aggregateCall(VAR_SAMP_NULLABLE, field);
-      case VARPOP:
-        return context.relBuilder.aggregateCall(VAR_POP_NULLABLE, field);
-      case STDDEV_POP:
-        return context.relBuilder.aggregateCall(STDDEV_POP_NULLABLE, field);
-      case STDDEV_SAMP:
-        return context.relBuilder.aggregateCall(STDDEV_SAMP_NULLABLE, field);
-        //            case PERCENTILE_APPROX:
-        //                return
-        // context.relBuilder.aggregateCall(SqlStdOperatorTable.PERCENTILE_CONT, field);
-      case TAKE:
-        return TransferUserDefinedAggFunction(
-            TakeAggFunction.class,
-            "TAKE",
-            UserDefinedFunctionUtils.getReturnTypeInferenceForArray(),
-            List.of(field),
-            argList,
-            context.relBuilder);
-      case PERCENTILE_APPROX:
-        List<RexNode> newArgList = new ArrayList<>(argList);
-        newArgList.add(context.rexBuilder.makeFlag(field.getType().getSqlTypeName()));
-        return TransferUserDefinedAggFunction(
-            PercentileApproxFunction.class,
-            "percentile_approx",
-            ReturnTypes.ARG0_FORCE_NULLABLE,
-            List.of(field),
-            newArgList,
-            context.relBuilder);
-      default:
-        throw new UnsupportedOperationException(
-            "Unexpected aggregation: " + functionName.getName().getFunctionName());
-    }
+    return PPLFuncImpTable.INSTANCE.resolveAgg(functionName, distinct, field, argList, context);
   }
 
   /** Get all uniq input references from a RexNode. */

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

@@ -76,7 +76,7 @@ public static RelBuilder.AggCall TransferUserDefinedAggFunction(
     return relBuilder.aggregateCall(sqlUDAF, addArgList);
   }
 
-  static SqlReturnTypeInference getReturnTypeInferenceForArray() {
+  public static SqlReturnTypeInference getReturnTypeInferenceForArray() {
     return opBinding -> {
       RelDataTypeFactory typeFactory = opBinding.getTypeFactory();
 

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

@@ -195,6 +195,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")),
 
@@ -336,6 +337,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 =