Merge pull request #135278 from cockroachdb/blathers/backport-release…

…-24.3-135228

release-24.3: sql: fix edge case causing suboptimal generic query plans

pkg/ccl/logictestccl/testdata/logic_test/generic

@@ -1220,3 +1220,100 @@ ALTER TABLE a RENAME TO b
 
 statement error pgcode 42P01 pq: relation \"a\" does not exist
 EXECUTE p
+
+statement ok
+DEALLOCATE p
+
+# Regression test for #135151. Do not build suboptimal generic plans when an
+# ideal generic plan (using the placeholder fast path) was previously planned.
+statement ok
+CREATE TABLE t135151 (
+  k INT PRIMARY KEY,
+  a INT,
+  b INT,
+  INDEX (a, b)
+);
+
+statement ok
+SET plan_cache_mode = force_custom_plan;
+
+statement ok
+PREPARE p AS SELECT k FROM t135151 WHERE a = $1 AND b = $2;
+
+query T
+EXPLAIN ANALYZE EXECUTE p(1, 2);
+----
+planning time: 10µs
+execution time: 100µs
+distribution: <hidden>
+vectorized: <hidden>
+plan type: generic, reused
+maximum memory usage: <hidden>
+network usage: <hidden>
+regions: <hidden>
+isolation level: serializable
+priority: normal
+quality of service: regular
+·
+• scan
+  sql nodes: <hidden>
+  kv nodes: <hidden>
+  regions: <hidden>
+  actual row count: 0
+  KV time: 0µs
+  KV contention time: 0µs
+  KV rows decoded: 0
+  KV bytes read: 0 B
+  KV gRPC calls: 0
+  estimated max memory allocated: 0 B
+  missing stats
+  table: t135151@t135151_a_b_idx
+  spans: [/1/2 - /1/2]
+
+statement ok
+ALTER TABLE t135151 INJECT STATISTICS '[
+  {
+    "columns": ["a"],
+    "created_at": "2018-01-01 1:00:00.00000+00:00",
+    "row_count": 10000,
+    "distinct_count": 10,
+    "avg_size": 1
+  },
+  {
+    "columns": ["b"],
+    "created_at": "2018-01-01 1:00:00.00000+00:00",
+    "row_count": 10000,
+    "distinct_count": 10,
+    "avg_size": 1
+  }
+]';
+
+query T
+EXPLAIN ANALYZE EXECUTE p(1, 2);
+----
+planning time: 10µs
+execution time: 100µs
+distribution: <hidden>
+vectorized: <hidden>
+plan type: custom
+maximum memory usage: <hidden>
+network usage: <hidden>
+regions: <hidden>
+isolation level: serializable
+priority: normal
+quality of service: regular
+·
+• scan
+  sql nodes: <hidden>
+  kv nodes: <hidden>
+  regions: <hidden>
+  actual row count: 0
+  KV time: 0µs
+  KV contention time: 0µs
+  KV rows decoded: 0
+  KV bytes read: 0 B
+  KV gRPC calls: 0
+  estimated max memory allocated: 0 B
+  estimated row count: 100 (1.0% of the table; stats collected <hidden> ago)
+  table: t135151@t135151_a_b_idx
+  spans: [/1/2 - /1/2]

pkg/sql/plan_opt.go

@@ -590,14 +590,10 @@ func (opc *optPlanningCtx) incPlanTypeTelemetry(cachedMemo *memo.Memo) {
 	}
 }
 
-// useGenericPlan returns true if a generic query plan should be used instead of
-// a custom plan.
-func (opc *optPlanningCtx) useGenericPlan() bool {
+// buildNonIdealGenericPlan returns true if we should attempt to build a
+// non-ideal generic query plan.
+func (opc *optPlanningCtx) buildNonIdealGenericPlan() bool {
 	prep := opc.p.stmt.Prepared
-	// Always use an ideal generic plan.
-	if prep.IdealGenericPlan {
-		return true
-	}
 	switch opc.p.SessionData().PlanCacheMode {
 	case sessiondatapb.PlanCacheModeForceGeneric:
 		return true
@@ -621,6 +617,17 @@ func (opc *optPlanningCtx) useGenericPlan() bool {
 	}
 }
 
+// reuseGenericPlan returns true if a cached generic query plan should be
+// reused. An ideal generic query plan is always reused, if it exists.
+func (opc *optPlanningCtx) reuseGenericPlan() bool {
+	prep := opc.p.stmt.Prepared
+	// Always use an ideal generic plan.
+	if prep.IdealGenericPlan {
+		return true
+	}
+	return opc.buildNonIdealGenericPlan()
+}
+
 // chooseValidPreparedMemo returns an optimized memo that is equal to, or built
 // from, baseMemo or genericMemo. It returns nil if both memos are stale. It
 // selects baseMemo or genericMemo based on the following rules, in order:
@@ -640,7 +647,7 @@ func (opc *optPlanningCtx) useGenericPlan() bool {
 //     new memo.
 func (opc *optPlanningCtx) chooseValidPreparedMemo(ctx context.Context) (*memo.Memo, error) {
 	prep := opc.p.stmt.Prepared
-	if opc.useGenericPlan() {
+	if opc.reuseGenericPlan() {
 		if prep.GenericMemo == nil {
 			// A generic plan does not yet exist.
 			return nil, nil
@@ -721,7 +728,7 @@ func (opc *optPlanningCtx) fetchPreparedMemo(ctx context.Context) (_ *memo.Memo,
 	// build a generic memo from it instead of building the memo from
 	// scratch.
 	opc.log(ctx, "rebuilding cached memo")
-	buildGeneric := opc.useGenericPlan()
+	buildGeneric := opc.buildNonIdealGenericPlan()
 	newMemo, typ, err := opc.buildReusableMemo(ctx, buildGeneric)
 	if err != nil {
 		return nil, err
@@ -738,7 +745,7 @@ func (opc *optPlanningCtx) fetchPreparedMemo(ctx context.Context) (_ *memo.Memo,
 			prep.Costs.SetGeneric(newMemo.RootExpr().(memo.RelExpr).Cost())
 			// Now that the cost of the generic plan is known, we need to
 			// re-evaluate the decision to use a generic or custom plan.
-			if !opc.useGenericPlan() {
+			if !opc.reuseGenericPlan() {
 				// The generic plan that we just built is too expensive, so we need
 				// to build a custom plan. We recursively call fetchPreparedMemo in
 				// case we have a custom plan that can be reused as a starting point