Add context to TryExpressionExtractor

soscpd · Nov 21, 2016 · 9cd88f0 · 9cd88f0
1 parent bce915b
commit 9cd88f0
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 33 deletions.
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/gen/CursorProcessorCompiler.java b/presto-main/src/main/java/com/facebook/presto/sql/gen/CursorProcessorCompiler.java
@@ -50,6 +50,7 @@
 import static com.facebook.presto.bytecode.ParameterizedType.type;
 import static com.facebook.presto.sql.gen.BytecodeUtils.generateWrite;
 import static com.facebook.presto.sql.gen.TryCodeGenerator.defineTryMethod;
+import static com.facebook.presto.sql.gen.TryExpressionExtractor.extractTryExpressions;
 import static java.lang.String.format;
 
 public class CursorProcessorCompiler
@@ -181,9 +182,7 @@ private Map<CallExpression, MethodDefinition> generateTryMethods(
             RowExpression projection,
             String methodPrefix)
     {
-        TryExpressionExtractor tryExtractor = new TryExpressionExtractor();
-        projection.accept(tryExtractor, null);
-        List<CallExpression> tryExpressions = tryExtractor.getTryExpressionsPreOrder();
+        List<CallExpression> tryExpressions = extractTryExpressions(projection);
 
         ImmutableMap.Builder<CallExpression, MethodDefinition> tryMethodMap = ImmutableMap.builder();
 

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/gen/JoinFilterFunctionCompiler.java b/presto-main/src/main/java/com/facebook/presto/sql/gen/JoinFilterFunctionCompiler.java
@@ -58,6 +58,7 @@
 import static com.facebook.presto.bytecode.expression.BytecodeExpressions.constantFalse;
 import static com.facebook.presto.sql.gen.BytecodeUtils.invoke;
 import static com.facebook.presto.sql.gen.TryCodeGenerator.defineTryMethod;
+import static com.facebook.presto.sql.gen.TryExpressionExtractor.extractTryExpressions;
 import static com.google.common.base.MoreObjects.toStringHelper;
 import static java.util.Objects.requireNonNull;
 
@@ -200,9 +201,7 @@ private Map<CallExpression, MethodDefinition> generateTryMethods(
             int leftBlocksSize,
             RowExpression filter)
     {
-        TryExpressionExtractor tryExtractor = new TryExpressionExtractor();
-        filter.accept(tryExtractor, null);
-        List<CallExpression> tryExpressions = tryExtractor.getTryExpressionsPreOrder();
+        List<CallExpression> tryExpressions = extractTryExpressions(filter);
 
         ImmutableMap.Builder<CallExpression, MethodDefinition> tryMethodMap = ImmutableMap.builder();
 

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/gen/PageProcessorCompiler.java b/presto-main/src/main/java/com/facebook/presto/sql/gen/PageProcessorCompiler.java
@@ -77,6 +77,7 @@
 import static com.facebook.presto.sql.gen.BytecodeUtils.generateWrite;
 import static com.facebook.presto.sql.gen.BytecodeUtils.loadConstant;
 import static com.facebook.presto.sql.gen.TryCodeGenerator.defineTryMethod;
+import static com.facebook.presto.sql.gen.TryExpressionExtractor.extractTryExpressions;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Verify.verify;
 import static com.google.common.collect.Iterables.concat;
@@ -738,9 +739,7 @@ private Map<CallExpression, MethodDefinition> generateTryMethods(
             RowExpression projection,
             String methodPrefix)
     {
-        TryExpressionExtractor tryExtractor = new TryExpressionExtractor();
-        projection.accept(tryExtractor, null);
-        List<CallExpression> tryExpressions = tryExtractor.getTryExpressionsPreOrder();
+        List<CallExpression> tryExpressions = extractTryExpressions(projection);
 
         ImmutableMap.Builder<CallExpression, MethodDefinition> tryMethodMap = ImmutableMap.builder();
 

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/gen/TryExpressionExtractor.java b/presto-main/src/main/java/com/facebook/presto/sql/gen/TryExpressionExtractor.java
@@ -13,8 +13,6 @@
  */
 package com.facebook.presto.sql.gen;
 
-import com.facebook.presto.bytecode.BytecodeNode;
-import com.facebook.presto.bytecode.Scope;
 import com.facebook.presto.sql.relational.CallExpression;
 import com.facebook.presto.sql.relational.ConstantExpression;
 import com.facebook.presto.sql.relational.InputReferenceExpression;
@@ -29,42 +27,60 @@
 import static com.google.common.collect.Iterables.getOnlyElement;
 
 public class TryExpressionExtractor
-        implements RowExpressionVisitor<Scope, BytecodeNode>
 {
-    private final ImmutableList.Builder<CallExpression> tryExpressions = ImmutableList.builder();
+    private TryExpressionExtractor()
+    {
+    }
 
-    @Override
-    public BytecodeNode visitInputReference(InputReferenceExpression node, Scope scope)
+    public static List<CallExpression> extractTryExpressions(RowExpression expression)
     {
-        // TODO: change such that CallExpressions only capture the inputs they actually depend on
-        return null;
+        Visitor tryOrLambdaExtractor = new Visitor();
+        expression.accept(tryOrLambdaExtractor, new Context());
+        return tryOrLambdaExtractor.getTryExpressionsPreOrder();
     }
 
-    @Override
-    public BytecodeNode visitCall(CallExpression call, Scope scope)
+    private static class Visitor
+            implements RowExpressionVisitor<Context, Void>
     {
-        if (call.getSignature().getName().equals(TRY)) {
-            checkState(call.getArguments().size() == 1, "try call expressions must have a single argument");
-            checkState(getOnlyElement(call.getArguments()) instanceof CallExpression, "try call expression argument must be a call expression");
+        private final ImmutableList.Builder<CallExpression> tryExpressions = ImmutableList.builder();
 
-            tryExpressions.add((CallExpression) getOnlyElement(call.getArguments()));
+        @Override
+        public Void visitInputReference(InputReferenceExpression node, Context context)
+        {
+            // TODO: change such that CallExpressions only capture the inputs they actually depend on
+            return null;
         }
 
-        for (RowExpression rowExpression : call.getArguments()) {
-            rowExpression.accept(this, null);
+        @Override
+        public Void visitCall(CallExpression call, Context context)
+        {
+            if (call.getSignature().getName().equals(TRY)) {
+                checkState(call.getArguments().size() == 1, "try call expressions must have a single argument");
+                checkState(getOnlyElement(call.getArguments()) instanceof CallExpression, "try call expression argument must be a call expression");
+
+                tryExpressions.add((CallExpression) getOnlyElement(call.getArguments()));
+            }
+
+            for (RowExpression rowExpression : call.getArguments()) {
+                rowExpression.accept(this, null);
+            }
+
+            return null;
         }
 
-        return null;
-    }
+        @Override
+        public Void visitConstant(ConstantExpression literal, Context context)
+        {
+            return null;
+        }
 
-    @Override
-    public BytecodeNode visitConstant(ConstantExpression literal, Scope scope)
-    {
-        return null;
+        public List<CallExpression> getTryExpressionsPreOrder()
+        {
+            return tryExpressions.build().reverse();
+        }
     }
 
-    public List<CallExpression> getTryExpressionsPreOrder()
+    private static class Context
     {
-        return tryExpressions.build().reverse();
     }
 }