[Presto] Handle RPC functions inside TRY expressions in RpcFunctionOptimizer

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

Author: feilong-liu

presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/RpcFunctionOptimizer.java

@@ -24,15 +24,19 @@
 import com.facebook.presto.spi.plan.ProjectNode;
 import com.facebook.presto.spi.relation.CallExpression;
 import com.facebook.presto.spi.relation.ConstantExpression;
+import com.facebook.presto.spi.relation.InputReferenceExpression;
 import com.facebook.presto.spi.relation.LambdaDefinitionExpression;
 import com.facebook.presto.spi.relation.RowExpression;
+import com.facebook.presto.spi.relation.RowExpressionVisitor;
+import com.facebook.presto.spi.relation.SpecialFormExpression;
 import com.facebook.presto.spi.relation.VariableReferenceExpression;
 import com.facebook.presto.sql.planner.TypeProvider;
 import com.facebook.presto.sql.planner.plan.RPCNode;
 import com.facebook.presto.sql.planner.plan.SimplePlanRewriter;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import io.airlift.slice.Slice;
 
@@ -232,6 +236,15 @@ public RowExpression rewriteCall(
                         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;
+                        }
+                    }
+
                     if (!rpcFunctionNames.contains(
                             node.getDisplayName().toLowerCase(Locale.ENGLISH))) {
                         return null;
@@ -258,12 +271,160 @@ public RowExpression rewriteLambda(
                         RowExpressionTreeRewriter<RpcExtractionContext> treeRewriter)
                 {
                     // Do not traverse into lambda bodies — RPC calls inside lambdas
-                    // have per-element semantics incompatible with RPCNode batching
+                    // have per-element semantics incompatible with RPCNode batching.
+                    // TRY lambdas are handled specially in rewriteCall above.
                     return node;
                 }
             };
         }
 
+        private static boolean isLambdaOrBindWithLambda(RowExpression expression)
+        {
+            if (expression instanceof LambdaDefinitionExpression) {
+                return true;
+            }
+            if (expression instanceof SpecialFormExpression
+                    && ((SpecialFormExpression) expression).getForm() == SpecialFormExpression.Form.BIND) {
+                List<RowExpression> bindArgs = ((SpecialFormExpression) expression).getArguments();
+                return bindArgs.get(bindArgs.size() - 1) instanceof LambdaDefinitionExpression;
+            }
+            return false;
+        }
+
+        private RowExpression rewriteTryWithRpcFunction(
+                CallExpression tryCall,
+                RpcExtractionContext context)
+        {
+            RowExpression tryArgument = tryCall.getArguments().get(0);
+            LambdaDefinitionExpression lambda;
+            List<RowExpression> bindVariables;
+
+            if (tryArgument instanceof LambdaDefinitionExpression) {
+                lambda = (LambdaDefinitionExpression) tryArgument;
+                bindVariables = ImmutableList.of();
+            }
+            else {
+                SpecialFormExpression bind = (SpecialFormExpression) tryArgument;
+                List<RowExpression> bindArgs = bind.getArguments();
+                lambda = (LambdaDefinitionExpression) bindArgs.get(bindArgs.size() - 1);
+                bindVariables = bindArgs.subList(0, bindArgs.size() - 1);
+            }
+
+            RowExpression body = lambda.getBody();
+            if (!bindVariables.isEmpty()) {
+                List<String> lambdaParams = lambda.getArguments();
+                ImmutableMap.Builder<String, RowExpression> substitutions = ImmutableMap.builder();
+                for (int i = 0; i < lambdaParams.size(); i++) {
+                    substitutions.put(lambdaParams.get(i), bindVariables.get(i));
+                }
+                body = VariableSubstitutor.substitute(body, substitutions.build());
+            }
+
+            if (!containsRpcFunction(body)) {
+                return null;
+            }
+
+            RowExpression rewrittenBody = RowExpressionTreeRewriter.rewriteWith(
+                    createRpcRewriter(), body, context);
+
+            return rewrittenBody;
+        }
+
+        private boolean containsRpcFunction(RowExpression expression)
+        {
+            if (expression instanceof CallExpression) {
+                CallExpression call = (CallExpression) expression;
+                if (rpcFunctionNames.contains(call.getDisplayName().toLowerCase(Locale.ENGLISH))) {
+                    return true;
+                }
+                return call.getArguments().stream().anyMatch(this::containsRpcFunction);
+            }
+            if (expression instanceof SpecialFormExpression) {
+                return ((SpecialFormExpression) expression).getArguments().stream()
+                        .anyMatch(this::containsRpcFunction);
+            }
+            return false;
+        }
+
+        private static class VariableSubstitutor
+                implements RowExpressionVisitor<RowExpression, Void>
+        {
+            private final Map<String, RowExpression> substitutions;
+
+            VariableSubstitutor(Map<String, RowExpression> substitutions)
+            {
+                this.substitutions = substitutions;
+            }
+
+            static RowExpression substitute(RowExpression expression, Map<String, RowExpression> substitutions)
+            {
+                return expression.accept(new VariableSubstitutor(substitutions), null);
+            }
+
+            @Override
+            public RowExpression visitCall(CallExpression call, Void context)
+            {
+                ImmutableList.Builder<RowExpression> newArgs = ImmutableList.builder();
+                boolean changed = false;
+                for (RowExpression arg : call.getArguments()) {
+                    RowExpression newArg = arg.accept(this, null);
+                    newArgs.add(newArg);
+                    changed |= newArg != arg;
+                }
+                return changed
+                        ? new CallExpression(call.getSourceLocation(), call.getDisplayName(), call.getFunctionHandle(), call.getType(), newArgs.build())
+                        : call;
+            }
+
+            @Override
+            public RowExpression visitInputReference(InputReferenceExpression reference, Void context)
+            {
+                return reference;
+            }
+
+            @Override
+            public RowExpression visitConstant(ConstantExpression literal, Void context)
+            {
+                return literal;
+            }
+
+            @Override
+            public RowExpression visitLambda(LambdaDefinitionExpression lambda, Void context)
+            {
+                RowExpression newBody = lambda.getBody().accept(this, null);
+                if (newBody.equals(lambda.getBody())) {
+                    return lambda;
+                }
+                return new LambdaDefinitionExpression(
+                        lambda.getSourceLocation(),
+                        lambda.getArgumentTypes(),
+                        lambda.getArguments(),
+                        newBody);
+            }
+
+            @Override
+            public RowExpression visitVariableReference(VariableReferenceExpression reference, Void context)
+            {
+                RowExpression replacement = substitutions.get(reference.getName());
+                return replacement != null ? replacement : reference;
+            }
+
+            @Override
+            public RowExpression visitSpecialForm(SpecialFormExpression specialForm, Void context)
+            {
+                ImmutableList.Builder<RowExpression> newArgs = ImmutableList.builder();
+                boolean changed = false;
+                for (RowExpression arg : specialForm.getArguments()) {
+                    RowExpression newArg = arg.accept(this, null);
+                    newArgs.add(newArg);
+                    changed |= newArg != arg;
+                }
+                return changed
+                        ? new SpecialFormExpression(specialForm.getForm(), specialForm.getType(), newArgs.build())
+                        : specialForm;
+            }
+        }
+
         private static class RpcExtractionContext
         {
             private final List<ExtractedRpcCall> extractedCalls = new ArrayList<>();