Extract outer symbols from lambda expressions

trinodb · Oct 2, 2024 · 77b7d83 · 77b7d83
1 parent 5ceaffb
commit 77b7d83
Showing 1 changed file with 17 additions and 6 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/SymbolsExtractor.java b/core/trino-main/src/main/java/io/trino/sql/planner/SymbolsExtractor.java
@@ -88,7 +88,7 @@ public static Set<Symbol> extractUnique(WindowNode.Function function)
     public static List<Symbol> extractAll(Expression expression)
     {
         ImmutableList.Builder<Symbol> builder = ImmutableList.builder();
-        new SymbolBuilderVisitor().process(expression, builder);
+        new SymbolBuilderVisitor().process(expression, new Context(ImmutableSet.of(), builder));
         return builder.build();
     }
 
@@ -132,20 +132,31 @@ public static Set<Symbol> extractOutputSymbols(PlanNode planNode, Lookup lookup)
     }
 
     private static class SymbolBuilderVisitor
-            extends DefaultTraversalVisitor<ImmutableList.Builder<Symbol>>
+            extends DefaultTraversalVisitor<Context>
     {
         @Override
-        protected Void visitReference(Reference node, ImmutableList.Builder<Symbol> builder)
+        protected Void visitReference(Reference node, Context context)
         {
-            builder.add(Symbol.from(node));
+            Symbol symbol = Symbol.from(node);
+            if (!context.lambdaArguments().contains(symbol)) {
+                context.builder().add(symbol);
+            }
             return null;
         }
 
         @Override
-        protected Void visitLambda(Lambda node, ImmutableList.Builder<Symbol> context)
+        protected Void visitLambda(Lambda node, Context context)
         {
-            // Symbols in lambda expression are bound to lambda arguments, so no need to extract them
+            Context lambdaContext = new Context(
+                    ImmutableSet.<Symbol>builder()
+                            .addAll(context.lambdaArguments())
+                            .addAll(node.arguments())
+                            .build(),
+                    context.builder());
+            process(node.body(), lambdaContext);
             return null;
         }
     }
+
+    private record Context(Set<Symbol> lambdaArguments, ImmutableList.Builder<Symbol> builder) {}
 }