diff --git a/pkg/frontend/computation_wrapper.go b/pkg/frontend/computation_wrapper.go index 693a4cdfc1a96..57bf3822e3452 100644 --- a/pkg/frontend/computation_wrapper.go +++ b/pkg/frontend/computation_wrapper.go @@ -1750,6 +1750,10 @@ 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, @@ -1757,9 +1761,8 @@ func specializePreparedExecutionPlan( 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 } diff --git a/pkg/frontend/computation_wrapper_test.go b/pkg/frontend/computation_wrapper_test.go index 7f12b94e13c91..225404d64e89c 100644 --- a/pkg/frontend/computation_wrapper_test.go +++ b/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) diff --git a/pkg/sql/plan/utils.go b/pkg/sql/plan/utils.go index 9c43fe772236d..b6186429f1a64 100644 --- a/pkg/sql/plan/utils.go +++ b/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)}}, } 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 c784e5b3d6bf4..0a1c72b3ef118 100644 --- a/pkg/sql/plan/visit_plan_rule_test.go +++ b/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")