From 2d430b659814f253382e262550cb792615b3bcbb Mon Sep 17 00:00:00 2001 From: Cao Kai Date: Fri, 28 Aug 2026 22:13:42 +0800 Subject: [PATCH 1/3] fix(frontend): reuse prepared specialization metadata --- pkg/frontend/computation_wrapper.go | 15 +++-- pkg/frontend/computation_wrapper_test.go | 70 +++++++++++++++++++++++- pkg/sql/plan/utils.go | 11 ++-- pkg/sql/plan/utils_test.go | 34 ++++++++++++ pkg/sql/plan/visit_plan_rule_test.go | 5 +- 5 files changed, 123 insertions(+), 12 deletions(-) diff --git a/pkg/frontend/computation_wrapper.go b/pkg/frontend/computation_wrapper.go index 14cd39572f06d..9b97feb4daeb0 100644 --- a/pkg/frontend/computation_wrapper.go +++ b/pkg/frontend/computation_wrapper.go @@ -1214,7 +1214,6 @@ func initExecuteStmtParamWithResolverInSession( prepareStmt.directResultParamPositionsSet = true prepareStmt.numericPrefixConsumer = preparedPlanHasNumericPrefixConsumer( newPreparePlan.Plan, len(newPreparePlan.ParamTypes)) - prepareStmt.directResultParamPositions = plan2.PreparedPlanDirectResultParamPositions(newPreparePlan.Plan) prepareStmt.hasPaginationParams = plan2.PreparedPlanHasPaginationParams(newPreparePlan.Plan) prepareStmt.hasLagLeadParams = len(plan2.PreparedLagLeadParamPositions(newPreparePlan.Plan)) > 0 prepareStmt.ColDefData = newColDefData @@ -1454,7 +1453,8 @@ func initExecuteStmtParamWithResolverInSession( (!binaryExecute || runtimeNumericPrefixCandidate || runtimeDirectResultCandidate || binaryLiteralPlan || prepareStmt.hasPaginationParams || needsRuntimeSpecialization) { runtimePlan, runtimeSpecialized, runtimePlanApplied, err = specializePreparedExecutionPlan( - reqCtx, executionPlan, cwft.paramVals, binaryExecute, runtimeDirectResultCandidate) + reqCtx, executionPlan, cwft.paramVals, binaryExecute, + needsRuntimeSpecialization, runtimeDirectResultCandidate) if err == nil && cacheableRuntimeQuery && runtimeSpecialized && runtimePlanApplied { err = plan2.RestorePreparedRuntimeParamRefs(reqCtx, runtimePlan) if err == nil { @@ -1702,11 +1702,15 @@ func preparedPlanHasStaticExactNumericPeer(preparePlan *plan2.Plan) bool { return found } +// specializePreparedExecutionPlan consumes both generation-cached static plan +// capability and execute-specific direct-result admission. It must not derive +// the static capability itself because this function is on the EXECUTE hot path. func specializePreparedExecutionPlan( ctx context.Context, executionPlan *plan2.Plan, paramVals []any, binaryExecute bool, + needsRuntimeSpecialization bool, directResultSpecialization bool, ) (*plan2.Plan, bool, bool, error) { if len(paramVals) == 0 || executionPlan == nil || @@ -1717,7 +1721,6 @@ func specializePreparedExecutionPlan( binaryLiteralPlan := binaryExecute && (executionPlan.GetDdl() != nil || executionPlan.GetDcl().GetSetVariables() != nil) needsNumericPrefix := plan2.PreparedPlanNeedsNumericPrefixSpecialization(executionPlan, paramVals) - needsRuntimeSpecialization := plan2.PreparedPlanNeedsRuntimeSpecialization(executionPlan) if !needsNumericPrefix && !directResultSpecialization && !binaryLiteralPlan && !plan2.PreparedPlanHasPaginationParams(executionPlan) && !needsRuntimeSpecialization { return executionPlan, false, false, nil @@ -2251,9 +2254,13 @@ func buildPlanForCompileRetry( ctx, retryPlan, preparedRetry.paramVals) return runtimePlan, err } + // A definition-change retry owns a newly built plan generation. Derive its + // static capability once here rather than reusing the preceding generation's + // decision or making the execute helper rescan every invocation. + needsRuntimeSpecialization := plan2.PreparedPlanNeedsRuntimeSpecialization(retryPlan) runtimePlan, _, applied, err := specializePreparedExecutionPlan( ctx, retryPlan, preparedRetry.paramVals, preparedRetry.binaryExecute, - preparedRetry.directResultSpecialization) + needsRuntimeSpecialization, preparedRetry.directResultSpecialization) if err != nil { return nil, err } diff --git a/pkg/frontend/computation_wrapper_test.go b/pkg/frontend/computation_wrapper_test.go index e0f732af54077..fb1f690d2f92a 100644 --- a/pkg/frontend/computation_wrapper_test.go +++ b/pkg/frontend/computation_wrapper_test.go @@ -1135,6 +1135,74 @@ func TestInitExecuteStmtParamSpecializesSQLExecuteCommonTypePlan(t *testing.T) { require.False(t, requiresV26) } +func buildPreparedRuntimePlanForFrontendTest(t testing.TB, sql string) *plan.Plan { + t.Helper() + optimizer := plan2.NewMockOptimizer(false) + stmts, err := mysql.Parse( + optimizer.CurrentContext().GetContext(), + fmt.Sprintf("prepare issue27807_stmt from '%s'", sql), + 1, + ) + require.NoError(t, err) + require.Len(t, stmts, 1) + defer stmts[0].Free() + + prepared, err := plan2.BuildPlan(optimizer.CurrentContext(), stmts[0], false) + require.NoError(t, err) + prepareControl := prepared.GetDcl().GetPrepare() + require.NotNil(t, prepareControl) + require.NotNil(t, prepareControl.GetPlan().GetQuery()) + return prepareControl.GetPlan() +} + +func preparedArithmeticDMLParamValues() []any { + return []any{ + plan2.ParamValue{ + Value: "1", RuntimeType: types.T_int64.ToType(), HasRuntimeType: true, + IsBinaryProtocol: true, PrepareParamKind: vector.PrepareParamInteger, + }, + plan2.ParamValue{ + Value: "7", RuntimeType: types.T_int64.ToType(), HasRuntimeType: true, + IsBinaryProtocol: true, PrepareParamKind: vector.PrepareParamInteger, + }, + } +} + +func TestSpecializePreparedExecutionPlanDoesNotRecomputeStaticCapability(t *testing.T) { + preparedPlan := buildPreparedRuntimePlanForFrontendTest(t, + "update nation set n_regionkey = n_regionkey + ? where n_nationkey = ?") + require.True(t, plan2.PreparedPlanNeedsRuntimeSpecialization(preparedPlan), + "the arithmetic write is the TPCC-shaped positive control") + + runtimePlan, specialized, applied, err := specializePreparedExecutionPlan( + context.Background(), preparedPlan, preparedArithmeticDMLParamValues(), + true, false, false) + require.NoError(t, err) + require.False(t, specialized) + require.False(t, applied) + require.Same(t, preparedPlan, runtimePlan, + "the execution helper must consume the generation-cached decision instead of rescanning the plan") +} + +func BenchmarkSpecializePreparedTPCCArithmeticUpdate(b *testing.B) { + preparedPlan := buildPreparedRuntimePlanForFrontendTest(b, + "update nation set n_regionkey = n_regionkey + ? where n_nationkey = ?") + needsRuntimeSpecialization := plan2.PreparedPlanNeedsRuntimeSpecialization(preparedPlan) + require.True(b, needsRuntimeSpecialization) + paramVals := preparedArithmeticDMLParamValues() + ctx := context.Background() + + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + _, _, _, err := specializePreparedExecutionPlan( + ctx, preparedPlan, paramVals, true, needsRuntimeSpecialization, false) + if err != nil { + b.Fatal(err) + } + } +} + func TestSpecializePreparedExecutionPlanSkipsIneligibleSQLPlan(t *testing.T) { ctx := context.Background() floatType := types.T_float32.ToType() @@ -1162,7 +1230,7 @@ func TestSpecializePreparedExecutionPlanSkipsIneligibleSQLPlan(t *testing.T) { plan2.ParamValue{ Value: "1.2345678", PrepareParamKind: vector.PrepareParamDecimal, EnableNumericPrefix: true, }, - }, false, false) + }, false, false, false) require.NoError(t, err) require.False(t, specialized) require.False(t, applied) diff --git a/pkg/sql/plan/utils.go b/pkg/sql/plan/utils.go index 833dd4bc95e1e..798a9ea622226 100644 --- a/pkg/sql/plan/utils.go +++ b/pkg/sql/plan/utils.go @@ -4392,7 +4392,9 @@ type ParamValue struct { MaterializedValue string // RetainParamRef records that a specialized query plan will be cached and // therefore must retain this parameter as runtime provenance even when the - // parameter itself is unrelated to numeric-prefix specialization. + // parameter itself is unrelated to numeric-prefix specialization. The + // frontend also sets it from generation-cached direct-result positions; + // generic parameter replacement must not rediscover those positions. RetainParamRef bool // EnableNumericPrefix records that the deployment-wide protocol version can // execute planner-injected MySQL numeric-prefix casts. Keep the negotiated @@ -5331,7 +5333,6 @@ func replaceParamVals( preserveDMLWriteArgs ...bool, ) (bool, error) { preserveDMLWrites := len(preserveDMLWriteArgs) > 0 && preserveDMLWriteArgs[0] - directResultPositions := PreparedPlanDirectResultParamPositions(plan0) params := make([]*Expr, len(paramVals)) var err error for i, val := range paramVals { @@ -5355,8 +5356,6 @@ func replaceParamVals( if hasRuntimeType { paramType = makePlan2Type(&runtimeType) } - _, directRuntimeResult := slices.BinarySearch(directResultPositions, int32(i)) - directRuntimeResult = directRuntimeResult && hasRuntimeType if val == nil { pc := &plan.Literal{ Isnull: true, @@ -5374,7 +5373,7 @@ func replaceParamVals( if err != nil { return false, err } - if numericPrefixSource || retainParamRef || directRuntimeResult { + if numericPrefixSource || retainParamRef { attachPreparedRuntimeParamSource(params[i], &plan.Expr{ Typ: paramType, Expr: &plan.Expr_P{P: &plan.ParamRef{Pos: int32(i)}}, }) @@ -5390,7 +5389,7 @@ func replaceParamVals( }, } } - if (numericPrefixSource || retainParamRef || directRuntimeResult) && params[i].GetLit() != nil { + if (numericPrefixSource || retainParamRef) && params[i].GetLit() != nil { params[i].GetLit().Src = &plan.Expr{ Typ: paramType, Expr: &plan.Expr_P{P: &plan.ParamRef{Pos: int32(i)}}, } diff --git a/pkg/sql/plan/utils_test.go b/pkg/sql/plan/utils_test.go index 3a0d39ddbfc90..0a82880e81b22 100644 --- a/pkg/sql/plan/utils_test.go +++ b/pkg/sql/plan/utils_test.go @@ -740,6 +740,40 @@ func TestPreparedPlanDirectResultParamPositions(t *testing.T) { })) } +func TestPreparedDirectResultParamRefRetentionOwnedByCaller(t *testing.T) { + makePlan := func() *plan.Plan { + return &plan.Plan{Plan: &plan.Plan_Query{Query: &plan.Query{ + StmtType: plan.Query_SELECT, + Steps: []int32{0}, + Nodes: []*plan.Node{{ + NodeType: plan.Node_VALUE_SCAN, + ProjectList: []*plan.Expr{{ + Typ: plan.Type{Id: int32(types.T_text)}, + Expr: &plan.Expr_P{P: &plan.ParamRef{Pos: 0}}, + }}, + }}, + }}} + } + fill := func(retain bool) *plan.Expr { + filled, specialized, err := FillValuesOfParamsInPlanWithSpecialization( + context.Background(), makePlan(), []any{ParamValue{ + Value: "42", RuntimeType: types.T_int64.ToType(), HasRuntimeType: true, + IsBinaryProtocol: true, RetainParamRef: retain, + }}) + require.NoError(t, err) + require.Equal(t, retain, specialized, + "only the caller-owned provenance mark may make a direct result cacheable") + return filled.GetQuery().Nodes[0].ProjectList[0] + } + + withoutOwnerMark := fill(false) + require.Nil(t, withoutOwnerMark.GetLit().GetSrc(), + "generic replacement must not rediscover direct-result positions") + withOwnerMark := fill(true) + require.NotNil(t, withOwnerMark.GetLit().GetSrc()) + require.Equal(t, int32(0), withOwnerMark.GetLit().GetSrc().GetP().Pos) +} + func TestPreparedDirectResultSpecializationUpdatesVisibleType(t *testing.T) { for _, test := range []struct { name string diff --git a/pkg/sql/plan/visit_plan_rule_test.go b/pkg/sql/plan/visit_plan_rule_test.go index 42d217f427c67..2a38c9e09437f 100644 --- a/pkg/sql/plan/visit_plan_rule_test.go +++ b/pkg/sql/plan/visit_plan_rule_test.go @@ -1573,7 +1573,10 @@ func TestFillValuesOfParamsSpecializationTracksBinaryExecutionDomains(t *testing require.False(t, specialized, "same-domain text execution should reuse the cached plan") _, specialized, err = FillValuesOfParamsInPlanWithSpecialization(ctx, direct, []any{ - ParamValue{Value: "5", RuntimeType: types.T_int64.ToType(), HasRuntimeType: true}, + ParamValue{ + Value: "5", RuntimeType: types.T_int64.ToType(), HasRuntimeType: true, + RetainParamRef: true, + }, }) require.NoError(t, err) require.True(t, specialized, "direct numeric result metadata must be specialized") From 1b4ab8c346494baaf2190f529790221e7b80bed4 Mon Sep 17 00:00:00 2001 From: Cao Kai Date: Fri, 28 Aug 2026 23:01:45 +0800 Subject: [PATCH 2/3] test(frontend): cover prepared retry generation metadata --- pkg/frontend/computation_wrapper_test.go | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/pkg/frontend/computation_wrapper_test.go b/pkg/frontend/computation_wrapper_test.go index a8b365e0a58c7..f94f6d68d1e3e 100644 --- a/pkg/frontend/computation_wrapper_test.go +++ b/pkg/frontend/computation_wrapper_test.go @@ -1240,6 +1240,23 @@ func TestSpecializePreparedExecutionPlanDoesNotRecomputeStaticCapability(t *test "the execution helper must consume the generation-cached decision instead of rescanning the plan") } +func TestBuildPlanForCompileRetryDerivesStaticCapabilityForNewGeneration(t *testing.T) { + ses, prepareStmt, _, execCtx := newPreparedExecuteEnvForSQL(t, 216, "select coalesce(?, ?) from dual") + defer prepareStmt.Close() + ses.SetSql("execute " + prepareStmt.Name) + originalPlan := prepareStmt.PreparePlan.GetDcl().GetPrepare().GetPlan() + require.True(t, plan2.PreparedPlanNeedsRuntimeSpecialization(originalPlan)) + + retryPlan, err := buildPlanForCompileRetry( + execCtx.reqCtx, ses, ses.GetTxnCompileCtx(), prepareStmt.PrepareStmt, false, + newPreparedExecutionRetry(preparedArithmeticDMLParamValues(), true)) + require.NoError(t, err) + require.Empty(t, queryParamPositions(retryPlan.GetQuery()), retryPlan.String()) + columns := plan2.GetResultColumnsFromPlan(retryPlan) + require.Len(t, columns, 1) + require.Equal(t, int32(types.T_int64), columns[0].Typ.Id) +} + func BenchmarkSpecializePreparedTPCCArithmeticUpdate(b *testing.B) { preparedPlan := buildPreparedRuntimePlanForFrontendTest(b, "update nation set n_regionkey = n_regionkey + ? where n_nationkey = ?") From dab9d1a1ebd779cab81a9d443d75b561765af880 Mon Sep 17 00:00:00 2001 From: Cao Kai Date: Fri, 28 Aug 2026 23:04:59 +0800 Subject: [PATCH 3/3] test(frontend): use invariant prepared fixture name --- pkg/frontend/computation_wrapper_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/frontend/computation_wrapper_test.go b/pkg/frontend/computation_wrapper_test.go index f94f6d68d1e3e..225404d64e89c 100644 --- a/pkg/frontend/computation_wrapper_test.go +++ b/pkg/frontend/computation_wrapper_test.go @@ -1196,7 +1196,7 @@ func buildPreparedRuntimePlanForFrontendTest(t testing.TB, sql string) *plan.Pla optimizer := plan2.NewMockOptimizer(false) stmts, err := mysql.Parse( optimizer.CurrentContext().GetContext(), - fmt.Sprintf("prepare issue27807_stmt from '%s'", sql), + fmt.Sprintf("prepare prepared_runtime_stmt from '%s'", sql), 1, ) require.NoError(t, err)