Refactor ExpressionOptimizer

prestodb · Dec 19, 2024 · 5409857 · 5409857
1 parent 3b1aff9
commit 5409857
Show file tree

Hide file tree

Showing 6 changed files with 23 additions and 13 deletions.
diff --git a/...ve-common/src/main/java/com/facebook/presto/hive/rule/BaseSubfieldExtractionRewriter.java b/...ve-common/src/main/java/com/facebook/presto/hive/rule/BaseSubfieldExtractionRewriter.java
@@ -436,18 +436,21 @@ private boolean isCandidate(Map<ColumnHandle, NullableValue> bindings)
 
             // Skip pruning if evaluation fails in a recoverable way. Failing here can cause
             // spurious query failures for partitions that would otherwise be filtered out.
-            Object optimized = null;
+            RowExpression optimized;
             try {
-                optimized = evaluator.getExpressionOptimizer().optimize(expression, OPTIMIZED, session, variableResolver);
+                optimized = evaluator.getExpressionOptimizer().optimize2(expression, OPTIMIZED, session, variableResolver);
             }
             catch (PrestoException e) {
                 propagateIfUnhandled(e);
                 return true;
             }
 
             // If any conjuncts evaluate to FALSE or null, then the whole predicate will never be true and so the partition should be pruned
-            return !Boolean.FALSE.equals(optimized) && optimized != null
-                    && (!(optimized instanceof ConstantExpression) || !((ConstantExpression) optimized).isNull());
+            if (!(optimized instanceof ConstantExpression)) {
+                return true;
+            }
+            ConstantExpression constantExpression = (ConstantExpression) optimized;
+            return !Boolean.FALSE.equals(constantExpression.getValue()) && !constantExpression.isNull();
         }
 
         private static void propagateIfUnhandled(PrestoException e)

diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestDomainTranslator.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestDomainTranslator.java
@@ -95,7 +95,7 @@ public RowExpression optimize(RowExpression rowExpression, Level level, Connecto
         }
 
         @Override
-        public Object optimize(RowExpression expression, Level level, ConnectorSession session, Function<VariableReferenceExpression, Object> variableResolver)
+        public RowExpression optimize(RowExpression expression, Level level, ConnectorSession session, Function<VariableReferenceExpression, Object> variableResolver)
         {
             throw new UnsupportedOperationException();
         }

diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestSubfieldExtractor.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestSubfieldExtractor.java
@@ -80,7 +80,7 @@ public RowExpression optimize(RowExpression rowExpression, Level level, Connecto
         }
 
         @Override
-        public Object optimize(RowExpression expression, Level level, ConnectorSession session, Function<VariableReferenceExpression, Object> variableResolver)
+        public RowExpression optimize(RowExpression expression, Level level, ConnectorSession session, Function<VariableReferenceExpression, Object> variableResolver)
         {
             throw new UnsupportedOperationException();
         }

diff --git a/...iceberg/src/main/java/com/facebook/presto/iceberg/optimizer/IcebergMetadataOptimizer.java b/...iceberg/src/main/java/com/facebook/presto/iceberg/optimizer/IcebergMetadataOptimizer.java
@@ -356,7 +356,7 @@ else if (scalarFunctionName.equals("least")) {
                         throw new PrestoException(StandardErrorCode.NOT_SUPPORTED, "unsupported function: " + scalarFunctionName);
                     }
 
-                    Object reducedValue = rowExpressionService.getExpressionOptimizer().optimize(
+                    RowExpression reducedValue = rowExpressionService.getExpressionOptimizer().optimize(
                             new CallExpression(
                                     Optional.empty(),
                                     scalarFunctionName,
@@ -366,7 +366,8 @@ else if (scalarFunctionName.equals("least")) {
                             Level.EVALUATED,
                             connectorSession,
                             variableReferenceExpression -> null);
-                    reducedArguments.add(new ConstantExpression(reducedValue, returnType));
+                    checkArgument(reducedValue instanceof ConstantExpression, "unexpected expression type: %s", reducedValue.getClass().getSimpleName());
+                    reducedArguments.add(reducedValue);
                 }
                 arguments = reducedArguments;
             }

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/relational/RowExpressionOptimizer.java b/presto-main/src/main/java/com/facebook/presto/sql/relational/RowExpressionOptimizer.java
@@ -46,9 +46,9 @@ public RowExpression optimize(RowExpression rowExpression, Level level, Connecto
     }
 
     @Override
-    public Object optimize(RowExpression expression, Level level, ConnectorSession session, Function<VariableReferenceExpression, Object> variableResolver)
+    public RowExpression optimize(RowExpression expression, Level level, ConnectorSession session, Function<VariableReferenceExpression, Object> variableResolver)
     {
         RowExpressionInterpreter interpreter = new RowExpressionInterpreter(expression, metadata.getFunctionAndTypeManager(), session, level);
-        return interpreter.optimize(variableResolver::apply);
+        return toRowExpression(expression.getSourceLocation(), interpreter.optimize(variableResolver::apply), expression.getType());
     }
 }
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/relation/ExpressionOptimizer.java b/presto-spi/src/main/java/com/facebook/presto/spi/relation/ExpressionOptimizer.java
@@ -20,11 +20,17 @@
 public interface ExpressionOptimizer
 {
     /**
-     * Optimize a RowExpression to
+     * Optimize a RowExpression to its simplest equivalent form.
      */
-    RowExpression optimize(RowExpression rowExpression, Level level, ConnectorSession session);
+    default RowExpression optimize(RowExpression rowExpression, Level level, ConnectorSession session)
+    {
+        return optimize(rowExpression, level, session, variable -> variable);
+    }
 
-    Object optimize(RowExpression expression, Level level, ConnectorSession session, Function<VariableReferenceExpression, Object> variableResolver);
+    /**
+     * Optimize a RowExpression to its simplest equivalent form, replacing VariableReferenceExpressions with their associated values.
+     */
+    RowExpression optimize(RowExpression expression, Level level, ConnectorSession session, Function<VariableReferenceExpression, Object> variableResolver);
 
     enum Level
     {