fix(optimizer): Handle RPC functions inside TRY expressions#27818
Open
feilong-liu wants to merge 1 commit into
Open
fix(optimizer): Handle RPC functions inside TRY expressions#27818feilong-liu wants to merge 1 commit into
feilong-liu wants to merge 1 commit into
Conversation
Contributor
Reviewer's GuideExtends RpcFunctionOptimizer to recognize and rewrite RPC functions that are wrapped inside $internal$try(BIND(..., lambda)) expressions, including substituting lambda parameters with bound variables before RPC extraction, and adds tests to validate TRY-wrapped and direct RPC behavior as well as non-RPC cases. File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
Contributor
There was a problem hiding this comment.
Hey - I've found 4 issues, and left some high level feedback:
- VariableSubstitutor applies the same substitution map inside nested lambdas and does not respect shadowing of lambda parameters, which can lead to incorrect rewrites when inner lambdas reuse parameter names; consider skipping traversal into nested lambdas or masking substitutions for their parameters.
- In VariableSubstitutor.visitSpecialForm, the new SpecialFormExpression is created without preserving the original sourceLocation, unlike visitCall which carries it through; it would be safer to use a constructor that keeps the existing source location for better error reporting and debugging.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- VariableSubstitutor applies the same substitution map inside nested lambdas and does not respect shadowing of lambda parameters, which can lead to incorrect rewrites when inner lambdas reuse parameter names; consider skipping traversal into nested lambdas or masking substitutions for their parameters.
- In VariableSubstitutor.visitSpecialForm, the new SpecialFormExpression is created without preserving the original sourceLocation, unlike visitCall which carries it through; it would be safer to use a constructor that keeps the existing source location for better error reporting and debugging.
## Individual Comments
### Comment 1
<location path="presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/RpcFunctionOptimizer.java" line_range="239-242" />
<code_context>
RpcExtractionContext context,
RowExpressionTreeRewriter<RpcExtractionContext> treeRewriter)
{
+ if (node.getDisplayName().equals("$internal$try")
+ && node.getArguments().size() == 1
+ && isLambdaOrBindWithLambda(node.getArguments().get(0))) {
+ RowExpression result = rewriteTryWithRpcFunction(node, context);
+ if (result != null) {
+ return result;
</code_context>
<issue_to_address>
**issue (bug_risk):** TRY wrapper is dropped when rewriting, which may change exception-handling semantics.
In `rewriteTryWithRpcFunction`, the returned expression is only the rewritten lambda body, so the original `$internal$try(...)` wrapper is lost and its semantics (catching exceptions and returning null instead of failing the row) disappear in this path. If the intent is just to pull RPC calls out of the TRY lambda, the rewritten expression should still be wrapped in a new `$internal$try` (or equivalent) rather than returning the body directly.
</issue_to_address>
### Comment 2
<location path="presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/RpcFunctionOptimizer.java" line_range="392-394" />
<code_context>
+ }
+
+ @Override
+ public RowExpression visitLambda(LambdaDefinitionExpression lambda, Void context)
+ {
+ RowExpression newBody = lambda.getBody().accept(this, null);
+ if (newBody.equals(lambda.getBody())) {
+ return lambda;
</code_context>
<issue_to_address>
**issue (bug_risk):** Variable substitution inside lambdas does not respect parameter shadowing.
`VariableSubstitutor.visitLambda` recurses into the lambda body without excluding variables bound by the lambda itself. If a substitution key matches a lambda parameter (in this or nested lambdas), those references will be replaced, violating lexical scoping and shadowing. Please either filter out substitutions whose keys are in `lambda.getArguments()` before recursing, or otherwise restrict where `VariableSubstitutor` is applied so nested lambdas are not affected incorrectly.
</issue_to_address>
### Comment 3
<location path="presto-main-base/src/test/java/com/facebook/presto/sql/planner/optimizations/TestRpcFunctionOptimizer.java" line_range="280-281" />
<code_context>
assertEquals(invokeParseDispatchBatchSize(rpcCall), 128);
}
+
+ @Test
+ public void testTryWithRpcFunctionProducesRpcNode()
+ {
+ PlanNodeIdAllocator idAllocator = new PlanNodeIdAllocator();
</code_context>
<issue_to_address>
**suggestion (testing):** Add coverage for `$internal$try` with a direct lambda argument (no BIND) since the optimizer supports that path.
`rewriteTryWithRpcFunction` has separate handling for `LambdaDefinitionExpression` and `BIND(..., lambda)`, but the test only covers the `BIND` case. Please add a test for `$internal$try((p) -> test_rpc_function(p, ...))` (lambda as the sole argument) and assert that optimization is applied and an `RPCNode` is produced, to cover that code path and prevent regressions.
Suggested implementation:
```java
assertEquals(invokeParseDispatchBatchSize(rpcCall), 128);
}
@Test
public void testTryWithRpcFunctionLambdaArgumentProducesRpcNode()
{
PlanNodeIdAllocator idAllocator = new PlanNodeIdAllocator();
VariableAllocator variableAllocator = new VariableAllocator();
VariableReferenceExpression inputVar = variableAllocator.newVariable("input_col", VARCHAR);
VariableReferenceExpression outputVar = variableAllocator.newVariable("output_col", VARCHAR);
// Build: $internal$try((p) -> test_rpc_function(p, "model", "system", "{}", '{"api_key":"test-key"}'))
VariableReferenceExpression lambdaParam = variableAllocator.newVariable("p", VARCHAR);
CallExpression rpcInsideLambda = new CallExpression(
"test_rpc_function",
createFunctionHandle("test_rpc_function"),
VARCHAR,
ImmutableList.of(
lambdaParam,
new ConstantExpression(Slices.utf8Slice("model"), VARCHAR),
new ConstantExpression(Slices.utf8Slice("system"), VARCHAR),
new ConstantExpression(Slices.utf8Slice("{}"), VARCHAR),
new ConstantExpression(Slices.utf8Slice("{\"api_key\":\"test-key\"}"), VARCHAR)));
LambdaDefinitionExpression lambda =
new LambdaDefinitionExpression(ImmutableList.of(lambdaParam), rpcInsideLambda);
CallExpression tryCall = new CallExpression(
"$internal$try",
createFunctionHandle("$internal$try"),
VARCHAR,
ImmutableList.of(lambda));
// Use the same optimization path as the BIND-based test and assert that an RpcNode is produced
PlanNode optimized = optimizeTryWithRpcFunctionPlan(
idAllocator,
variableAllocator,
inputVar,
outputVar,
tryCall);
// The top-level node for the project on the RPC result should be an RpcNode
assertTrue(containsRpcNode(optimized));
}
@Test
public void testTryWithRpcFunctionProducesRpcNode()
{
PlanNodeIdAllocator idAllocator = new PlanNodeIdAllocator();
VariableAllocator variableAllocator = new VariableAllocator();
VariableReferenceExpression inputVar = variableAllocator.newVariable("input_col", VARCHAR);
VariableReferenceExpression outputVar = variableAllocator.newVariable("output_col", VARCHAR);
// Build: $internal$try(BIND(input_col, (p) -> test_rpc_function(p, "model", "system", "{}")))
CallExpression rpcInsideLambda = new CallExpression(
"test_rpc_function",
createFunctionHandle("test_rpc_function"),
VARCHAR,
```
To wire this in fully and compile against your existing code:
1. Ensure the following imports are present at the top of `TestRpcFunctionOptimizer.java` (some are likely already there from the existing tests):
- `import com.facebook.presto.common.block.Slices;`
- `import com.facebook.presto.spi.relation.CallExpression;`
- `import com.facebook.presto.spi.relation.ConstantExpression;`
- `import com.facebook.presto.spi.relation.LambdaDefinitionExpression;`
- `import com.facebook.presto.spi.relation.VariableReferenceExpression;`
- `import com.facebook.presto.sql.planner.PlanNodeIdAllocator;`
- `import com.facebook.presto.sql.planner.plan.PlanNode;`
- `import com.facebook.presto.sql.planner.VariableAllocator;`
- `import com.google.common.collect.ImmutableList;`
2. The new test assumes there is an `optimizeTryWithRpcFunctionPlan(...)` helper and a `containsRpcNode(PlanNode)` helper, or equivalent, already used by `testTryWithRpcFunctionProducesRpcNode`. If your existing test uses differently named helpers or inlines this logic:
- Factor that logic into a private helper:
```java
private PlanNode optimizeTryWithRpcFunctionPlan(
PlanNodeIdAllocator idAllocator,
VariableAllocator variableAllocator,
VariableReferenceExpression inputVar,
VariableReferenceExpression outputVar,
CallExpression tryCall)
{
// Move the optimization pipeline currently in testTryWithRpcFunctionProducesRpcNode here
}
private boolean containsRpcNode(PlanNode node)
{
// Reuse the existing assertions / traversal that currently checks for RpcNode in your BIND-based test
}
```
- Update both `testTryWithRpcFunctionProducesRpcNode` and the new `testTryWithRpcFunctionLambdaArgumentProducesRpcNode` to call these helpers.
3. If the function handle for `$internal$try` is created differently in your codebase, adjust `createFunctionHandle("$internal$try")` to match your existing pattern (for example, using `FUNCTION_MANAGER.lookupFunction(...)` or a precomputed handle).
4. Align the argument list for `test_rpc_function` (names, order, and types) with whatever the existing BIND-based test uses so both tests exercise the same RPC function signature.
</issue_to_address>
### Comment 4
<location path="presto-main-base/src/test/java/com/facebook/presto/sql/planner/optimizations/TestRpcFunctionOptimizer.java" line_range="289-298" />
<code_context>
+ VariableReferenceExpression inputVar = variableAllocator.newVariable("input_col", VARCHAR);
+ VariableReferenceExpression outputVar = variableAllocator.newVariable("output_col", VARCHAR);
+
+ // Build: $internal$try(BIND(input_col, (p) -> test_rpc_function(p, "model", "system", "{}")))
+ CallExpression rpcInsideLambda = new CallExpression(
+ "test_rpc_function",
+ createFunctionHandle("test_rpc_function"),
+ VARCHAR,
+ ImmutableList.of(
+ new VariableReferenceExpression(Optional.empty(), "p", VARCHAR),
+ varchar("model"),
+ varchar("system"),
+ varchar("{}")));
+
+ LambdaDefinitionExpression lambda = new LambdaDefinitionExpression(
+ Optional.empty(),
+ ImmutableList.of(VARCHAR),
</code_context>
<issue_to_address>
**suggestion (testing):** Consider a test with multiple bound variables / lambda parameters to validate `VariableSubstitutor` behavior.
Since this path assumes a positional mapping between lambda parameters and BIND variables, the current single-parameter tests don’t exercise that mapping. Please add a test with multiple parameters (e.g., `$internal$try(BIND(input_col1, input_col2, (a, b) -> test_rpc_function(a, b, ...)))`) and assert that the resulting plan’s `RPCNode` is correctly wired to both variables. This will validate `VariableSubstitutor` for multi-parameter lambdas.
Suggested implementation:
```java
VariableAllocator variableAllocator = new VariableAllocator();
VariableReferenceExpression inputVar = variableAllocator.newVariable("input_col", VARCHAR);
VariableReferenceExpression inputVar2 = variableAllocator.newVariable("input_col2", VARCHAR);
VariableReferenceExpression outputVar = variableAllocator.newVariable("output_col", VARCHAR);
// Build: $internal$try(BIND(input_col, input_col2, (a, b) -> test_rpc_function(a, b, "model", "system", "{}")))
```
```java
CallExpression rpcInsideLambda = new CallExpression(
"test_rpc_function",
createFunctionHandle("test_rpc_function"),
VARCHAR,
ImmutableList.of(
new VariableReferenceExpression(Optional.empty(), "a", VARCHAR),
new VariableReferenceExpression(Optional.empty(), "b", VARCHAR),
varchar("model"),
varchar("system"),
varchar("{}")));
LambdaDefinitionExpression lambda = new LambdaDefinitionExpression(
Optional.empty(),
ImmutableList.of(VARCHAR, VARCHAR),
ImmutableList.of("a", "b"),
rpcInsideLambda);
```
```java
SpecialFormExpression bind = new SpecialFormExpression(
SpecialFormExpression.Form.BIND,
VARCHAR,
ImmutableList.of(inputVar, inputVar2, lambda));
```
To fully validate multi-parameter behavior and ensure the test remains correct, you should also:
1. Update any assertions later in `testTryWithRpcFunctionProducesRpcNode` that:
- Assume there is only a single input column or a single BIND argument.
- Inspect the arguments of the resulting `RPCNode` (or equivalent) and currently only check wiring for `inputVar`.
2. Extend those assertions to:
- Verify that the `RPCNode` (or the expression feeding it) has two inputs corresponding to `inputVar` and `inputVar2`.
- Confirm that the ordering of the inputs matches the lambda parameters `a` and `b` (i.e., the first argument in the RPC call corresponds to `inputVar`, the second to `inputVar2`).
3. If other tests rely on the previous single-parameter shape of this plan, consider:
- Either adjusting them to the new two-parameter shape, or
- Introducing a separate, new test for the multi-parameter case and reverting this one to the single-parameter version, using the same pattern as above but in a new `@Test` method.
You’ll need to adapt the concrete assertion code to match how `TestRpcFunctionOptimizer` currently inspects plans and `RPCNode` instances.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
…timizer (prestodb#27818) Summary: `TRY(rpc_func(...))` in SQL is lowered to `$internal$try(BIND(vars, lambda))`, wrapping the RPC function inside a lambda. The `RpcFunctionOptimizer` skips lambda bodies (to avoid breaking per-element semantics in `transform()`/`filter()` lambdas), so RPC functions inside TRY were never rewritten to use `RPCNode`. This caused `VeloxRuntimeError` at runtime because the stub function was called directly. This diff adds special handling for `$internal$try(BIND(..., lambda))` in the optimizer: 1. Detects the `$internal$try` pattern in `rewriteCall` 2. Unwraps the BIND+lambda structure 3. Substitutes lambda parameters with bound variables (so the body uses plan-level variables) 4. Rewrites the substituted body to extract RPC functions into `RPCNode` Follows the same pattern as `PlanRemoteProjections.processInternalTry` which solves the identical problem for remote function planning. Differential Revision: D105342187
bdfdf4f to
bd4cf98
Compare
feilong-liu
commented
May 15, 2026
Contributor
Author
feilong-liu
left a comment
There was a problem hiding this comment.
Re: the two overall review items from sourcery-ai:
-
Lambda shadowing — Already addressed.
VariableSubstitutor.visitLambdafilters out substitutions whose keys match lambda parameters before recursing into the body. -
visitSpecialFormsourceLocation — Already correct. The constructor call usesspecialForm.getSourceLocation():
new SpecialFormExpression(specialForm.getSourceLocation(), specialForm.getForm(), specialForm.getType(), newArgs.build())All four inline comments are also addressed (two by code fixes, two by existing tests testTryWithDirectLambdaRpcFunction and testTryWithMultipleBoundVariables).
zhichenxu-meta
approved these changes
May 15, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
TRY(rpc_func(...))in SQL is lowered to$internal$try(BIND(vars, lambda)),wrapping the RPC function inside a lambda. The
RpcFunctionOptimizerskipslambda bodies to avoid breaking per-element semantics in
transform()/filter()lambdas, so RPC functions inside TRY were never rewritten to use
RPCNode.This caused a
VeloxRuntimeErrorat runtime because the stub function wascalled directly instead of being dispatched through the RPC batching layer.
This change adds special handling for the
$internal$try(BIND(..., lambda))pattern in the optimizer:
$internal$trypattern inrewriteCallplan-level variables
RPCNodeFollows the same pattern as
PlanRemoteProjections.processInternalTrywhichsolves the identical problem for remote function planning.
Motivation and Context
When users write
TRY(rpc_function(...))to gracefully handle RPC failures,the query fails with a
VeloxRuntimeErrorbecause the RPC function is neverextracted into an
RPCNode. The optimizer's lambda-skipping logic, which iscorrect for
transform()/filter()lambdas, inadvertently blocks rewritingof TRY-wrapped RPC calls.
Impact
Queries using
TRY(rpc_function(...))that previously failed at runtime willnow execute correctly. No public API changes. No performance impact on queries
that don't use TRY with RPC functions.
Test Plan
Added unit tests in
TestRpcFunctionOptimizer:testTryWithRpcFunctionProducesRpcNode— TRY(BIND(var, lambda)) with RPC function produces RPCNodetestTryWithoutRpcFunctionIsUnchanged— TRY with non-RPC function is not modifiedtestDirectRpcFunctionStillWorks— regression test for direct RPC calls (no TRY)testTryWithDirectLambdaRpcFunction— TRY with direct lambda (no BIND)testTryWithMultipleBoundVariables— TRY with multiple BIND variablesContributor checklist
Release Notes