Supports serilizing external OpenSearch UDFs

Signed-off-by: Yuanchun Shen <[email protected]>

Author: yuancu

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

@@ -5,8 +5,7 @@
 
 package org.opensearch.sql.opensearch.executor;
 
-import static org.opensearch.sql.expression.function.BuiltinFunctionName.DISTINCT_COUNT_APPROX;
-
+import com.google.common.base.Suppliers;
 import java.security.AccessController;
 import java.security.PrivilegedAction;
 import java.sql.PreparedStatement;
@@ -17,16 +16,21 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 import org.apache.calcite.plan.RelOptUtil;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.RelRoot;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeField;
 import org.apache.calcite.runtime.Hook;
 import org.apache.calcite.sql.SqlExplainLevel;
+import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.SqlOperatorTable;
 import org.apache.calcite.sql.type.ReturnTypes;
 import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.calcite.sql.util.ListSqlOperatorTable;
 import org.apache.calcite.sql.validate.SqlUserDefinedAggFunction;
 import org.apache.calcite.sql.validate.SqlUserDefinedFunction;
 import org.apache.logging.log4j.LogManager;
@@ -273,8 +277,9 @@ private void buildResultSet(
   private void registerOpenSearchFunctions() {
     if (client instanceof OpenSearchNodeClient) {
       SqlUserDefinedFunction geoIpFunction =
-          new GeoIpFunction(client.getNodeClient()).toUDF("GEOIP");
+          new GeoIpFunction(client.getNodeClient()).toUDF(BuiltinFunctionName.GEOIP.name());
       PPLFuncImpTable.INSTANCE.registerExternalOperator(BuiltinFunctionName.GEOIP, geoIpFunction);
+      OperatorTable.addOperator(BuiltinFunctionName.GEOIP.name(), geoIpFunction);
     } else {
       logger.info(
           "Function [GEOIP] not registered: incompatible client type {}",
@@ -284,10 +289,37 @@ private void registerOpenSearchFunctions() {
     SqlUserDefinedAggFunction approxDistinctCountFunction =
         UserDefinedFunctionUtils.createUserDefinedAggFunction(
             DistinctCountApproxAggFunction.class,
-            DISTINCT_COUNT_APPROX.toString(),
+            BuiltinFunctionName.DISTINCT_COUNT_APPROX.name(),
             ReturnTypes.BIGINT_FORCE_NULLABLE,
             null);
     PPLFuncImpTable.INSTANCE.registerExternalAggOperator(
-        DISTINCT_COUNT_APPROX, approxDistinctCountFunction);
+        BuiltinFunctionName.DISTINCT_COUNT_APPROX, approxDistinctCountFunction);
+    OperatorTable.addOperator(
+        BuiltinFunctionName.DISTINCT_COUNT_APPROX.name(), approxDistinctCountFunction);
+  }
+
+  /**
+   * Dynamic SqlOperatorTable that allows adding operators after initialization. Similar to
+   * PPLBuiltinOperator.instance() or SqlStdOperatorTable.instance().
+   */
+  public static class OperatorTable extends ListSqlOperatorTable {
+    private static final Supplier<OperatorTable> INSTANCE =
+        Suppliers.memoize(() -> (OperatorTable) new OperatorTable().init());
+    // Use map instead of list to avoid duplicated elements if the class is initialized multiple
+    // times
+    private static final Map<String, SqlOperator> operators = new ConcurrentHashMap<>();
+
+    public static SqlOperatorTable instance() {
+      return INSTANCE.get();
+    }
+
+    private ListSqlOperatorTable init() {
+      setOperators(buildIndex(operators.values()));
+      return this;
+    }
+
+    public static synchronized void addOperator(String name, SqlOperator operator) {
+      operators.put(name, operator);
+    }
   }
 }

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/serde/RelJsonSerializer.java

@@ -12,6 +12,7 @@
 import java.io.ByteArrayOutputStream;
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
+import java.io.Serializable;
 import java.util.Base64;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
@@ -31,6 +32,7 @@
 import org.opensearch.sql.calcite.CalcitePlanContext;
 import org.opensearch.sql.data.type.ExprType;
 import org.opensearch.sql.expression.function.PPLBuiltinOperators;
+import org.opensearch.sql.opensearch.executor.OpenSearchExecutionEngine;
 import org.opensearch.sql.opensearch.util.OpenSearchRelOptUtil;
 
 /**
@@ -39,7 +41,7 @@
  * <p>This serializer:
  * <li>Uses Calcite's RelJson class to convert RexNode and RelDataType to/from JSON string
  * <li>Manages required OpenSearch field mapping information Note: OpenSearch ExprType subclasses
- *     implement {@link java.io.Serializable} and are handled through standard Java serialization.
+ *     implement {@link Serializable} and are handled through standard Java serialization.
  */
 @Getter
 public class RelJsonSerializer {
@@ -52,13 +54,7 @@ public class RelJsonSerializer {
   private static final ObjectMapper mapper = new ObjectMapper();
   private static final TypeReference<LinkedHashMap<String, Object>> TYPE_REF =
       new TypeReference<>() {};
-  private static final SqlOperatorTable pplSqlOperatorTable =
-      SqlOperatorTables.chain(
-          PPLBuiltinOperators.instance(),
-          SqlStdOperatorTable.instance(),
-          // Add a list of necessary SqlLibrary if needed
-          SqlLibraryOperatorTableFactory.INSTANCE.getOperatorTable(
-              SqlLibrary.MYSQL, SqlLibrary.BIG_QUERY, SqlLibrary.SPARK, SqlLibrary.POSTGRESQL));
+  private static volatile SqlOperatorTable pplSqlOperatorTable;
 
   static {
     mapper.configure(DeserializationFeature.USE_BIG_DECIMAL_FOR_FLOATS, true);
@@ -68,6 +64,27 @@ public RelJsonSerializer(RelOptCluster cluster) {
     this.cluster = cluster;
   }
 
+  private static SqlOperatorTable getPplSqlOperatorTable() {
+    if (pplSqlOperatorTable == null) {
+      synchronized (RelJsonSerializer.class) {
+        if (pplSqlOperatorTable == null) {
+          pplSqlOperatorTable =
+              SqlOperatorTables.chain(
+                  PPLBuiltinOperators.instance(),
+                  SqlStdOperatorTable.instance(),
+                  OpenSearchExecutionEngine.OperatorTable.instance(),
+                  // Add a list of necessary SqlLibrary if needed
+                  SqlLibraryOperatorTableFactory.INSTANCE.getOperatorTable(
+                      SqlLibrary.MYSQL,
+                      SqlLibrary.BIG_QUERY,
+                      SqlLibrary.SPARK,
+                      SqlLibrary.POSTGRESQL));
+        }
+      }
+    }
+    return pplSqlOperatorTable;
+  }
+
   /**
    * Serializes Calcite expressions and field types into a map object string.
    *
@@ -136,7 +153,8 @@ public Map<String, Object> deserialize(String struct) {
       Map<String, Object> rowTypeMap = mapper.readValue((String) objectMap.get(ROW_TYPE), TYPE_REF);
       RelDataType rowType = relJson.toType(cluster.getTypeFactory(), rowTypeMap);
       OpenSearchRelInputTranslator inputTranslator = new OpenSearchRelInputTranslator(rowType);
-      relJson = relJson.withInputTranslator(inputTranslator).withOperatorTable(pplSqlOperatorTable);
+      relJson =
+          relJson.withInputTranslator(inputTranslator).withOperatorTable(getPplSqlOperatorTable());
       Map<String, Object> exprMap = mapper.readValue((String) objectMap.get(EXPR), TYPE_REF);
       RexNode rexNode = relJson.toRex(cluster, exprMap);