feat(plugin-native-sidecar): Add native row expression optimizer

presto-docs/src/main/sphinx/plugin/native-sidecar-plugin.rst

@@ -106,3 +106,16 @@ Property Name                                Description
                                              query plans against native engine, ensuring execution compatibility.
 ============================================ ===================================================================== ==============================
 
+Expression optimizer
+-----------------
+
+These properties must be configured in ``etc/expression-manager/native.properties`` to use the native expression optimizer of the ``NativeSidecarPlugin``.
+
+============================================ ===================================================================== ==============================
+Property Name                                Description                                                           Value
+============================================ ===================================================================== ==============================
+``expression-manager-factory.name``          Identifier for the expression optimizer. Enables optimization of      `native`
+                                             expressions using the native expression optimizer.
+============================================ ===================================================================== ==============================
+
+To enable the native expression optimizer for your session, set the expression_optimizer_name session property to native: ``SET SESSION expression_optimizer_name = 'native'``

presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java

@@ -855,6 +855,11 @@ public CatalogSchemaName getDefaultNamespace()
         return defaultNamespace;
     }
 
+    public HandleResolver getHandleResolver()
+    {
+        return handleResolver;
+    }
+
     protected Type getType(UserDefinedType userDefinedType)
     {
         // Distinct type

presto-main-base/src/main/java/com/facebook/presto/metadata/HandleJsonModule.java

@@ -23,6 +23,18 @@
 public class HandleJsonModule
         implements Module
 {
+    private final HandleResolver handleResolver;
+
+    public HandleJsonModule()
+    {
+        this(null);
+    }
+
+    public HandleJsonModule(HandleResolver handleResolver)
+    {
+        this.handleResolver = handleResolver;
+    }
+
     @Override
     public void configure(Binder binder)
     {
@@ -40,6 +52,11 @@ public void configure(Binder binder)
         jsonBinder(binder).addModuleBinding().to(PartitioningHandleJacksonModule.class);
         jsonBinder(binder).addModuleBinding().to(FunctionHandleJacksonModule.class);
 
-        binder.bind(HandleResolver.class).in(Scopes.SINGLETON);
+        if (handleResolver == null) {
+            binder.bind(HandleResolver.class).in(Scopes.SINGLETON);
+        }
+        else {
+            binder.bind(HandleResolver.class).toInstance(handleResolver);
+        }
     }
 }

presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataManager.java

@@ -83,6 +83,7 @@
 import com.facebook.presto.sql.analyzer.FunctionsConfig;
 import com.facebook.presto.sql.analyzer.SemanticException;
 import com.facebook.presto.sql.analyzer.TypeSignatureProvider;
+import com.facebook.presto.testing.TestProcedureRegistry;
 import com.facebook.presto.transaction.TransactionManager;
 import com.facebook.presto.type.TypeDeserializer;
 import com.google.common.annotations.VisibleForTesting;
@@ -301,6 +302,21 @@ public static MetadataManager createTestMetadataManager(TransactionManager trans
                 procedureRegistry);
     }
 
