release-23.2.4-rc: opt/memo: use virtual column stats in statistics builder

pkg/sql/exec_util.go

@@ -3757,6 +3757,10 @@ func (m *sessionDataMutator) SetDistSQLPlanGatewayBias(val int64) {
 	m.data.DistsqlPlanGatewayBias = val
 }
 
+func (m *sessionDataMutator) SetOptimizerUseVirtualComputedColumnStats(val bool) {
+	m.data.OptimizerUseVirtualComputedColumnStats = val
+}
+
 // Utility functions related to scrubbing sensitive information on SQL Stats.
 
 // quantizeCounts ensures that the Count field in the

pkg/sql/logictest/testdata/logic_test/information_schema

@@ -5578,6 +5578,7 @@ optimizer_use_lock_op_for_serializable                     off
 optimizer_use_multicol_stats                               on
 optimizer_use_not_visible_indexes                          off
 optimizer_use_provided_ordering_fix                        on
+optimizer_use_virtual_computed_column_stats                off
 override_multi_region_zone_config                          off
 parallelize_multi_key_lookup_joins_enabled                 off
 password_encryption                                        scram-sha-256

pkg/sql/logictest/testdata/logic_test/pg_catalog

@@ -2895,6 +2895,7 @@ optimizer_use_lock_op_for_serializable                     off                 N
 optimizer_use_multicol_stats                               on                  NULL      NULL        NULL        string
 optimizer_use_not_visible_indexes                          off                 NULL      NULL        NULL        string
 optimizer_use_provided_ordering_fix                        on                  NULL      NULL        NULL        string
+optimizer_use_virtual_computed_column_stats                off                 NULL      NULL        NULL        string
 override_multi_region_zone_config                          off                 NULL      NULL        NULL        string
 parallelize_multi_key_lookup_joins_enabled                 off                 NULL      NULL        NULL        string
 password_encryption                                        scram-sha-256       NULL      NULL        NULL        string
@@ -3068,6 +3069,7 @@ optimizer_use_lock_op_for_serializable                     off                 N
 optimizer_use_multicol_stats                               on                  NULL  user     NULL      on                  on
 optimizer_use_not_visible_indexes                          off                 NULL  user     NULL      off                 off
 optimizer_use_provided_ordering_fix                        on                  NULL  user     NULL      on                  on
+optimizer_use_virtual_computed_column_stats                off                 NULL  user     NULL      off                 off
 override_multi_region_zone_config                          off                 NULL  user     NULL      off                 off
 parallelize_multi_key_lookup_joins_enabled                 off                 NULL  user     NULL      false               false
 password_encryption                                        scram-sha-256       NULL  user     NULL      scram-sha-256       scram-sha-256
@@ -3240,6 +3242,7 @@ optimizer_use_lock_op_for_serializable                     NULL    NULL     NULL
 optimizer_use_multicol_stats                               NULL    NULL     NULL     NULL        NULL
 optimizer_use_not_visible_indexes                          NULL    NULL     NULL     NULL        NULL
 optimizer_use_provided_ordering_fix                        NULL    NULL     NULL     NULL        NULL
+optimizer_use_virtual_computed_column_stats                NULL    NULL     NULL     NULL        NULL
 override_multi_region_zone_config                          NULL    NULL     NULL     NULL        NULL
 parallelize_multi_key_lookup_joins_enabled                 NULL    NULL     NULL     NULL        NULL
 password_encryption                                        NULL    NULL     NULL     NULL        NULL

pkg/sql/logictest/testdata/logic_test/show_source

@@ -133,6 +133,7 @@ optimizer_use_lock_op_for_serializable                     off
 optimizer_use_multicol_stats                               on
 optimizer_use_not_visible_indexes                          off
 optimizer_use_provided_ordering_fix                        on
+optimizer_use_virtual_computed_column_stats                off
 override_multi_region_zone_config                          off
 parallelize_multi_key_lookup_joins_enabled                 off
 password_encryption                                        scram-sha-256

pkg/sql/opt/exec/execbuilder/testdata/explain_redact

@@ -700,7 +700,7 @@ upsert bc
 query T
 EXPLAIN (OPT, MEMO, REDACT) INSERT INTO bc SELECT a::float + 1 FROM a ON CONFLICT (b) DO UPDATE SET b = bc.b + 100
 ----
