Move geoip function to opensearch package & register it in OpenSearchExecutionEngine

Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>

Author: yuancu

core/build.gradle

@@ -63,8 +63,6 @@ dependencies {
     api 'org.apache.calcite:calcite-linq4j:1.38.0'
     api project(':common')
     implementation "com.github.seancfoley:ipaddress:5.4.2"
-    implementation group: 'org.opensearch', name:'geospatial-client', version: "${opensearch_build}"
-    implementation group: 'org.opensearch.client', name: 'opensearch-rest-high-level-client', version: "${opensearch_version}"
 
     annotationProcessor('org.immutables:value:2.8.8')
     compileOnly('org.immutables:value-annotations:2.8.8')

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

@@ -21,7 +21,6 @@
 import org.opensearch.sql.calcite.utils.CalciteToolsHelper;
 import org.opensearch.sql.executor.QueryType;
 import org.opensearch.sql.expression.function.FunctionProperties;
-import org.opensearch.transport.client.node.NodeClient;
 
 public class CalcitePlanContext {
 
@@ -45,17 +44,13 @@ public class CalcitePlanContext {
 
   private final Stack<RexCorrelVariable> correlVar = new Stack<>();
 
-  /** Certain functions (e.g., GEOIP) depend on NodeClient for performing RPC calls. * */
-  @Getter private final NodeClient client;
-
-  private CalcitePlanContext(FrameworkConfig config, QueryType queryType, NodeClient client) {
+  private CalcitePlanContext(FrameworkConfig config, QueryType queryType) {
     this.config = config;
     this.queryType = queryType;
     this.connection = CalciteToolsHelper.connect(config, TYPE_FACTORY);
     this.relBuilder = CalciteToolsHelper.create(config, TYPE_FACTORY, connection);
     this.rexBuilder = new ExtendedRexBuilder(relBuilder.getRexBuilder());
     this.functionProperties = new FunctionProperties(QueryType.PPL);
-    this.client = client;
   }
 
   public RexNode resolveJoinCondition(
@@ -87,12 +82,7 @@ public Optional<RexCorrelVariable> peekCorrelVar() {
     }
   }
 
-  public static CalcitePlanContext create(
-      FrameworkConfig config, QueryType queryType, NodeClient client) {
-    return new CalcitePlanContext(config, queryType, client);
-  }
-
   public static CalcitePlanContext create(FrameworkConfig config, QueryType queryType) {
-    return new CalcitePlanContext(config, queryType, null);
+    return new CalcitePlanContext(config, queryType);
   }
 }

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

@@ -65,7 +65,6 @@
 import org.opensearch.sql.exception.SemanticCheckException;
 import org.opensearch.sql.expression.function.BuiltinFunctionName;
 import org.opensearch.sql.expression.function.PPLFuncImpTable;
-import org.opensearch.sql.expression.function.udf.ip.GeoIpFunction;
 
 @RequiredArgsConstructor
 public class CalciteRexNodeVisitor extends AbstractNodeVisitor<RexNode, CalcitePlanContext> {
@@ -349,13 +348,6 @@ public RexNode visitFunction(Function node, CalcitePlanContext context) {
     List<RexNode> arguments =
         node.getFuncArgs().stream().map(arg -> analyze(arg, context)).collect(Collectors.toList());
 
-    // GEOIP needs NodeClient to perform RPC calls. Therefore, we handle it separately.
-    if (BuiltinFunctionName.GEOIP.equals(BuiltinFunctionName.of(node.getFuncName()).orElse(null))) {
-      PPLFuncImpTable.FunctionImp geoIpImpl =
-          (builder, args) ->
-              builder.makeCall(new GeoIpFunction(context.getClient()).toUDF("GEOIP"), args);
-      return geoIpImpl.resolve(context.rexBuilder, arguments.toArray(new RexNode[0]));
-    }
     RexNode resolvedNode =
         PPLFuncImpTable.INSTANCE.resolveSafe(
             context.rexBuilder, node.getFuncName(), arguments.toArray(new RexNode[0]));

core/src/main/java/org/opensearch/sql/executor/QueryService.java

@@ -44,7 +44,6 @@
 import org.opensearch.sql.planner.Planner;
 import org.opensearch.sql.planner.logical.LogicalPlan;
 import org.opensearch.sql.planner.physical.PhysicalPlan;
-import org.opensearch.transport.client.node.NodeClient;
 
 /** The low level interface of core engine. */
 @RequiredArgsConstructor
@@ -60,7 +59,6 @@ public class QueryService {
 
   private DataSourceService dataSourceService;
   private Settings settings;
-  private NodeClient nodeClient;
 
   /** Execute the {@link UnresolvedPlan}, using {@link ResponseListener} to get response.<br> */
   public void execute(
@@ -96,7 +94,7 @@ public void executeWithCalcite(
           (PrivilegedAction<Void>)
               () -> {
                 CalcitePlanContext context =
-                    CalcitePlanContext.create(buildFrameworkConfig(), queryType, nodeClient);
+                    CalcitePlanContext.create(buildFrameworkConfig(), queryType);
                 RelNode relNode = analyze(plan, context);
                 RelNode optimized = optimize(relNode);
                 RelNode calcitePlan = convertToCalcitePlan(optimized);
@@ -128,7 +126,7 @@ public void explainWithCalcite(
           (PrivilegedAction<Void>)
               () -> {
                 CalcitePlanContext context =
-                    CalcitePlanContext.create(buildFrameworkConfig(), queryType, nodeClient);
+                    CalcitePlanContext.create(buildFrameworkConfig(), queryType);
                 RelNode relNode = analyze(plan, context);
                 RelNode optimized = optimize(relNode);
                 RelNode calcitePlan = convertToCalcitePlan(optimized);

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

@@ -152,13 +152,29 @@ default List<RelDataType> getParams() {
     INSTANCE = new PPLFuncImpTable(builder);
   }
 
-  private final ImmutableMap<BuiltinFunctionName, PairList<CalciteFuncSignature, FunctionImp>> map;
+  private final ImmutableMap<BuiltinFunctionName, PairList<CalciteFuncSignature, FunctionImp>>
+      functionRegistry;
+  private final Map<BuiltinFunctionName, PairList<CalciteFuncSignature, FunctionImp>>
+      externalFunctionRegistry;
 
   private PPLFuncImpTable(Builder builder) {
     final ImmutableMap.Builder<BuiltinFunctionName, PairList<CalciteFuncSignature, FunctionImp>>
         mapBuilder = ImmutableMap.builder();
     builder.map.forEach((k, v) -> mapBuilder.put(k, v.immutable()));
-    this.map = ImmutableMap.copyOf(mapBuilder.build());
+    this.functionRegistry = ImmutableMap.copyOf(mapBuilder.build());
+    this.externalFunctionRegistry = new HashMap<>();
+  }
+
+  /**
+   * 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 function
+   */
+  public void registerExternalFunction(BuiltinFunctionName functionName, FunctionImp functionImp) {
+    CalciteFuncSignature signature =
+        new CalciteFuncSignature(functionName.getName(), functionImp.getParams());
+    externalFunctionRegistry.put(functionName, PairList.of(signature, functionImp));
   }
 
   public @Nullable RexNode resolveSafe(
@@ -180,7 +196,11 @@ public RexNode resolve(final RexBuilder builder, final String functionName, RexN
 
   public RexNode resolve(
       final RexBuilder builder, final BuiltinFunctionName functionName, RexNode... args) {
-    final PairList<CalciteFuncSignature, FunctionImp> implementList = map.get(functionName);
+    PairList<CalciteFuncSignature, FunctionImp> implementList = functionRegistry.get(functionName);
+    // If the function is not part of the built-in registry, check the external registry
+    if (implementList == null) {
+      implementList = externalFunctionRegistry.get(functionName);
+    }
     if (implementList == null || implementList.isEmpty()) {
       throw new IllegalStateException(String.format("Cannot resolve function: %s", functionName));
     }

datasources/src/main/java/org/opensearch/sql/datasources/service/DataSourceServiceImpl.java

@@ -11,7 +11,6 @@
 import java.util.*;
 import java.util.stream.Collectors;
 
-import lombok.Getter;
 import org.opensearch.sql.datasource.DataSourceService;
 import org.opensearch.sql.datasource.RequestContext;
 import org.opensearch.sql.datasource.model.DataSource;
@@ -20,9 +19,7 @@
 import org.opensearch.sql.datasources.auth.DataSourceUserAuthorizationHelper;
 import org.opensearch.sql.datasources.exceptions.DataSourceNotFoundException;
 import org.opensearch.sql.datasources.exceptions.DatasourceDisabledException;
-import org.opensearch.sql.legacy.esdomain.OpenSearchClient;
 import org.opensearch.sql.storage.DataSourceFactory;
-import org.opensearch.transport.client.node.NodeClient;
 
 /**
  * Default implementation of {@link DataSourceService}. It is per-jvm single instance.

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

@@ -372,7 +372,7 @@ public QueryPlanFactory queryPlanFactory(ExecutionEngine executionEngine) {
       Planner planner = new Planner(LogicalPlanOptimizer.create());
       // NodeClient is not used in integration test, so we pass null
       QueryService queryService =
-          new QueryService(analyzer, executionEngine, planner, dataSourceService, settings, null);
+          new QueryService(analyzer, executionEngine, planner, dataSourceService, settings);
       return new QueryPlanFactory(queryService);
     }
   }

opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java

@@ -18,7 +18,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicReference;
-import lombok.RequiredArgsConstructor;
 import org.apache.calcite.plan.RelOptUtil;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.RelRoot;
@@ -38,21 +37,33 @@
 import org.opensearch.sql.executor.ExecutionEngine.Schema.Column;
 import org.opensearch.sql.executor.Explain;
 import org.opensearch.sql.executor.pagination.PlanSerializer;
+import org.opensearch.sql.expression.function.BuiltinFunctionName;
+import org.opensearch.sql.expression.function.PPLFuncImpTable;
 import org.opensearch.sql.opensearch.client.OpenSearchClient;
 import org.opensearch.sql.opensearch.executor.protector.ExecutionProtector;
+import org.opensearch.sql.opensearch.functions.GeoIpFunction;
 import org.opensearch.sql.opensearch.util.JdbcOpenSearchDataTypeConvertor;
 import org.opensearch.sql.planner.physical.PhysicalPlan;
 import org.opensearch.sql.storage.TableScanOperator;
 
 /** OpenSearch execution engine implementation. */
-@RequiredArgsConstructor
 public class OpenSearchExecutionEngine implements ExecutionEngine {
 
   private final OpenSearchClient client;
 
   private final ExecutionProtector executionProtector;
   private final PlanSerializer planSerializer;
 
+  public OpenSearchExecutionEngine(
+      OpenSearchClient client,
+      ExecutionProtector executionProtector,
+      PlanSerializer planSerializer) {
+    this.client = client;
+    this.executionProtector = executionProtector;
+    this.planSerializer = planSerializer;
+    registerOpenSearchFunctions();
+  }
+
   @Override
   public void execute(PhysicalPlan physicalPlan, ResponseListener<QueryResponse> listener) {
     execute(physicalPlan, ExecutionContext.emptyExecutionContext(), listener);
@@ -224,4 +235,12 @@ private void buildResultSet(
     QueryResponse response = new QueryResponse(schema, values, null);
     listener.onResponse(response);
   }
+
+  /** Registers opensearch-dependent functions */
+  private void registerOpenSearchFunctions() {
+    PPLFuncImpTable.FunctionImp geoIpImpl =
+        (builder, args) ->
+            builder.makeCall(new GeoIpFunction(client.getNodeClient()).toUDF("GEOIP"), args);
+    PPLFuncImpTable.INSTANCE.registerExternalFunction(BuiltinFunctionName.GEOIP, geoIpImpl);
+  }
 }

…ssion/function/udf/ip/GeoIpFunction.java …/opensearch/functions/GeoIpFunction.javacore/src/main/java/org/opensearch/sql/expression/function/udf/ip/GeoIpFunction.java renamed to opensearch/src/main/java/org/opensearch/sql/opensearch/functions/GeoIpFunction.java core/src/main/java/org/opensearch/sql/expression/function/udf/ip/GeoIpFunction.java renamed to opensearch/src/main/java/org/opensearch/sql/opensearch/functions/GeoIpFunction.java

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-package org.opensearch.sql.expression.function.udf.ip;
+package org.opensearch.sql.opensearch.functions;
 
 import java.util.*;
 import java.util.stream.Collectors;

plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java

@@ -98,22 +98,13 @@ public SQLService sqlService(QueryManager queryManager, QueryPlanFactory queryPl
   /** {@link QueryPlanFactory}. */
   @Provides
   public QueryPlanFactory queryPlanFactory(
-      DataSourceService dataSourceService,
-      ExecutionEngine executionEngine,
-      Settings settings,
-      OpenSearchClient client) {
+      DataSourceService dataSourceService, ExecutionEngine executionEngine, Settings settings) {
     Analyzer analyzer =
         new Analyzer(
             new ExpressionAnalyzer(functionRepository), dataSourceService, functionRepository);
     Planner planner = new Planner(LogicalPlanOptimizer.create());
     QueryService queryService =
-        new QueryService(
-            analyzer,
-            executionEngine,
-            planner,
-            dataSourceService,
-            settings,
-            client.getNodeClient());
+        new QueryService(analyzer, executionEngine, planner, dataSourceService, settings);
     return new QueryPlanFactory(queryService);
   }
 }