Merge pull request #44749 from RaduBerinde/backport19.2-44668

release-19.2: opt: fix remaining filters when using partitioned constraints

pkg/ccl/logictestccl/testdata/logic_test/partitioning_constrained_scans

@@ -238,3 +238,25 @@ query T
 SELECT feature_name FROM crdb_internal.feature_usage WHERE feature_name='sql.partitioning.partition-constrained-scan' AND usage_count > 0
 ----
 sql.partitioning.partition-constrained-scan
+
+# Regression test for #44154: a remaining filter that is not identical to an
+# input filter should not be dropped.
+statement ok
+CREATE TABLE t0(c0 BOOL UNIQUE, c1 BOOL CHECK (true))
+
+statement ok
+INSERT INTO t0(c0) VALUES (true)
+
+query T
+EXPLAIN (OPT) SELECT * FROM t0 WHERE t0.c0 AND (c1 OR (c0 > false AND c0 < false))
+----
+select
+ ├── index-join t0
+ │    └── scan t0@t0_c0_key
+ │         └── constraint: /1: [/true - /true]
+ └── filters
+      └── c1 OR (c0 < false)
+
+query BB
+SELECT * FROM t0 WHERE t0.c0 AND (c1 OR (c0 > false AND c0 < false))
+----

pkg/sql/opt/idxconstraint/index_constraints.go

