Merge pull request #103412 from msirek/backport23.1-100242

release-23.1: xform: constrain index scans on computed columns in more cases

pkg/sql/exec_util.go

@@ -3522,6 +3522,10 @@ func (m *sessionDataMutator) SetOptimizerHoistUncorrelatedEqualitySubqueries(val
 	m.data.OptimizerHoistUncorrelatedEqualitySubqueries = val
 }
 
+func (m *sessionDataMutator) SetOptimizerUseImprovedComputedColumnFiltersDerivation(val bool) {
+	m.data.OptimizerUseImprovedComputedColumnFiltersDerivation = val
+}
+
 func (m *sessionDataMutator) SetEnableCreateStatsUsingExtremes(val bool) {
 	m.data.EnableCreateStatsUsingExtremes = val
 }

pkg/sql/opt/BUILD.bazel

@@ -69,6 +69,7 @@ go_test(
         "//pkg/sql/opt/cat",
         "//pkg/sql/opt/memo",
         "//pkg/sql/opt/norm",
+        "//pkg/sql/opt/props",
         "//pkg/sql/opt/testutils/testcat",
         "//pkg/sql/privilege",
         "//pkg/sql/sem/catid",

pkg/sql/opt/lookupjoin/constraint_builder_test.go

@@ -106,7 +106,9 @@ func TestLookupConstraints(t *testing.T) {
 							return 0, opt.ColSet{}, err
 						}
 						compExpr := f.Memo().RootExpr().(opt.ScalarExpr)
-						md.TableMeta(tableID).AddComputedCol(colID, compExpr)
+						var sharedProps props.Shared
+						memo.BuildSharedProps(compExpr, &sharedProps, &evalCtx)
+						md.TableMeta(tableID).AddComputedCol(colID, compExpr, sharedProps.OuterCols)
 					}
 				}
 				return tableID, cols, nil

pkg/sql/opt/memo/memo.go