-memo (optimized, ~33KB, required=[presentation: info:19] [distribution: test])
+memo (optimized, ~34KB, required=[presentation: info:19] [distribution: test])
  ├── G1: (explain G2 [distribution: test])
  │    └── [presentation: info:19] [distribution: test]
  │         ├── best: (explain G2="[distribution: test]" [distribution: test])
@@ -1690,7 +1690,7 @@ project
 query T
 EXPLAIN (OPT, MEMO, REDACT) SELECT * FROM bc WHERE b >= 1.0 AND b < 2.0
 ----
-memo (optimized, ~11KB, required=[presentation: info:5] [distribution: test])
+memo (optimized, ~12KB, required=[presentation: info:5] [distribution: test])
  ├── G1: (explain G2 [presentation: b:1,c:2] [distribution: test])
  │    └── [presentation: info:5] [distribution: test]
  │         ├── best: (explain G2="[presentation: b:1,c:2] [distribution: test]" [presentation: b:1,c:2] [distribution: test])
@@ -1841,7 +1841,7 @@ select
 query T
 EXPLAIN (OPT, MEMO, REDACT) SELECT * FROM g WHERE (g || 'abc') LIKE '%ggg%'
 ----
-memo (optimized, ~12KB, required=[presentation: info:7] [distribution: test])
+memo (optimized, ~13KB, required=[presentation: info:7] [distribution: test])
  ├── G1: (explain G2 [presentation: g:1] [distribution: test])
  │    └── [presentation: info:7] [distribution: test]
  │         ├── best: (explain G2="[presentation: g:1] [distribution: test]" [presentation: g:1] [distribution: test])
@@ -2439,7 +2439,7 @@ project
 query T
 EXPLAIN (OPT, MEMO, REDACT) SELECT * FROM bc JOIN f ON b = f + 1
 ----
-memo (optimized, ~26KB, required=[presentation: info:10] [distribution: test])
+memo (optimized, ~27KB, required=[presentation: info:10] [distribution: test])
  ├── G1: (explain G2 [presentation: b:1,c:2,f:5] [distribution: test])
  │    └── [presentation: info:10] [distribution: test]
  │         ├── best: (explain G2="[presentation: b:1,c:2,f:5] [distribution: test]" [presentation: b:1,c:2,f:5] [distribution: test])
@@ -2824,7 +2824,7 @@ project
 query T
 EXPLAIN (OPT, MEMO, REDACT) SELECT f, g FROM f, LATERAL (SELECT count(DISTINCT c + f + 1) * 2 AS g FROM bc WHERE b * f < 10)
 ----
-memo (optimized, ~32KB, required=[presentation: info:12] [distribution: test])
+memo (optimized, ~33KB, required=[presentation: info:12] [distribution: test])
  ├── G1: (explain G2 [presentation: f:1,g:11] [distribution: test])
  │    └── [presentation: info:12] [distribution: test]
  │         ├── best: (explain G2="[presentation: f:1,g:11] [distribution: test]" [presentation: f:1,g:11] [distribution: test])
@@ -3103,7 +3103,7 @@ anti-join-apply
 query T
 EXPLAIN (OPT, MEMO, REDACT) SELECT * FROM a WHERE a > ALL (SELECT c::int + 2 FROM bc WHERE b > a::float * 3)
 ----
-memo (optimized, ~21KB, required=[presentation: info:10] [distribution: test])
+memo (optimized, ~22KB, required=[presentation: info:10] [distribution: test])
  ├── G1: (explain G2 [presentation: a:1] [distribution: test])
  │    └── [presentation: info:10] [distribution: test]
  │         ├── best: (explain G2="[presentation: a:1] [distribution: test]" [presentation: a:1] [distribution: test])
@@ -3807,7 +3807,7 @@ project
 query T
 EXPLAIN (OPT, MEMO, REDACT) SELECT min(c || 'cccc') OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING EXCLUDE CURRENT ROW) FROM cd
 ----
-memo (optimized, ~11KB, required=[presentation: info:8] [distribution: test])
+memo (optimized, ~12KB, required=[presentation: info:8] [distribution: test])
  ├── G1: (explain G2 [presentation: min:6] [distribution: test])
  │    └── [presentation: info:8] [distribution: test]
  │         ├── best: (explain G2="[presentation: min:6] [distribution: test]" [presentation: min:6] [distribution: test])