@@ -1061,35 +1061,53 @@ func (c *indexConstraintCtx) simplifyFilter(
 type Instance struct {
 	indexConstraintCtx
 
-	filters memo.FiltersExpr
+	requiredFilters memo.FiltersExpr
+	// allFilters includes requiredFilters along with optional filters that don't
+	// need to generate remaining filters (see Instance.Init()).
+	allFilters memo.FiltersExpr
 
 	constraint   constraint.Constraint
 	consolidated constraint.Constraint
 	tight        bool
 	initialized  bool
 }
 
-// Init processes the filter and calculates the spans.
+// Init processes the filters and calculates the spans.
+//
+// Optional filters are filters that can be used for determining spans, but
+// they need not generate remaining filters. This is e.g. used for check
+// constraints that can help generate better spans but don't actually need to be
+// enforced.
 func (ic *Instance) Init(
-	filters memo.FiltersExpr,
+	requiredFilters memo.FiltersExpr,
+	optionalFilters memo.FiltersExpr,
 	columns []opt.OrderingColumn,
 	notNullCols opt.ColSet,
 	isInverted bool,
 	evalCtx *tree.EvalContext,
 	factory *norm.Factory,
 ) {
-	ic.filters = filters
+	ic.requiredFilters = requiredFilters
+	if len(optionalFilters) == 0 {
+		ic.allFilters = requiredFilters
+	} else {
+		// Force allocation of a bigger slice.
+		// TODO(radu): we should keep the required and optional filters separate and
+		// add a helper for iterating through all of them.
+		ic.allFilters = requiredFilters[:len(requiredFilters):len(requiredFilters)]
+		ic.allFilters = append(ic.allFilters, optionalFilters...)
+	}
 	ic.indexConstraintCtx.init(columns, notNullCols, isInverted, evalCtx, factory)
-	constraints := make([]*constraint.Constraint, 0, 1)
 	if isInverted {
-		ic.tight, constraints = ic.makeInvertedIndexSpansForExpr(
-			&ic.filters, constraints, false,
+		tight, constraints := ic.makeInvertedIndexSpansForExpr(
+			&ic.allFilters, nil /* constraints */, false,
 		)
 		// makeInvertedIndexSpansForExpr is guaranteed to add at least one
 		// constraint. It may be an "unconstrained" constraint.
+		ic.tight = tight
 		ic.constraint = *constraints[0]
 	} else {
-		ic.tight = ic.makeSpansForExpr(0 /* offset */, &ic.filters, &ic.constraint)
+		ic.tight = ic.makeSpansForExpr(0 /* offset */, &ic.allFilters, &ic.constraint)
 	}
 	// Note: we only consolidate spans at the end; consolidating partial results
 	// can lead to worse spans, for example:
@@ -1125,8 +1143,7 @@ func (ic *Instance) Constraint() *constraint.Constraint {
 // AllInvertedIndexConstraints returns all constraints that can be created on
 // the specified inverted index. Only works for inverted indexes.
 func (ic *Instance) AllInvertedIndexConstraints() ([]*constraint.Constraint, error) {
-	allConstraints := make([]*constraint.Constraint, 0)
-	_, allConstraints = ic.makeInvertedIndexSpansForExpr(&ic.filters, allConstraints, true)
+	_, allConstraints := ic.makeInvertedIndexSpansForExpr(&ic.allFilters, nil /* constraints */, true /* allPaths */)
 
 	return allConstraints, nil
 }
@@ -1139,15 +1156,15 @@ func (ic *Instance) RemainingFilters() memo.FiltersExpr {
 		return memo.TrueFilter
 	}
 	if ic.constraint.IsUnconstrained() {
-		return ic.filters
+		return ic.requiredFilters
 	}
 	var newFilters memo.FiltersExpr
-	for i := range ic.filters {
+	for i := range ic.requiredFilters {
 		prefix := ic.getMaxSimplifyPrefix(&ic.constraint)
-		condition := ic.simplifyFilter(ic.filters[i].Condition, &ic.constraint, prefix)
+		condition := ic.simplifyFilter(ic.requiredFilters[i].Condition, &ic.constraint, prefix)
 		if condition.Op() != opt.TrueOp {
 			if newFilters == nil {
-				newFilters = make(memo.FiltersExpr, 0, len(ic.filters))
+				newFilters = make(memo.FiltersExpr, 0, len(ic.requiredFilters))
 			}
 			newFilters = append(newFilters, memo.FiltersItem{Condition: condition})
 		}

pkg/sql/opt/idxconstraint/index_constraints_test.go

@@ -67,7 +67,6 @@ func TestIndexConstraints(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 
 	datadriven.Walk(t, "testdata", func(t *testing.T, path string) {
-		ctx := context.Background()
 		semaCtx := tree.MakeSemaContext()
 		evalCtx := tree.MakeTestingEvalContext(cluster.MakeTestingClusterSettings())
 
@@ -124,36 +123,30 @@ func TestIndexConstraints(t *testing.T) {
 
 			switch d.Cmd {
 			case "index-constraints":
-				typedExpr, err := testutils.ParseScalarExpr(d.Input, iVarHelper.Container())
-				if err != nil {
-					d.Fatalf(t, "%v", err)
-				}
-
-				if normalizeTypedExpr {
-					typedExpr, err = evalCtx.NormalizeExpr(typedExpr)
+				// Allow specifying optional filters using the "optional:" delimiter.
+				var filters, optionalFilters memo.FiltersExpr
+				if idx := strings.Index(d.Input, "optional:"); idx >= 0 {
+					optional := d.Input[idx+len("optional:"):]
+					optionalFilters, err = buildFilters(
+						optional, &semaCtx, &evalCtx, &f, &iVarHelper, normalizeTypedExpr,
+					)
 					if err != nil {
 						d.Fatalf(t, "%v", err)
 					}
+					d.Input = d.Input[:idx]
+				}
+				if filters, err = buildFilters(
+					d.Input, &semaCtx, &evalCtx, &f, &iVarHelper, normalizeTypedExpr,
+				); err != nil {
+					d.Fatalf(t, "%v", err)
 				}
 
 				varNames := make([]string, len(varTypes))
 				for i := range varNames {
 					varNames[i] = fmt.Sprintf("@%d", i+1)
 				}
-				b := optbuilder.NewScalar(ctx, &semaCtx, &evalCtx, &f)
-				b.AllowUnsupportedExpr = true
-				err = b.Build(typedExpr)
-				if err != nil {
-					return fmt.Sprintf("error: %v\n", err)
-				}
-				root := f.Memo().RootExpr().(opt.ScalarExpr)
-				filters := memo.FiltersExpr{{Condition: root}}
-				if _, ok := root.(*memo.TrueExpr); ok {
-					filters = memo.TrueFilter
-				}
-
 				var ic idxconstraint.Instance
-				ic.Init(filters, indexCols, notNullCols, invertedIndex, &evalCtx, &f)
+				ic.Init(filters, optionalFilters, indexCols, notNullCols, invertedIndex, &evalCtx, &f)
 				result := ic.Constraint()
 				var buf bytes.Buffer
 				for i := 0; i < result.Spans.Count(); i++ {
@@ -232,6 +225,7 @@ func BenchmarkIndexConstraints(b *testing.B) {
 		testCases = append(testCases, tc)
 	}
 
+	semaCtx := tree.MakeSemaContext()
 	evalCtx := tree.MakeTestingEvalContext(cluster.MakeTestingClusterSettings())
 
 	for _, tc := range testCases {
@@ -249,26 +243,15 @@ func BenchmarkIndexConstraints(b *testing.B) {
 			indexCols, notNullCols := parseIndexColumns(b, md, strings.Split(tc.indexInfo, ", "))
 
 			iVarHelper := tree.MakeTypesOnlyIndexedVarHelper(varTypes)
-			typedExpr, err := testutils.ParseScalarExpr(tc.expr, iVarHelper.Container())
+			filters, err := buildFilters(tc.expr, &semaCtx, &evalCtx, &f, &iVarHelper, false /* normalizeTypedExpr */)
 			if err != nil {
 				b.Fatal(err)
 			}
 
-			semaCtx := tree.MakeSemaContext()
-			evalCtx := tree.MakeTestingEvalContext(cluster.MakeTestingClusterSettings())
-			bld := optbuilder.NewScalar(context.Background(), &semaCtx, &evalCtx, &f)
-
-			err = bld.Build(typedExpr)
-			if err != nil {
-				b.Fatal(err)
-			}
-			nd := f.Memo().RootExpr()
-			filters := memo.FiltersExpr{{Condition: nd.(opt.ScalarExpr)}}
-
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
 				var ic idxconstraint.Instance
-				ic.Init(filters, indexCols, notNullCols, false /*isInverted */, &evalCtx, &f)
+				ic.Init(filters, nil /* optionalFilters */, indexCols, notNullCols, false /*isInverted */, &evalCtx, &f)
 				_ = ic.Constraint()
 				_ = ic.RemainingFilters()
 			}
@@ -316,3 +299,37 @@ func parseIndexColumns(
 	}
 	return columns, notNullCols
 }
+
+func buildFilters(
+	input string,
+	semaCtx *tree.SemaContext,
+	evalCtx *tree.EvalContext,
+	f *norm.Factory,
+	iVarHelper *tree.IndexedVarHelper,
+	normalizeTypedExpr bool,
+) (memo.FiltersExpr, error) {
+	if input == "" {
+		return memo.TrueFilter, nil
+	}
+	typedExpr, err := testutils.ParseScalarExpr(input, iVarHelper.Container())
+	if err != nil {
+		return memo.FiltersExpr{}, err
+	}
+	if normalizeTypedExpr {
+		typedExpr, err = evalCtx.NormalizeExpr(typedExpr)
+		if err != nil {
+			return memo.FiltersExpr{}, err
+		}
+	}
+
+	b := optbuilder.NewScalar(context.Background(), semaCtx, evalCtx, f)
+	b.AllowUnsupportedExpr = true
+	if err := b.Build(typedExpr); err != nil {
+		return memo.FiltersExpr{}, err
+	}
+	root := f.Memo().RootExpr().(opt.ScalarExpr)
+	if _, ok := root.(*memo.TrueExpr); ok {
+		return memo.TrueFilter, nil
+	}
+	return memo.FiltersExpr{{Condition: root}}, nil
+}

pkg/sql/opt/idxconstraint/testdata/optional

@@ -0,0 +1,16 @@
+index-constraints vars=(int, int) index=(@1)
+@1 > 2 AND @1 < 4 AND @2 = 2
+----
+[/3 - /3]
+Remaining filter: @2 = 2
+
+index-constraints vars=(int, int) index=(@1)
+@1 > 2
+optional: @1 < 4 AND @2 = 2
+----
+[/3 - /3]
+
+index-constraints vars=(int, int) index=(@1 desc, @2 desc)
+optional: @1 >= 2 AND @1 <= 4 AND @2 IN (1, 2, 3)
+----
+[/4/3 - /2/1]

pkg/sql/opt/memo/expr.go

@@ -234,21 +234,6 @@ func (n *FiltersExpr) Deduplicate() {
 	*n = dedup
 }
 
-// RetainCommonFilters retains only the filters found in n and other.
-func (n *FiltersExpr) RetainCommonFilters(other FiltersExpr) {
-	// TODO(ridwanmsharif): Faster intersection using a map
-	common := (*n)[:0]
-	for _, filter := range *n {
-		for _, otherFilter := range other {
-			if filter.Condition == otherFilter.Condition {
-				common = append(common, filter)
-				break
-			}
-		}
-	}
-	*n = common
-}
-
 // RemoveCommonFilters removes the filters found in other from n.
 func (n *FiltersExpr) RemoveCommonFilters(other FiltersExpr) {
 	// TODO(ridwanmsharif): Faster intersection using a map

pkg/sql/opt/memo/memo.go

@@ -223,10 +223,8 @@ func (m *Memo) SetRoot(e RelExpr, phys *physical.Required) {
 }
 
 // SetScalarRoot stores the root memo expression when it is a scalar expression.
+// Used only for testing.
 func (m *Memo) SetScalarRoot(scalar opt.ScalarExpr) {
-	if m.rootExpr != nil {
-		panic(errors.AssertionFailedf("cannot set scalar root multiple times"))
-	}
 	m.rootExpr = scalar
 }