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

…-24.2-129050

release-24.2: opt: disable `GenerateParameterizedJoin` when forcing custom plans

pkg/sql/opt/norm/BUILD.bazel

@@ -91,6 +91,7 @@ go_test(
         "//pkg/sql/sem/catconstants",
         "//pkg/sql/sem/eval",
         "//pkg/sql/sem/tree",
+        "//pkg/sql/sessiondatapb",
         "//pkg/sql/types",
         "//pkg/testutils/datapathutils",
         "//pkg/util/leaktest",

pkg/sql/opt/norm/factory_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/catconstants"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+	"github.com/cockroachdb/cockroach/pkg/sql/sessiondatapb"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 )
 
@@ -84,6 +85,7 @@ func TestCopyAndReplace(t *testing.T) {
 	}
 
 	evalCtx := eval.MakeTestingEvalContext(cluster.MakeTestingClusterSettings())
+	evalCtx.SessionData().PlanCacheMode = sessiondatapb.PlanCacheModeAuto
 
 	var o xform.Optimizer
 	testutils.BuildQuery(t, &o, cat, &evalCtx, "SELECT * FROM ab WHERE a = $1")

pkg/sql/opt/testutils/opttester/BUILD.bazel

@@ -45,6 +45,7 @@ go_library(
         "//pkg/sql/sem/eval",
         "//pkg/sql/sem/tree",
         "//pkg/sql/sem/volatility",
+        "//pkg/sql/sessiondatapb",
         "//pkg/sql/stats",
         "//pkg/testutils/datapathutils",
         "//pkg/testutils/floatcmp",

pkg/sql/opt/testutils/opttester/opt_tester.go

@@ -60,6 +60,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/volatility"
+	"github.com/cockroachdb/cockroach/pkg/sql/sessiondatapb"
 	"github.com/cockroachdb/cockroach/pkg/sql/stats"
 	"github.com/cockroachdb/cockroach/pkg/testutils/datapathutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/floatcmp"
@@ -168,6 +169,9 @@ type Flags struct {
 	// memo/check_expr.go.
 	DisableCheckExpr bool
 
+	// Generic enables optimizations for generic query plans.
+	Generic bool
+
 	// ExploreTraceRule restricts the ExploreTrace output to only show the effects
 	// of a specific rule.
 	ExploreTraceRule opt.RuleName
@@ -477,6 +481,11 @@ func New(catalog cat.Catalog, sqlStr string) *OptTester {
 //
 //   - disable-check-expr: skips the assertions in memo/check_expr.go.
 //
+//   - generic: enables optimizations for generic query plans.
+//     NOTE: This flag sets the plan_cache_mode session setting to "auto", which
+//     cannot be done via the "set" flag because it requires a CCL license,
+//     which optimizer tests are not set up to utilize.
+//
 //   - rule: used with exploretrace; the value is the name of a rule. When
 //     specified, the exploretrace output is filtered to only show expression
 //     changes due to that specific rule.
@@ -1003,6 +1012,9 @@ func (f *Flags) Set(arg datadriven.CmdArg) error {
 	case "disable-check-expr":
 		f.DisableCheckExpr = true
 
+	case "generic":
+		f.evalCtx.SessionData().PlanCacheMode = sessiondatapb.PlanCacheModeAuto
+
 	case "rule":
 		if len(arg.Vals) != 1 {
 			return fmt.Errorf("rule requires one argument")

pkg/sql/opt/xform/BUILD.bazel

@@ -50,6 +50,7 @@ go_library(
         "//pkg/sql/sem/eval",
         "//pkg/sql/sem/tree",
         "//pkg/sql/sem/volatility",
+        "//pkg/sql/sessiondatapb",
         "//pkg/sql/types",
         "//pkg/util/buildutil",
         "//pkg/util/cancelchecker",

pkg/sql/opt/xform/generic_funcs.go

@@ -17,9 +17,16 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/volatility"
+	"github.com/cockroachdb/cockroach/pkg/sql/sessiondatapb"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 )
 
+// GenericRulesEnabled returns true if rules for optimizing generic query plans
+// are enabled, based on the plan_cache_mode session setting.
+func (c *CustomFuncs) GenericRulesEnabled() bool {
+	return c.e.evalCtx.SessionData().PlanCacheMode != sessiondatapb.PlanCacheModeForceCustom
+}
+
 // HasPlaceholdersOrStableExprs returns true if the given relational expression's subtree has
 // at least one placeholder.
 func (c *CustomFuncs) HasPlaceholdersOrStableExprs(e memo.RelExpr) bool {

pkg/sql/opt/xform/rules/generic.opt

@@ -29,7 +29,9 @@
 #
 [GenerateParameterizedJoin, Explore]
 (Select
-    $scan:(Scan $scanPrivate:*) & (IsCanonicalScan $scanPrivate)
+    $scan:(Scan $scanPrivate:*) &
+        (GenericRulesEnabled) &
+        (IsCanonicalScan $scanPrivate)
     $filters:* &
         (HasPlaceholdersOrStableExprs (Root)) &
         (Let

pkg/sql/opt/xform/testdata/external/hibernate

@@ -886,7 +886,7 @@ project
       └── filters
            └── min:14 = $1 [outer=(14), constraints=(/14: (/NULL - ]), fd=()-->(14)]
 
-opt
+opt generic
 select
     person0_.id as id1_2_,
     person0_.address as address2_2_,
@@ -951,7 +951,7 @@ project
       └── filters
            └── max:16 = 0 [outer=(16), constraints=(/16: [/0 - /0]; tight), fd=()-->(16)]
 
-opt
+opt generic
 select
     person0_.id as id1_2_,
     person0_.address as address2_2_,
@@ -1016,7 +1016,7 @@ project
            │         └── filters (true)
            └── filters (true)
 
-opt
+opt generic
 select
     person0_.id as id1_2_,
     person0_.address as address2_2_,

pkg/sql/opt/xform/testdata/external/nova

@@ -107,7 +107,7 @@ create table instance_type_extra_specs
 )
 ----
 
-opt
+opt generic
 select anon_1.flavors_created_at as anon_1_flavors_created_at,
        anon_1.flavors_updated_at as anon_1_flavors_updated_at,
        anon_1.flavors_id as anon_1_flavors_id,
@@ -710,7 +710,7 @@ sort
            └── filters
                 └── instance_type_extra_specs_1.instance_type_id:22 = instance_types.id:1 [outer=(1,22), constraints=(/1: (/NULL - ]; /22: (/NULL - ]), fd=(1)==(22), (22)==(1)]
 
-opt
+opt generic
 select anon_1.instance_types_created_at as anon_1_instance_types_created_at,
        anon_1.instance_types_updated_at as anon_1_instance_types_updated_at,
        anon_1.instance_types_deleted_at as anon_1_instance_types_deleted_at,
@@ -861,7 +861,7 @@ project
       │         └── instance_type_extra_specs_1.deleted:37 = $7 [outer=(37), constraints=(/37: (/NULL - ]), fd=()-->(37)]
       └── filters (true)
 
-opt
+opt generic
 select anon_1.instance_types_created_at as anon_1_instance_types_created_at,
        anon_1.instance_types_updated_at as anon_1_instance_types_updated_at,
        anon_1.instance_types_deleted_at as anon_1_instance_types_deleted_at,
@@ -1026,7 +1026,7 @@ project
       │         └── instance_type_extra_specs_1.deleted:37 = $7 [outer=(37), constraints=(/37: (/NULL - ]), fd=()-->(37)]
       └── filters (true)
 
-opt
+opt generic
 select anon_1.flavors_created_at as anon_1_flavors_created_at,
        anon_1.flavors_updated_at as anon_1_flavors_updated_at,
        anon_1.flavors_id as anon_1_flavors_id,
@@ -1166,7 +1166,7 @@ project
       │    └── filters (true)
       └── filters (true)
 
-opt
+opt generic
 select anon_1.flavors_created_at as anon_1_flavors_created_at,
        anon_1.flavors_updated_at as anon_1_flavors_updated_at,
        anon_1.flavors_id as anon_1_flavors_id,
@@ -1609,7 +1609,7 @@ sort
            └── filters
                 └── instance_type_extra_specs_1.instance_type_id:36 = instance_types.id:1 [outer=(1,36), constraints=(/1: (/NULL - ]; /36: (/NULL - ]), fd=(1)==(36), (36)==(1)]
 
-opt
+opt generic
 select anon_1.instance_types_created_at as anon_1_instance_types_created_at,
        anon_1.instance_types_updated_at as anon_1_instance_types_updated_at,
        anon_1.instance_types_deleted_at as anon_1_instance_types_deleted_at,
@@ -2357,7 +2357,7 @@ sort
            └── filters
                 └── instance_type_extra_specs_1.instance_type_id:50 = instance_types.id:1 [outer=(1,50), constraints=(/1: (/NULL - ]; /50: (/NULL - ]), fd=(1)==(50), (50)==(1)]
 
-opt
+opt generic
 select anon_1.instance_types_created_at as anon_1_instance_types_created_at,
        anon_1.instance_types_updated_at as anon_1_instance_types_updated_at,
        anon_1.instance_types_deleted_at as anon_1_instance_types_deleted_at,
@@ -2511,7 +2511,7 @@ project
       │         └── instance_type_extra_specs_1.deleted:37 = $7 [outer=(37), constraints=(/37: (/NULL - ]), fd=()-->(37)]
       └── filters (true)
 
-opt
+opt generic
 select anon_1.flavors_created_at as anon_1_flavors_created_at,
        anon_1.flavors_updated_at as anon_1_flavors_updated_at,
        anon_1.flavors_id as anon_1_flavors_id,

pkg/sql/opt/xform/testdata/rules/generic

@@ -15,7 +15,7 @@ CREATE TABLE t (
 # GenerateParameterizedJoin
 # --------------------------------------------------
 
-opt expect=GenerateParameterizedJoin
+opt generic expect=GenerateParameterizedJoin
 SELECT * FROM t WHERE k = $1
 ----
 project
@@ -42,7 +42,7 @@ project
       │    └── ($1,)
       └── filters (true)
 
-opt expect=GenerateParameterizedJoin
+opt generic expect=GenerateParameterizedJoin
 SELECT * FROM t WHERE k = $1::INT
 ----
 project
@@ -69,7 +69,7 @@ project
       │    └── ($1,)
       └── filters (true)
 
-opt expect=GenerateParameterizedJoin
+opt generic expect=GenerateParameterizedJoin
 SELECT * FROM t WHERE i = $1 AND s = $2 AND b = $3
 ----
 project
@@ -103,7 +103,7 @@ project
 
 # A placeholder referenced multiple times in the filters should only appear once
 # in the Values expression.
-opt expect=GenerateParameterizedJoin
+opt generic expect=GenerateParameterizedJoin
 SELECT * FROM t WHERE k = $1 AND i = $1
 ----
 project
@@ -133,7 +133,7 @@ project
 
 # The generated join should not be reordered and merge joins should not be
 # explored on it.
-opt expect=GenerateParameterizedJoin expect-not=(ReorderJoins,GenerateMergeJoins)
+opt generic expect=GenerateParameterizedJoin expect-not=(ReorderJoins,GenerateMergeJoins)
 SELECT * FROM t WHERE i = $1
 ----
 project
@@ -165,7 +165,7 @@ project
       │    └── filters (true)
       └── filters (true)
 
-opt expect=GenerateParameterizedJoin
+opt generic expect=GenerateParameterizedJoin
 SELECT * FROM t WHERE k = (SELECT i FROM t WHERE k = $1)
 ----
 project
@@ -209,7 +209,7 @@ project
 
 # TODO(mgartner): The rule doesn't apply because the filters do not reference
 # the placeholder directly. Consider ways to handle cases like this.
-opt
+opt generic
 SELECT * FROM t WHERE k = (SELECT $1::INT)
 ----
 project
@@ -242,7 +242,7 @@ exec-ddl
 CREATE INDEX partial_idx ON t(t) WHERE t IS NOT NULL
 ----
 
-opt expect=GenerateParameterizedJoin
+opt generic expect=GenerateParameterizedJoin
 SELECT * FROM t WHERE t = $1
 ----
 project
@@ -282,7 +282,7 @@ exec-ddl
 CREATE INDEX partial_idx ON t(i, t) WHERE i IS NOT NULL AND t IS NOT NULL
 ----
 
-opt expect=GenerateParameterizedJoin
+opt generic expect=GenerateParameterizedJoin
 SELECT * FROM t WHERE i = $1 AND t = $2
 ----
 project
@@ -322,7 +322,7 @@ exec-ddl
 CREATE INDEX partial_idx ON t(s) WHERE k = i
 ----
 
-opt
+opt generic
 SELECT * FROM t@partial_idx WHERE s = $1 AND k = $2 AND i = $2
 ----
 project
@@ -362,7 +362,7 @@ exec-ddl
 DROP INDEX partial_idx
 ----
 
-opt no-stable-folds expect=GenerateParameterizedJoin
+opt generic no-stable-folds expect=GenerateParameterizedJoin
 SELECT * FROM t WHERE t = now()
 ----
 project
@@ -394,7 +394,7 @@ project
       │    └── filters (true)
       └── filters (true)
 
-opt no-stable-folds expect=GenerateParameterizedJoin
+opt generic no-stable-folds expect=GenerateParameterizedJoin
 SELECT * FROM t WHERE i = $1 AND t = now()
 ----
 project
@@ -426,7 +426,7 @@ project
       │    └── filters (true)
       └── filters (true)
 
-opt no-stable-folds expect=GenerateParameterizedJoin
+opt generic no-stable-folds expect=GenerateParameterizedJoin
 SELECT * FROM t WHERE i = $1 AND t > now()
 ----
 project
@@ -461,7 +461,7 @@ project
       │    └── filters (true)
       └── filters (true)
 
-opt no-stable-folds expect=GenerateParameterizedJoin
+opt generic no-stable-folds expect=GenerateParameterizedJoin
 SELECT * FROM t WHERE i = $1 AND t = now() + $2
 ----
 project
@@ -506,7 +506,7 @@ project
            │    └── filters (true)
            └── filters (true)
 
-opt no-stable-folds expect=GenerateParameterizedJoin
+opt generic no-stable-folds expect=GenerateParameterizedJoin
 SELECT * FROM t WHERE i = $1 AND t = now() + '1 hr'::INTERVAL
 ----
 project
@@ -552,7 +552,7 @@ project
            └── filters (true)
 
 # TODO(mgartner): Apply the rule to stable, non-leaf expressions.
-opt no-stable-folds
+opt generic no-stable-folds
 SELECT * FROM t WHERE t = '2024-01-01 12:00:00'::TIMESTAMP::TIMESTAMPTZ
 ----
 select
@@ -571,7 +571,7 @@ select
 # arguments.
 # TODO(mgartner): We should be able to relax this restriction as long as all the
 # arguments are constants or placeholders.
-opt no-stable-folds expect=GenerateParameterizedJoin
+opt generic no-stable-folds expect=GenerateParameterizedJoin
 SELECT * FROM t WHERE i = $1 AND s = quote_literal(1::INT)
 ----
 project
@@ -607,7 +607,7 @@ project
 # A stable function is not included in the Values expression if its arguments
 # reference a column from the table. This would create an illegal outer column
 # reference in a non-apply-join.
-opt no-stable-folds expect=GenerateParameterizedJoin
+opt generic no-stable-folds expect=GenerateParameterizedJoin
 SELECT * FROM t WHERE i = $1 AND s = quote_literal(i)
 ----
 project
@@ -642,7 +642,7 @@ project
 
 # The rule does not match if there are no placeholders or stable expressions in
 # the filters.
-opt expect-not=GenerateParameterizedJoin
+opt generic expect-not=GenerateParameterizedJoin
 SELECT * FROM t WHERE i = 1 AND s = 'foo'
 ----
 index-join t
@@ -654,3 +654,21 @@ index-join t
       ├── constraint: /2/3/4/1: [/1/'foo' - /1/'foo']
       ├── key: (1)
       └── fd: ()-->(2,3), (1)-->(4)
+
+# The rule does not match if generic optimizations are disabled.
+opt expect-not=GenerateParameterizedJoin
+SELECT * FROM t WHERE k = $1 AND s = quote_literal(1::INT)
+----
+select
+ ├── columns: k:1!null i:2 s:3!null b:4 t:5
+ ├── cardinality: [0 - 1]
+ ├── has-placeholder
+ ├── key: ()
+ ├── fd: ()-->(1-5)
+ ├── scan t
+ │    ├── columns: k:1!null i:2 s:3 b:4 t:5
+ │    ├── key: (1)
+ │    └── fd: (1)-->(2-5)
+ └── filters
+      ├── k:1 = $1 [outer=(1), constraints=(/1: (/NULL - ]), fd=()-->(1)]
+      └── s:3 = e'\'1\'' [outer=(3), constraints=(/3: [/e'\'1\'' - /e'\'1\'']; tight), fd=()-->(3)]