pkg/sql/opt/memo/memo.go

@@ -26,6 +26,14 @@ import (
 	"github.com/cockroachdb/redact"
 )
 
+// replaceFunc is the callback function passed to norm.Factory.Replace. It is
+// copied here from norm.ReplaceFunc to avoid a circular dependency.
+type ReplaceFunc func(e opt.Expr) opt.Expr
+
+// replacer is a wrapper around norm.Factory.Replace, so that we can call it
+// without creating a circular dependency.
+type replacer func(e opt.Expr, replace ReplaceFunc) opt.Expr
+
 // Memo is a data structure for efficiently storing a forest of query plans.
 // Conceptually, the memo is composed of a numbered set of equivalency classes
 // called groups where each group contains a set of logically equivalent
@@ -117,6 +125,10 @@ type Memo struct {
 	// most one instance of each expression in the memo.
 	interner interner
 
+	// replacer is a wrapper around norm.Factory.Replace, used by statistics
+	// builder to rewrite some expressions when calculating stats.
+	replacer replacer
+
 	// logPropsBuilder is inlined in the memo so that it can be reused each time
 	// scalar or relational properties need to be built.
 	logPropsBuilder logicalPropsBuilder
@@ -174,6 +186,7 @@ type Memo struct {
 	sharedLockingForSerializable               bool
 	useLockOpForSerializable                   bool
 	useProvidedOrderingFix                     bool
+	useVirtualComputedColumnStats              bool
 
 	// txnIsoLevel is the isolation level under which the plan was created. This
 	// affects the planning of some locking operations, so it must be included in
@@ -246,12 +259,17 @@ func (m *Memo) Init(ctx context.Context, evalCtx *eval.Context) {
 		sharedLockingForSerializable:               evalCtx.SessionData().SharedLockingForSerializable,
 		useLockOpForSerializable:                   evalCtx.SessionData().OptimizerUseLockOpForSerializable,
 		useProvidedOrderingFix:                     evalCtx.SessionData().OptimizerUseProvidedOrderingFix,
+		useVirtualComputedColumnStats:              evalCtx.SessionData().OptimizerUseVirtualComputedColumnStats,
 		txnIsoLevel:                                evalCtx.TxnIsoLevel,
 	}
 	m.metadata.Init()
 	m.logPropsBuilder.init(ctx, evalCtx, m)
 }
 
+func (m *Memo) SetReplacer(replacer replacer) {
+	m.replacer = replacer
+}
+
 // AllowUnconstrainedNonCoveringIndexScan indicates whether unconstrained
 // non-covering index scans are enabled.
 func (m *Memo) AllowUnconstrainedNonCoveringIndexScan() bool {
@@ -392,6 +410,7 @@ func (m *Memo) IsStale(
 		m.sharedLockingForSerializable != evalCtx.SessionData().SharedLockingForSerializable ||
 		m.useLockOpForSerializable != evalCtx.SessionData().OptimizerUseLockOpForSerializable ||
 		m.useProvidedOrderingFix != evalCtx.SessionData().OptimizerUseProvidedOrderingFix ||
+		m.useVirtualComputedColumnStats != evalCtx.SessionData().OptimizerUseVirtualComputedColumnStats ||
 		m.txnIsoLevel != evalCtx.TxnIsoLevel {
 		return true, nil
 	}

pkg/sql/opt/memo/memo_test.go

@@ -424,6 +424,12 @@ func TestMemoIsStale(t *testing.T) {
 	evalCtx.SessionData().OptimizerUseProvidedOrderingFix = false
 	notStale()
 
+	// Stale optimizer_use_virtual_computed_column_stats.
+	evalCtx.SessionData().OptimizerUseVirtualComputedColumnStats = true
+	stale()
+	evalCtx.SessionData().OptimizerUseVirtualComputedColumnStats = false
+	notStale()
+
 	// User no longer has access to view.
 	catalog.View(tree.NewTableNameWithSchema("t", catconstants.PublicSchemaName, "abcview")).Revoked = true
 	_, err = o.Memo().IsStale(ctx, &evalCtx, catalog)