Fix improperly scoped argument binding in lambda expressions
@@ -55,7 +55,6 @@ import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.collect.ImmutableMap.toImmutableMap;
import static java.lang.String.format;
import static java.util.Locale.ENGLISH;
-import static java.util.Objects.requireNonNull;
import static java.util.function.Function.identity;
public final class SqlFunctionUtils
@@ -226,7 +225,7 @@ public final class SqlFunctionUtils
for (int i = 0; i < argumentNames.size(); i++) {
argumentBindings.put(argumentNames.get(i), argumentValues.get(i));
}
- return ExpressionTreeRewriter.rewriteWith(new ExpressionFunctionVisitor(argumentBindings.build()), function);
+ return ExpressionTreeRewriter.rewriteWith(new ExpressionFunctionVisitor(), function, argumentBindings.build());
}
public static RowExpression bindFunctionArguments(RowExpression function, List<Optional<String>> argumentNames, List<RowExpression> argumentValues)
@@ -262,20 +261,26 @@ public final class SqlFunctionUtils
}
private static class ExpressionFunctionVisitor
- extends ExpressionRewriter<Void>
+ extends ExpressionRewriter<Map<String, Expression>>
{
- private final Map<String, Expression> argumentBindings;
-
- public ExpressionFunctionVisitor(Map<String, Expression> argumentBindings)
+ @Override
+ public Expression rewriteLambdaExpression(LambdaExpression lambda, Map<String, Expression> context, ExpressionTreeRewriter<Map<String, Expression>> treeRewriter)
{
- this.argumentBindings = requireNonNull(argumentBindings, "argumentBindings is null");
+ ImmutableList<String> lambdaStringArguments = lambda.getArguments().stream()
+ .map(x -> x.getName().getValue())
+ .collect(toImmutableList());
+ ImmutableMap<String, Expression> lambdaContext = context.entrySet().stream()
+ .filter(entry -> !lambdaStringArguments.contains(entry.getKey()))
+ .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
+ Expression rewrittenBody = treeRewriter.rewrite(lambda.getBody(), lambdaContext);
+ return new LambdaExpression(lambda.getArguments(), rewrittenBody);
}
@Override
- public Expression rewriteIdentifier(Identifier node, Void context, ExpressionTreeRewriter<Void> treeRewriter)
+ public Expression rewriteIdentifier(Identifier node, Map<String, Expression> context, ExpressionTreeRewriter<Map<String, Expression>> treeRewriter)
{
- if (argumentBindings.containsKey(node.getValue())) {
- return argumentBindings.get(node.getValue());
+ if (context.containsKey(node.getValue())) {
+ return context.get(node.getValue());
}
return node;
}