+    public static MetadataManager createTestMetadataManager(FunctionAndTypeManager functionAndTypeManager)
+    {
+        BlockEncodingManager blockEncodingManager = new BlockEncodingManager();
+        return new MetadataManager(
+                functionAndTypeManager,
+                blockEncodingManager,
+                createTestingSessionPropertyManager(),
+                new SchemaPropertyManager(),
+                new TablePropertyManager(),
+                new ColumnPropertyManager(),
+                new AnalyzePropertyManager(),
+                functionAndTypeManager.getTransactionManager(),
+                new TestProcedureRegistry());
+    }
+
     @Override
     public final void verifyComparableOrderableContract()
     {

presto-main-base/src/main/java/com/facebook/presto/sql/expressions/ExpressionOptimizerManager.java

@@ -13,12 +13,14 @@
  */
 package com.facebook.presto.sql.expressions;
 
+import com.facebook.airlift.log.Logger;
 import com.facebook.presto.FullConnectorSession;
 import com.facebook.presto.Session;
 import com.facebook.presto.metadata.FunctionAndTypeManager;
 import com.facebook.presto.nodeManager.PluginNodeManager;
 import com.facebook.presto.spi.ConnectorSession;
 import com.facebook.presto.spi.NodeManager;
+import com.facebook.presto.spi.RowExpressionSerde;
 import com.facebook.presto.spi.relation.ExpressionOptimizer;
 import com.facebook.presto.spi.relation.ExpressionOptimizerProvider;
 import com.facebook.presto.spi.sql.planner.ExpressionOptimizerContext;
@@ -46,29 +48,32 @@
 public class ExpressionOptimizerManager
         implements ExpressionOptimizerProvider
 {
+    private static final Logger log = Logger.get(ExpressionOptimizerManager.class);
     public static final String DEFAULT_EXPRESSION_OPTIMIZER_NAME = "default";
     private static final File EXPRESSION_MANAGER_CONFIGURATION_DIRECTORY = new File("etc/expression-manager/");
     private static final String EXPRESSION_MANAGER_FACTORY_NAME = "expression-manager-factory.name";
 
     private final NodeManager nodeManager;
     private final FunctionAndTypeManager functionAndTypeManager;
+    private final RowExpressionSerde rowExpressionSerde;
     private final FunctionResolution functionResolution;
     private final File configurationDirectory;
 
     private final Map<String, ExpressionOptimizerFactory> expressionOptimizerFactories = new ConcurrentHashMap<>();
     private final Map<String, ExpressionOptimizer> expressionOptimizers = new ConcurrentHashMap<>();
 
     @Inject
-    public ExpressionOptimizerManager(PluginNodeManager nodeManager, FunctionAndTypeManager functionAndTypeManager)
+    public ExpressionOptimizerManager(PluginNodeManager nodeManager, FunctionAndTypeManager functionAndTypeManager, RowExpressionSerde rowExpressionSerde)
     {
-        this(nodeManager, functionAndTypeManager, EXPRESSION_MANAGER_CONFIGURATION_DIRECTORY);
+        this(nodeManager, functionAndTypeManager, rowExpressionSerde, EXPRESSION_MANAGER_CONFIGURATION_DIRECTORY);
     }
 
-    public ExpressionOptimizerManager(PluginNodeManager nodeManager, FunctionAndTypeManager functionAndTypeManager, File configurationDirectory)
+    public ExpressionOptimizerManager(PluginNodeManager nodeManager, FunctionAndTypeManager functionAndTypeManager, RowExpressionSerde rowExpressionSerde, File configurationDirectory)
     {
         requireNonNull(nodeManager, "nodeManager is null");
         this.nodeManager = requireNonNull(nodeManager, "nodeManager is null");
         this.functionAndTypeManager = requireNonNull(functionAndTypeManager, "functionAndTypeManager is null");
+        this.rowExpressionSerde = requireNonNull(rowExpressionSerde, "rowExpressionSerde is null");
         this.functionResolution = new FunctionResolution(functionAndTypeManager.getFunctionAndTypeResolver());
         this.configurationDirectory = requireNonNull(configurationDirectory, "configurationDirectory is null");
         expressionOptimizers.put(DEFAULT_EXPRESSION_OPTIMIZER_NAME, new RowExpressionOptimizer(functionAndTypeManager));
@@ -88,7 +93,7 @@ public void loadExpressionOptimizerFactories()
         }
     }
 
-    private void loadExpressionOptimizerFactory(File configurationFile)
+    public void loadExpressionOptimizerFactory(File configurationFile)
             throws IOException
     {
         String optimizerName = getNameWithoutExtension(configurationFile.getName());
@@ -104,13 +109,16 @@ private void loadExpressionOptimizerFactory(File configurationFile)
 
     public void loadExpressionOptimizerFactory(String factoryName, String optimizerName, Map<String, String> properties)
     {
+        requireNonNull(factoryName, "factoryName is null");
         checkArgument(expressionOptimizerFactories.containsKey(factoryName),
                 "ExpressionOptimizerFactory %s is not registered, registered factories: ", factoryName, expressionOptimizerFactories.keySet());
 
+        log.info("-- Loading expression optimizer [%s] --", optimizerName);
         ExpressionOptimizer optimizer = expressionOptimizerFactories.get(factoryName).createOptimizer(
                 properties,
-                new ExpressionOptimizerContext(nodeManager, functionAndTypeManager, functionResolution));
+                new ExpressionOptimizerContext(nodeManager, rowExpressionSerde, functionAndTypeManager, functionResolution));
         expressionOptimizers.put(optimizerName, optimizer);
+        log.info("-- Added expression optimizer [%s] --", optimizerName);
     }
 
     public void addExpressionOptimizerFactory(ExpressionOptimizerFactory expressionOptimizerFactory)

presto-main-base/src/main/java/com/facebook/presto/sql/expressions/JsonCodecRowExpressionSerde.java

@@ -0,0 +1,46 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.sql.expressions;
+
+import com.facebook.airlift.json.JsonCodec;
+import com.facebook.presto.spi.RowExpressionSerde;
+import com.facebook.presto.spi.relation.RowExpression;
+
+import javax.inject.Inject;
+
+import static java.util.Objects.requireNonNull;
+
+public class JsonCodecRowExpressionSerde
+        implements RowExpressionSerde
+{
+    private final JsonCodec<RowExpression> codec;
+
+    @Inject
+    public JsonCodecRowExpressionSerde(JsonCodec<RowExpression> codec)
+    {
+        this.codec = requireNonNull(codec, "codec is null");
+    }
+
+    @Override
+    public String serialize(RowExpression expression)
+    {
+        return codec.toJson(expression);
+    }
+
+    @Override
+    public RowExpression deserialize(String data)
+    {
+        return codec.fromJson(data);
+    }
+}

presto-main-base/src/main/java/com/facebook/presto/sql/relational/RowExpressionOptimizer.java

@@ -13,43 +13,24 @@
  */
 package com.facebook.presto.sql.relational;
 
-import com.facebook.presto.common.CatalogSchemaName;
-import com.facebook.presto.common.QualifiedObjectName;
-import com.facebook.presto.expressions.RowExpressionRewriter;
-import com.facebook.presto.expressions.RowExpressionTreeRewriter;
 import com.facebook.presto.metadata.FunctionAndTypeManager;
 import com.facebook.presto.metadata.Metadata;
 import com.facebook.presto.spi.ConnectorSession;
-import com.facebook.presto.spi.PrestoException;
-import com.facebook.presto.spi.function.FunctionHandle;
-import com.facebook.presto.spi.function.FunctionMetadata;
-import com.facebook.presto.spi.relation.CallExpression;
 import com.facebook.presto.spi.relation.ExpressionOptimizer;
 import com.facebook.presto.spi.relation.RowExpression;
 import com.facebook.presto.spi.relation.VariableReferenceExpression;
-import com.facebook.presto.sql.analyzer.TypeSignatureProvider;
 import com.facebook.presto.sql.planner.RowExpressionInterpreter;
-import com.google.common.collect.ImmutableList;
-import jakarta.annotation.Nullable;
 
-import java.util.IdentityHashMap;
-import java.util.Map;
 import java.util.function.Function;
 
-import static com.facebook.presto.common.Utils.checkState;
-import static com.facebook.presto.metadata.BuiltInTypeAndFunctionNamespaceManager.JAVA_BUILTIN_NAMESPACE;
-import static com.facebook.presto.spi.StandardErrorCode.FUNCTION_NOT_FOUND;
 import static com.facebook.presto.spi.relation.ExpressionOptimizer.Level.OPTIMIZED;
 import static com.facebook.presto.sql.planner.LiteralEncoder.toRowExpression;
-import static com.google.common.collect.ImmutableList.toImmutableList;
-import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 
 public final class RowExpressionOptimizer
         implements ExpressionOptimizer
 {
     private final FunctionAndTypeManager functionAndTypeManager;
-    private final CatalogSchemaName defaultNamespace;
 
     public RowExpressionOptimizer(Metadata metadata)
     {
@@ -59,133 +40,21 @@ public RowExpressionOptimizer(Metadata metadata)
     public RowExpressionOptimizer(FunctionAndTypeManager functionAndTypeManager)
     {
         this.functionAndTypeManager = requireNonNull(functionAndTypeManager, "functionAndTypeManager is null");
-        this.defaultNamespace = functionAndTypeManager.getDefaultNamespace();
     }
 
     @Override
     public RowExpression optimize(RowExpression rowExpression, Level level, ConnectorSession session)
     {
         if (level.ordinal() <= OPTIMIZED.ordinal()) {
-            return getRowExpression(rowExpression, level, session, null);
+            return toRowExpression(rowExpression.getSourceLocation(), new RowExpressionInterpreter(rowExpression, functionAndTypeManager, session, level).optimize(), rowExpression.getType());
         }
         throw new IllegalArgumentException("Not supported optimization level: " + level);
     }
 
     @Override
     public RowExpression optimize(RowExpression expression, Level level, ConnectorSession session, Function<VariableReferenceExpression, Object> variableResolver)
     {
-        return getRowExpression(expression, level, session, variableResolver);
-    }
-
-    private RowExpression getRowExpression(RowExpression expression, Level level, ConnectorSession session, @Nullable Function<VariableReferenceExpression, Object> variableResolver)
-    {
-        BuiltInNamespaceRewriter visitor = new BuiltInNamespaceRewriter();
-        RowExpressionInterpreter interpreter = new RowExpressionInterpreter(
-                visitor.convertToInterpreterNamespace(expression),
-                functionAndTypeManager,
-                session,
-                level);
-        return visitor.restoreOriginalNamespaces(toRowExpression(
-                expression.getSourceLocation(),
-                interpreter.optimize(variableResolver != null ? variableResolver::apply : null),
-                expression.getType()));
-    }
-
-    /**
-     * TODO: GIANT HACK
-     * This class is a hack and should eventually be removed.  It is used to ensure consistent constant folding behavior when the built-in
-     * function namespace has been switched (for example, to native.default. in the case of native functions).  This will no longer be needed
-     * when the native sidecar is capable of providing its own expression optimizer.
-     */
-    private class BuiltInNamespaceRewriter
-    {
-        private final Map<FunctionHandle, FunctionHandle> defaultToOriginalFunctionHandles = new IdentityHashMap<>();
-
-        public RowExpression convertToInterpreterNamespace(RowExpression expression)
-        {
-            if (defaultNamespace.equals(JAVA_BUILTIN_NAMESPACE)) {
-                // No need to replace built-in namespaces if the default namespace is already the Java built-in namespace
-                return expression;
-            }
-            return RowExpressionTreeRewriter.rewriteWith(new ReplaceBuiltInNamespaces(), expression, null);
-        }
-
-        public RowExpression restoreOriginalNamespaces(RowExpression expression)
-        {
-            if (defaultToOriginalFunctionHandles.isEmpty()) {
-                return expression;
-            }
-            return RowExpressionTreeRewriter.rewriteWith(new ReplaceOriginalNamespaces(), expression, null);
-        }
-
-        private class ReplaceBuiltInNamespaces
-                extends RowExpressionRewriter<Void>
-        {
-            @Override
-            public RowExpression rewriteCall(CallExpression call, Void context, RowExpressionTreeRewriter<Void> treeRewriter)
-            {
-                FunctionHandle functionHandle = call.getFunctionHandle();
-                FunctionMetadata functionMetadata = functionAndTypeManager.getFunctionMetadata(functionHandle);
-                if (!functionMetadata.getImplementationType().canBeEvaluatedInCoordinator()) {
-                    checkState(!functionHandle.getCatalogSchemaName().equals(JAVA_BUILTIN_NAMESPACE),
-                            format("FunctionHandle %s is already in the Java built-in namespace (%s), yet is marked as ineligible to be evaluated in the coordinator", functionHandle, functionHandle.getCatalogSchemaName()));
-
-                    // Replace the namespace with the Java built-in namespace
-                    FunctionHandle javaNamespaceFunctionHandle;
-                    try {
-                        javaNamespaceFunctionHandle = functionAndTypeManager.lookupFunction(
-                                QualifiedObjectName.valueOf(JAVA_BUILTIN_NAMESPACE, call.getDisplayName()),
-                                functionHandle.getArgumentTypes().stream().map(TypeSignatureProvider::new).collect(toImmutableList()));
-                    }
-                    catch (PrestoException e) {
-                        if (e.getErrorCode().equals(FUNCTION_NOT_FOUND.toErrorCode())) {
-                            // If the function is not found in the Java built-in namespace, let default rewriter handle it
-                            return null;
-                        }
-                        throw e; // Rethrow other exceptions
-                    }
-
-                    checkState(functionAndTypeManager.getFunctionMetadata(javaNamespaceFunctionHandle).getImplementationType().canBeEvaluatedInCoordinator(),
-                            format("FunctionHandle %s in the Java built-in namespace (%s) is not eligible to be evaluated in the coordinator", javaNamespaceFunctionHandle, JAVA_BUILTIN_NAMESPACE));
-
-                    defaultToOriginalFunctionHandles.put(javaNamespaceFunctionHandle, functionHandle);
-                    ImmutableList<RowExpression> rewrittenArgs = call.getArguments().stream()
-                            .map(arg -> treeRewriter.rewrite(arg, context))
-                            .collect(toImmutableList());
-                    return new CallExpression(
-                            call.getSourceLocation(),
-                            call.getDisplayName(),
-                            javaNamespaceFunctionHandle,
-                            call.getType(),
-                            rewrittenArgs);
-                }
-
-                // Return null to let the default rewriter handle it (which will rewrite children automatically)
-                return null;
-            }
-        }
-
-        private class ReplaceOriginalNamespaces
-                extends RowExpressionRewriter<Void>
-        {
-            @Override
-            public RowExpression rewriteCall(CallExpression call, Void context, RowExpressionTreeRewriter<Void> treeRewriter)
-            {
-                if (defaultToOriginalFunctionHandles.containsKey(call.getFunctionHandle())) {
-                    FunctionHandle originalFunctionHandle = defaultToOriginalFunctionHandles.get(call.getFunctionHandle());
-                    ImmutableList<RowExpression> rewrittenArgs = call.getArguments().stream()
-                            .map(arg -> treeRewriter.rewrite(arg, context))
-                            .collect(toImmutableList());
-                    return new CallExpression(
-                            call.getSourceLocation(),
-                            call.getDisplayName(),
-                            originalFunctionHandle,
-                            call.getType(),
-                            rewrittenArgs);
-                }
-                // Return null to let the default rewriter handle it (which will rewrite children automatically)
-                return null;
-            }
-        }
+        RowExpressionInterpreter interpreter = new RowExpressionInterpreter(expression, functionAndTypeManager, session, level);
+        return toRowExpression(expression.getSourceLocation(), interpreter.optimize(variableResolver::apply), expression.getType());
     }
 }