@@ -135,35 +135,36 @@ type Memo struct {
 	// planning. We need to cross-check these before reusing a cached memo.
 	// NOTE: If you add new fields here, be sure to add them to the relevant
 	//       fields in explain_bundle.go.
-	reorderJoinsLimit                      int
-	zigzagJoinEnabled                      bool
-	useForecasts                           bool
-	useHistograms                          bool
-	useMultiColStats                       bool
-	useNotVisibleIndex                     bool
-	localityOptimizedSearch                bool
-	safeUpdates                            bool
-	preferLookupJoinsForFKs                bool
-	saveTablesPrefix                       string
-	dateStyle                              pgdate.DateStyle
-	intervalStyle                          duration.IntervalStyle
-	propagateInputOrdering                 bool
-	disallowFullTableScans                 bool
-	largeFullScanRows                      float64
-	nullOrderedLast                        bool
-	costScansWithDefaultColSize            bool
-	allowUnconstrainedNonCoveringIndexScan bool
-	testingOptimizerRandomSeed             int64
-	testingOptimizerCostPerturbation       float64
-	testingOptimizerDisableRuleProbability float64
-	enforceHomeRegion                      bool
-	variableInequalityLookupJoinEnabled    bool
-	allowOrdinalColumnReferences           bool
-	useImprovedDisjunctionStats            bool
-	useLimitOrderingForStreamingGroupBy    bool
-	useImprovedSplitDisjunctionForJoins    bool
-	alwaysUseHistograms                    bool
-	hoistUncorrelatedEqualitySubqueries    bool
+	reorderJoinsLimit                          int
+	zigzagJoinEnabled                          bool
+	useForecasts                               bool
+	useHistograms                              bool
+	useMultiColStats                           bool
+	useNotVisibleIndex                         bool
+	localityOptimizedSearch                    bool
+	safeUpdates                                bool
+	preferLookupJoinsForFKs                    bool
+	saveTablesPrefix                           string
+	dateStyle                                  pgdate.DateStyle
+	intervalStyle                              duration.IntervalStyle
+	propagateInputOrdering                     bool
+	disallowFullTableScans                     bool
+	largeFullScanRows                          float64
+	nullOrderedLast                            bool
+	costScansWithDefaultColSize                bool
+	allowUnconstrainedNonCoveringIndexScan     bool
+	testingOptimizerRandomSeed                 int64
+	testingOptimizerCostPerturbation           float64
+	testingOptimizerDisableRuleProbability     float64
+	enforceHomeRegion                          bool
+	variableInequalityLookupJoinEnabled        bool
+	allowOrdinalColumnReferences               bool
+	useImprovedDisjunctionStats                bool
+	useLimitOrderingForStreamingGroupBy        bool
+	useImprovedSplitDisjunctionForJoins        bool
+	alwaysUseHistograms                        bool
+	hoistUncorrelatedEqualitySubqueries        bool
+	useImprovedComputedColumnFiltersDerivation bool
 
 	// curRank is the highest currently in-use scalar expression rank.
 	curRank opt.ScalarRank
@@ -193,36 +194,37 @@ func (m *Memo) Init(ctx context.Context, evalCtx *eval.Context) {
 	// This initialization pattern ensures that fields are not unwittingly
 	// reused. Field reuse must be explicit.
 	*m = Memo{
-		metadata:                               m.metadata,
-		reorderJoinsLimit:                      int(evalCtx.SessionData().ReorderJoinsLimit),
-		zigzagJoinEnabled:                      evalCtx.SessionData().ZigzagJoinEnabled,
-		useForecasts:                           evalCtx.SessionData().OptimizerUseForecasts,
-		useHistograms:                          evalCtx.SessionData().OptimizerUseHistograms,
-		useMultiColStats:                       evalCtx.SessionData().OptimizerUseMultiColStats,
-		useNotVisibleIndex:                     evalCtx.SessionData().OptimizerUseNotVisibleIndexes,
-		localityOptimizedSearch:                evalCtx.SessionData().LocalityOptimizedSearch,
-		safeUpdates:                            evalCtx.SessionData().SafeUpdates,
-		preferLookupJoinsForFKs:                evalCtx.SessionData().PreferLookupJoinsForFKs,
-		saveTablesPrefix:                       evalCtx.SessionData().SaveTablesPrefix,
-		dateStyle:                              evalCtx.SessionData().GetDateStyle(),
-		intervalStyle:                          evalCtx.SessionData().GetIntervalStyle(),
-		propagateInputOrdering:                 evalCtx.SessionData().PropagateInputOrdering,
-		disallowFullTableScans:                 evalCtx.SessionData().DisallowFullTableScans,
-		largeFullScanRows:                      evalCtx.SessionData().LargeFullScanRows,
-		nullOrderedLast:                        evalCtx.SessionData().NullOrderedLast,
-		costScansWithDefaultColSize:            evalCtx.SessionData().CostScansWithDefaultColSize,
-		allowUnconstrainedNonCoveringIndexScan: evalCtx.SessionData().UnconstrainedNonCoveringIndexScanEnabled,
-		testingOptimizerRandomSeed:             evalCtx.SessionData().TestingOptimizerRandomSeed,
-		testingOptimizerCostPerturbation:       evalCtx.SessionData().TestingOptimizerCostPerturbation,
-		testingOptimizerDisableRuleProbability: evalCtx.SessionData().TestingOptimizerDisableRuleProbability,
-		enforceHomeRegion:                      evalCtx.SessionData().EnforceHomeRegion,
-		variableInequalityLookupJoinEnabled:    evalCtx.SessionData().VariableInequalityLookupJoinEnabled,
-		allowOrdinalColumnReferences:           evalCtx.SessionData().AllowOrdinalColumnReferences,
-		useImprovedDisjunctionStats:            evalCtx.SessionData().OptimizerUseImprovedDisjunctionStats,
-		useLimitOrderingForStreamingGroupBy:    evalCtx.SessionData().OptimizerUseLimitOrderingForStreamingGroupBy,
-		useImprovedSplitDisjunctionForJoins:    evalCtx.SessionData().OptimizerUseImprovedSplitDisjunctionForJoins,
-		alwaysUseHistograms:                    evalCtx.SessionData().OptimizerAlwaysUseHistograms,
-		hoistUncorrelatedEqualitySubqueries:    evalCtx.SessionData().OptimizerHoistUncorrelatedEqualitySubqueries,
+		metadata:                                   m.metadata,
+		reorderJoinsLimit:                          int(evalCtx.SessionData().ReorderJoinsLimit),
+		zigzagJoinEnabled:                          evalCtx.SessionData().ZigzagJoinEnabled,
+		useForecasts:                               evalCtx.SessionData().OptimizerUseForecasts,
+		useHistograms:                              evalCtx.SessionData().OptimizerUseHistograms,
+		useMultiColStats:                           evalCtx.SessionData().OptimizerUseMultiColStats,
+		useNotVisibleIndex:                         evalCtx.SessionData().OptimizerUseNotVisibleIndexes,
+		localityOptimizedSearch:                    evalCtx.SessionData().LocalityOptimizedSearch,
+		safeUpdates:                                evalCtx.SessionData().SafeUpdates,
+		preferLookupJoinsForFKs:                    evalCtx.SessionData().PreferLookupJoinsForFKs,
+		saveTablesPrefix:                           evalCtx.SessionData().SaveTablesPrefix,
+		dateStyle:                                  evalCtx.SessionData().GetDateStyle(),
+		intervalStyle:                              evalCtx.SessionData().GetIntervalStyle(),
+		propagateInputOrdering:                     evalCtx.SessionData().PropagateInputOrdering,
+		disallowFullTableScans:                     evalCtx.SessionData().DisallowFullTableScans,
+		largeFullScanRows:                          evalCtx.SessionData().LargeFullScanRows,
+		nullOrderedLast:                            evalCtx.SessionData().NullOrderedLast,
+		costScansWithDefaultColSize:                evalCtx.SessionData().CostScansWithDefaultColSize,
+		allowUnconstrainedNonCoveringIndexScan:     evalCtx.SessionData().UnconstrainedNonCoveringIndexScanEnabled,
+		testingOptimizerRandomSeed:                 evalCtx.SessionData().TestingOptimizerRandomSeed,
+		testingOptimizerCostPerturbation:           evalCtx.SessionData().TestingOptimizerCostPerturbation,
+		testingOptimizerDisableRuleProbability:     evalCtx.SessionData().TestingOptimizerDisableRuleProbability,
+		enforceHomeRegion:                          evalCtx.SessionData().EnforceHomeRegion,
+		variableInequalityLookupJoinEnabled:        evalCtx.SessionData().VariableInequalityLookupJoinEnabled,
+		allowOrdinalColumnReferences:               evalCtx.SessionData().AllowOrdinalColumnReferences,
+		useImprovedDisjunctionStats:                evalCtx.SessionData().OptimizerUseImprovedDisjunctionStats,
+		useLimitOrderingForStreamingGroupBy:        evalCtx.SessionData().OptimizerUseLimitOrderingForStreamingGroupBy,
+		useImprovedSplitDisjunctionForJoins:        evalCtx.SessionData().OptimizerUseImprovedSplitDisjunctionForJoins,
+		alwaysUseHistograms:                        evalCtx.SessionData().OptimizerAlwaysUseHistograms,
+		hoistUncorrelatedEqualitySubqueries:        evalCtx.SessionData().OptimizerHoistUncorrelatedEqualitySubqueries,
+		useImprovedComputedColumnFiltersDerivation: evalCtx.SessionData().OptimizerUseImprovedComputedColumnFiltersDerivation,
 	}
 	m.metadata.Init()
 	m.logPropsBuilder.init(ctx, evalCtx, m)
@@ -365,7 +367,8 @@ func (m *Memo) IsStale(
 		m.useLimitOrderingForStreamingGroupBy != evalCtx.SessionData().OptimizerUseLimitOrderingForStreamingGroupBy ||
 		m.useImprovedSplitDisjunctionForJoins != evalCtx.SessionData().OptimizerUseImprovedSplitDisjunctionForJoins ||
 		m.alwaysUseHistograms != evalCtx.SessionData().OptimizerAlwaysUseHistograms ||
-		m.hoistUncorrelatedEqualitySubqueries != evalCtx.SessionData().OptimizerHoistUncorrelatedEqualitySubqueries {
+		m.hoistUncorrelatedEqualitySubqueries != evalCtx.SessionData().OptimizerHoistUncorrelatedEqualitySubqueries ||
+		m.useImprovedComputedColumnFiltersDerivation != evalCtx.SessionData().OptimizerUseImprovedComputedColumnFiltersDerivation {
 		return true, nil
 	}
 

pkg/sql/opt/memo/memo_test.go

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

pkg/sql/opt/metadata.go

@@ -652,6 +652,7 @@ func (md *Metadata) DuplicateTable(
 	// Create new computed column expressions by remapping the column IDs in
 	// each ScalarExpr.
 	var computedCols map[ColumnID]ScalarExpr
+	var referencedColsInComputedExpressions ColSet
 	if len(tabMeta.ComputedCols) > 0 {
 		computedCols = make(map[ColumnID]ScalarExpr, len(tabMeta.ComputedCols))
 		for colID, e := range tabMeta.ComputedCols {
@@ -661,6 +662,9 @@ func (md *Metadata) DuplicateTable(
 			}
 			computedCols[ColumnID(newColID)] = remapColumnIDs(e, colMap)
 		}
+		// Add columns present in newScalarExpr to referencedColsInComputedExpressions.
+		referencedColsInComputedExpressions =
+			tabMeta.ColsInComputedColsExpressions.CopyAndMaybeRemap(colMap)
 	}
 
 	// Create new partial index predicate expressions by remapping the column
@@ -693,15 +697,16 @@ func (md *Metadata) DuplicateTable(
 	}
 
 	newTabMeta := TableMeta{
-		MetaID:                   newTabID,
-		Table:                    tabMeta.Table,
-		Alias:                    tabMeta.Alias,
-		IgnoreForeignKeys:        tabMeta.IgnoreForeignKeys,
-		Constraints:              constraints,
-		ComputedCols:             computedCols,
-		partialIndexPredicates:   partialIndexPredicates,
-		indexPartitionLocalities: tabMeta.indexPartitionLocalities,
-		checkConstraintsStats:    checkConstraintsStats,
+		MetaID:                        newTabID,
+		Table:                         tabMeta.Table,
+		Alias:                         tabMeta.Alias,
+		IgnoreForeignKeys:             tabMeta.IgnoreForeignKeys,
+		Constraints:                   constraints,
+		ComputedCols:                  computedCols,
+		ColsInComputedColsExpressions: referencedColsInComputedExpressions,
+		partialIndexPredicates:        partialIndexPredicates,
+		indexPartitionLocalities:      tabMeta.indexPartitionLocalities,
+		checkConstraintsStats:         checkConstraintsStats,
 	}
 	md.tables = append(md.tables, newTabMeta)
 	regionConfig, ok := md.TableAnnotation(tabID, regionConfigAnnID).(*multiregion.RegionConfig)

pkg/sql/opt/metadata_test.go

@@ -23,6 +23,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/cat"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/norm"
+	"github.com/cockroachdb/cockroach/pkg/sql/opt/props"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/testutils/testcat"
 	"github.com/cockroachdb/cockroach/pkg/sql/privilege"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/catid"
@@ -82,7 +83,9 @@ func TestMetadata(t *testing.T) {
 	}
 	tabMeta := md.TableMeta(tabID)
 	tabMeta.SetConstraints(scalar)
-	tabMeta.AddComputedCol(cmpID, scalar)
+	var sharedProps props.Shared
+	memo.BuildSharedProps(scalar, &sharedProps, &evalCtx)
+	tabMeta.AddComputedCol(cmpID, scalar, sharedProps.OuterCols)
 	tabMeta.AddPartialIndexPredicate(0, scalar)
 	if md.AddSequence(&testcat.Sequence{SeqID: 100}) != seqID {
 		t.Fatalf("unexpected sequence id")
@@ -143,6 +146,10 @@ func TestMetadata(t *testing.T) {
 		t.Fatalf("expected computed column expression to be copied")
 	}
 
+	if !tabMeta.ColsInComputedColsExpressions.Equals(tabMetaNew.ColsInComputedColsExpressions) {
+		t.Fatalf("expected computed column expression referenced columns to be copied")
+	}
+
 	partialIdxPredPtr := reflect.ValueOf(tabMeta.PartialIndexPredicatesUnsafe()).Pointer()
 	newPartialIdxPredPtr := reflect.ValueOf(tabMetaNew.PartialIndexPredicatesUnsafe()).Pointer()
 	if newPartialIdxPredPtr == partialIdxPredPtr {
@@ -382,21 +389,27 @@ func TestIndexColumns(t *testing.T) {
 
 // TestDuplicateTable tests that we can extract a set of columns from an index ordinal.
 func TestDuplicateTable(t *testing.T) {
+	evalCtx := eval.MakeTestingEvalContext(cluster.MakeTestingClusterSettings())
+	var f norm.Factory
+	f.Init(context.Background(), &evalCtx, nil /* catalog */)
+	md := f.Metadata()
 	cat := testcat.New()
 	_, err := cat.ExecuteDDL("CREATE TABLE a (b BOOL, b2 BOOL, INDEX (b2) WHERE b)")
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	var md opt.Metadata
 	tn := tree.NewUnqualifiedTableName("a")
 	a := md.AddTable(cat.Table(tn), tn)
 	b := a.ColumnID(0)
 	b2 := a.ColumnID(1)
 
 	tabMeta := md.TableMeta(a)
 	tabMeta.SetConstraints(&memo.VariableExpr{Col: b})
-	tabMeta.AddComputedCol(b2, &memo.VariableExpr{Col: b})
+	scalar := &memo.VariableExpr{Col: b}
+	var sharedProps props.Shared
+	memo.BuildSharedProps(scalar, &sharedProps, &evalCtx)
+	tabMeta.AddComputedCol(b2, scalar, sharedProps.OuterCols)
 	tabMeta.AddPartialIndexPredicate(1, &memo.VariableExpr{Col: b})
 
 	// remap is a simple function that can only remap column IDs in a
@@ -434,6 +447,18 @@ func TestDuplicateTable(t *testing.T) {
 		t.Errorf("expected computed column to reference new column ID %d, got %d", dupB, col)
 	}
 
+	if tabMeta.ColsInComputedColsExpressions.Equals(dupTabMeta.ColsInComputedColsExpressions) {
+		t.Fatalf("expected computed column expression referenced columns to hold new column ids")
+	}
+
+	if dupTabMeta.ColsInComputedColsExpressions.Empty() {
+		t.Fatalf("expected computed column expression referenced columns to not be empty")
+	}
+
+	if tabMeta.ColsInComputedColsExpressions.Len() != dupTabMeta.ColsInComputedColsExpressions.Len() {
+		t.Fatalf("expected same number of computed column expression referenced columns")
+	}
+
 	pred, isPartialIndex := dupTabMeta.PartialIndexPredicate(1)
 	if !isPartialIndex {
 		t.Fatalf("expected partial index predicates to be duplicated")

pkg/sql/opt/optbuilder/select.go

@@ -876,7 +876,7 @@ func (b *Builder) addComputedColsForTable(tabMeta *opt.TableMeta) {
 			var sharedProps props.Shared
 			memo.BuildSharedProps(scalar, &sharedProps, b.evalCtx)
 			if !sharedProps.VolatilitySet.HasStable() && !sharedProps.VolatilitySet.HasVolatile() {
-				tabMeta.AddComputedCol(colID, scalar)
+				tabMeta.AddComputedCol(colID, scalar, sharedProps.OuterCols)
 			}
 		}
 	}

pkg/sql/opt/table_meta.go

@@ -174,6 +174,10 @@ type TableMeta struct {
 	// Computed columns with non-immutable operators are omitted.
 	ComputedCols map[ColumnID]ScalarExpr
 
+	// ColsInComputedColsExpressions is the set of all columns referenced in the
+	// expressions used to build the column data of computed columns.
+	ColsInComputedColsExpressions ColSet
+
 	// partialIndexPredicates is a map from index ordinals on the table to
 	// *FiltersExprs representing the predicate on the corresponding partial
 	// index. If an index is not a partial index, it will not have an entry in
@@ -249,6 +253,7 @@ func (tm *TableMeta) copyFrom(from *TableMeta, copyScalarFn func(Expr) Expr) {
 		for col, e := range from.ComputedCols {
 			tm.ComputedCols[col] = copyScalarFn(e).(ScalarExpr)
 		}
+		tm.ColsInComputedColsExpressions = from.ColsInComputedColsExpressions
 	}
 
 	if from.partialIndexPredicates != nil {
@@ -335,12 +340,15 @@ func (tm *TableMeta) SetConstraints(constraints ScalarExpr) {
 	tm.Constraints = constraints
 }
 
-// AddComputedCol adds a computed column expression to the table's metadata.
-func (tm *TableMeta) AddComputedCol(colID ColumnID, computedCol ScalarExpr) {
+// AddComputedCol adds a computed column expression to the table's metadata and
+// also adds any referenced columns in the `computedCol` expression to the
+// table's metadata.
+func (tm *TableMeta) AddComputedCol(colID ColumnID, computedCol ScalarExpr, outerCols ColSet) {
 	if tm.ComputedCols == nil {
 		tm.ComputedCols = make(map[ColumnID]ScalarExpr)
 	}
 	tm.ComputedCols[colID] = computedCol
+	tm.ColsInComputedColsExpressions.UnionWith(outerCols)
 }
 
 // ComputedColExpr returns the computed expression for the given column, if it