fix(frontend): reuse prepared specialization metadata

pkg/frontend/computation_wrapper.go

@@ -1750,16 +1750,19 @@ func preparedPlanHasStaticExactNumericPeer(preparePlan *plan2.Plan) bool {
 	return found
 }
 
+// specializePreparedExecutionPlan consumes generation-cached static plan
+// capability plus execute-specific numeric-overload and 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,
 	forceNumericOverload bool,
 	directResultSpecialization bool,
-	forceSpecialization ...bool,
+	needsRuntimeSpecialization bool,
 ) (*plan2.Plan, bool, bool, error) {
-	needsForcedSpecialization := len(forceSpecialization) > 0 && forceSpecialization[0]
 	if len(paramVals) == 0 || executionPlan == nil ||
 		(executionPlan.GetQuery() == nil && executionPlan.GetDdl() == nil &&
 			executionPlan.GetDcl().GetSetVariables() == nil) {
@@ -1769,8 +1772,6 @@ func specializePreparedExecutionPlan(
 		(executionPlan.GetDdl() != nil || executionPlan.GetDcl().GetSetVariables() != nil)
 	needsNumericPrefix := !forceNumericOverload &&
 		plan2.PreparedPlanNeedsNumericPrefixSpecialization(executionPlan, paramVals)
-	needsRuntimeSpecialization := needsForcedSpecialization ||
-		plan2.PreparedPlanNeedsRuntimeSpecialization(executionPlan)
 	if !forceNumericOverload && !needsNumericPrefix && !directResultSpecialization && !binaryLiteralPlan &&
 		!plan2.PreparedPlanHasPaginationParams(executionPlan) && !needsRuntimeSpecialization {
 		return executionPlan, false, false, nil
@@ -2309,10 +2310,14 @@ 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)
+	forceNumericOverload := len(plan2.PreparedPlanNumericFallbackParamPositions(retryPlan)) > 0
 	runtimePlan, _, applied, err := specializePreparedExecutionPlan(
 		ctx, retryPlan, preparedRetry.paramVals, preparedRetry.binaryExecute,
-		len(plan2.PreparedPlanNumericFallbackParamPositions(retryPlan)) > 0,
-		preparedRetry.directResultSpecialization)
+		forceNumericOverload, preparedRetry.directResultSpecialization, needsRuntimeSpecialization)
 	if err != nil {
 		return nil, err
 	}

pkg/frontend/computation_wrapper_test.go

@@ -1191,6 +1191,91 @@ 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 prepared_runtime_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, 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 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 = ?")
+	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, false, false, needsRuntimeSpecialization)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
 func TestSpecializePreparedExecutionPlanSkipsIneligibleSQLPlan(t *testing.T) {
 	ctx := context.Background()
 	floatType := types.T_float32.ToType()
@@ -1218,7 +1303,7 @@ func TestSpecializePreparedExecutionPlanSkipsIneligibleSQLPlan(t *testing.T) {
 		plan2.ParamValue{
 			Value: "1.2345678", PrepareParamKind: vector.PrepareParamDecimal, EnableNumericPrefix: true,
 		},
-	}, false, false, false)
+	}, false, false, false, false)
 	require.NoError(t, err)
 	require.False(t, specialized)
 	require.False(t, applied)

pkg/sql/plan/utils.go

@@ -4492,7 +4492,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
@@ -5431,7 +5433,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 {
@@ -5455,8 +5456,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,
@@ -5474,7 +5473,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)}},
 					})
@@ -5490,7 +5489,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)}},
 			}

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

pkg/sql/plan/visit_plan_rule_test.go

@@ -1577,7 +1577,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")