refactor code

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

Author: xinyual

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

@@ -10,7 +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.udf.udaf.AggTransferFunctionMap.AGG_FUNCTION_MAP;
 
 import com.google.common.collect.ImmutableList;
 import java.util.ArrayList;
@@ -28,8 +27,8 @@
 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.AggHandler;
 import org.opensearch.sql.expression.function.BuiltinFunctionName;
+import org.opensearch.sql.expression.function.PPLFuncImpTable;
 
 public interface PlanUtils {
 
@@ -220,13 +219,7 @@ static RelBuilder.AggCall makeAggCall(
       boolean distinct,
       RexNode field,
       List<RexNode> argList) {
-    AggHandler handler = AGG_FUNCTION_MAP.get(functionName);
-
-    if (handler == null) {
-      throw new UnsupportedOperationException("Unexpected aggregation: " + functionName);
-    }
-
-    return handler.apply(distinct, field, argList, context);
+    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/udf/udaf/AggTransferFunctionMap.java

@@ -5,10 +5,14 @@
 
 package org.opensearch.sql.expression.function;
 
+import static org.opensearch.sql.calcite.utils.CalciteToolsHelper.*;
+import static org.opensearch.sql.calcite.utils.CalciteToolsHelper.STDDEV_POP_NULLABLE;
 import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.TYPE_FACTORY;
 import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.getLegacyTypeName;
+import static org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils.TransferUserDefinedAggFunction;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.*;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import java.math.BigDecimal;
 import java.util.ArrayList;
@@ -24,11 +28,23 @@
 import org.apache.calcite.sql.fun.SqlLibraryOperators;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
 import org.apache.calcite.sql.fun.SqlTrimFunction.Flag;
+import org.apache.calcite.sql.type.ReturnTypes;
 import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.calcite.tools.RelBuilder;
 import org.apache.commons.lang3.tuple.Pair;
+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.calcite.utils.UserDefinedFunctionUtils;
 import org.opensearch.sql.executor.QueryType;
 
 public class PPLFuncImpTable {
+  /** A lambda function interface which could apply parameters to get AggCall. */
+  @FunctionalInterface
+  public interface AggHandler {
+    RelBuilder.AggCall apply(
+        boolean distinct, RexNode field, List<RexNode> argList, CalcitePlanContext context);
+  }
 
   public interface FunctionImp {
     RelDataType ANY_TYPE = TYPE_FACTORY.createSqlType(SqlTypeName.ANY);
@@ -88,7 +104,9 @@ default List<RelDataType> getParams() {
   static {
     final Builder builder = new Builder();
     builder.populate();
-    INSTANCE = new PPLFuncImpTable(builder);
+    final AggBuilder aggBuilder = new AggBuilder();
+    aggBuilder.populate();
+    INSTANCE = new PPLFuncImpTable(builder, aggBuilder);
   }
 
   /**
@@ -107,12 +125,32 @@ default List<RelDataType> getParams() {
   private final Map<BuiltinFunctionName, List<Pair<CalciteFuncSignature, FunctionImp>>>
       externalFunctionRegistry;
 
-  private PPLFuncImpTable(Builder builder) {
+  /**
+   * The registry for built-in agg functions. Agg Functions defined by the PPL specification, whose
+   * implementations are independent of any specific data storage, should be registered here
+   * internally.
+   */
+  private final ImmutableMap<BuiltinFunctionName, AggHandler> aggFunctionRegistry;
+
+  /**
+   * The external agg function registry. Agg Functions whose implementations depend on a specific
+   * data engine should be registered here. This reduces coupling between the core module and
+   * particular storage backends.
+   */
+  private final Map<BuiltinFunctionName, AggHandler> aggExternalFunctionRegistry;
+
+  private PPLFuncImpTable(Builder builder, AggBuilder aggBuilder) {
     final ImmutableMap.Builder<BuiltinFunctionName, List<Pair<CalciteFuncSignature, FunctionImp>>>
         mapBuilder = ImmutableMap.builder();
     builder.map.forEach((k, v) -> mapBuilder.put(k, List.copyOf(v)));
     this.functionRegistry = ImmutableMap.copyOf(mapBuilder.build());
     this.externalFunctionRegistry = new HashMap<>();
+
+    final ImmutableMap.Builder<BuiltinFunctionName, AggHandler> aggMapBuilder =
+        ImmutableMap.builder();
+    aggBuilder.map.forEach(aggMapBuilder::put);
+    this.aggFunctionRegistry = ImmutableMap.copyOf(aggMapBuilder.build());
+    this.aggExternalFunctionRegistry = new HashMap<>();
   }
 
   /**
@@ -132,6 +170,33 @@ public void registerExternalFunction(BuiltinFunctionName functionName, FunctionI
     }
   }
 
+  /**
+   * Register a function implementation from external services dynamically.
+   *
+   * @param functionName the name of the function, has to be defined in BuiltinFunctionName
+   * @param functionImp the implementation of the agg function
+   */
+  public void registerExternalAggFunction(
+      BuiltinFunctionName functionName, AggHandler functionImp) {
+    aggExternalFunctionRegistry.put(functionName, functionImp);
+  }
+
+  public RelBuilder.AggCall resolveAgg(
+      BuiltinFunctionName functionName,
+      boolean distinct,
+      RexNode field,
+      List<RexNode> argList,
+      CalcitePlanContext context) {
+    AggHandler handler = aggExternalFunctionRegistry.get(functionName);
+    if (handler == null) {
+      handler = aggFunctionRegistry.get(functionName);
+    }
+    if (handler == null) {
+      throw new IllegalStateException(String.format("Cannot resolve function: %s", functionName));
+    }
+    return handler.apply(distinct, field, argList, context);
+  }
+
   public RexNode resolve(final RexBuilder builder, final String functionName, RexNode... args) {
     Optional<BuiltinFunctionName> funcNameOpt = BuiltinFunctionName.of(functionName);
     if (funcNameOpt.isEmpty()) {
@@ -462,4 +527,70 @@ public List<RelDataType> getParams() {
       return List.of(boolType, boolType);
     }
   }
+
+  private static class AggBuilder {
+    private final Map<BuiltinFunctionName, AggHandler> map = new HashMap<>();
+
+    void register(BuiltinFunctionName functionName, AggHandler aggHandler) {
+      map.put(functionName, aggHandler);
+    }
+
+    void populate() {
+      register(MAX, (distinct, field, argList, ctx) -> ctx.relBuilder.max(field));
+      register(MIN, (distinct, field, argList, ctx) -> ctx.relBuilder.min(field));
+
+      register(AVG, (distinct, field, argList, ctx) -> ctx.relBuilder.avg(distinct, null, field));
+
+      register(
+          COUNT,
+          (distinct, field, argList, ctx) ->
+              ctx.relBuilder.count(
+                  distinct, null, field == null ? ImmutableList.of() : ImmutableList.of(field)));
+      register(SUM, (distinct, field, argList, ctx) -> ctx.relBuilder.sum(distinct, null, field));
+
+      register(
+          VARSAMP,
+          (distinct, field, argList, ctx) ->
+              ctx.relBuilder.aggregateCall(VAR_SAMP_NULLABLE, field));
+
+      register(
+          VARPOP,
+          (distinct, field, argList, ctx) -> ctx.relBuilder.aggregateCall(VAR_POP_NULLABLE, field));
+
+      register(
+          STDDEV_SAMP,
+          (distinct, field, argList, ctx) ->
+              ctx.relBuilder.aggregateCall(STDDEV_SAMP_NULLABLE, field));
+
+      register(
+          STDDEV_POP,
+          (distinct, field, argList, ctx) ->
+              ctx.relBuilder.aggregateCall(STDDEV_POP_NULLABLE, field));
+
+      register(
+          TAKE,
+          (distinct, field, argList, ctx) ->
+              TransferUserDefinedAggFunction(
+                  TakeAggFunction.class,
+                  "TAKE",
+                  UserDefinedFunctionUtils.getReturnTypeInferenceForArray(),
+                  List.of(field),
+                  argList,
+                  ctx.relBuilder));
+
+      register(
+          PERCENTILE_APPROX,
+          (distinct, field, argList, ctx) -> {
+            List<RexNode> newArgList = new ArrayList<>(argList);
+            newArgList.add(ctx.rexBuilder.makeFlag(field.getType().getSqlTypeName()));
+            return TransferUserDefinedAggFunction(
+                PercentileApproxFunction.class,
+                "percentile_approx",
+                ReturnTypes.ARG0_FORCE_NULLABLE,
+                List.of(field),
+                newArgList,
+                ctx.relBuilder);
+          });
+    }
+  }
 }

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

@@ -30,8 +30,6 @@
 import org.apache.calcite.sql.type.ReturnTypes;
 import org.opensearch.sql.ast.statement.Explain.ExplainFormat;
 import org.opensearch.sql.calcite.CalcitePlanContext;
-import org.opensearch.sql.calcite.udf.udaf.AggTransferFunctionMap;
-import org.opensearch.sql.calcite.udf.udaf.AggTransferFunctionMap.*;
 import org.opensearch.sql.calcite.utils.CalciteToolsHelper.OpenSearchRelRunners;
 import org.opensearch.sql.common.response.ResponseListener;
 import org.opensearch.sql.data.model.ExprTupleValue;
@@ -253,7 +251,7 @@ private void registerOpenSearchFunctions() {
             builder.makeCall(new GeoIpFunction(client.getNodeClient()).toUDF("GEOIP"), args);
     PPLFuncImpTable.INSTANCE.registerExternalFunction(BuiltinFunctionName.GEOIP, geoIpImpl);
 
-    AggTransferFunctionMap.AGG_FUNCTION_MAP.put(
+    PPLFuncImpTable.INSTANCE.registerExternalAggFunction(
         DISTINCT_COUNT_APPROX,
         (distinct, field, argList, ctx) ->
             TransferUserDefinedAggFunction(