opt: explore splitting disjunction over same column

pkg/sql/opt/xform/select_funcs.go

@@ -1941,7 +1941,7 @@ func (c *CustomFuncs) SplitDisjunction(
 // An "interesting" pair of expressions is one where:
 //
 //  1. The column sets of both expressions in the pair are not
-//     equal.
+//     equal, and
 //  2. Two index scans can potentially be constrained by both expressions in
 //     the pair.
 //
@@ -1960,6 +1960,13 @@ func (c *CustomFuncs) SplitDisjunction(
 // There is no possible "interesting" pair here because the left and right sides
 // of the disjunction share the same columns.
 //
+// There is one exceptional case when a pair could be interesting even with
+// equal column sets for both expressions: when the table itself contains
+// multiple partial indexes with different predicates referencing the same
+// column. In this case we might be able to use different partial indexes for
+// both expressions, and so consider a pair interesting even with equal column
+// sets.
+//
 // findInterestingDisjunctionPair groups all sub-expressions adjacent to the
 // input's top-level OrExpr into left and right expression groups. These two
 // groups form the new filter expressions on the left and right side of the
@@ -1997,11 +2004,14 @@ func (c *CustomFuncs) findInterestingDisjunctionPair(
 		// not match) on.
 		if leftColSet.Empty() {
 			leftColSet = cols
+			leftExprs = append(leftExprs, expr)
+			return
 		}
 
-		// If the current expression ColSet matches leftColSet, add the expr to
-		// the left group. Otherwise, add it to the right group.
-		if leftColSet.Equals(cols) {
+		// If the current expression ColSet matches leftColSet (and we're not using
+		// the exception for multiple referencing partial index predicates) add the
+		// expr to the left group. Otherwise, add it to the right group.
+		if leftColSet.Equals(cols) && !c.multiplePartialIndexesReferencing(sp, leftColSet) {
 			leftExprs = append(leftExprs, expr)
 		} else {
 			rightColSet.UnionWith(cols)
@@ -2095,6 +2105,56 @@ func (c *CustomFuncs) canMaybeConstrainIndexWithCols(
 	return false
 }
 
+// multiplePartialIndexesReferencing returns true if at least one of the columns
+// is referenced by the predicates of multiple partial indexes. For example,
+// given this table:
+//
+//	CREATE TABLE abc (
+//	  a INT NOT NULL,
+//	  b INT NOT NULL,
+//	  c INT NOT NULL,
+//	  INDEX (a) WHERE b > 10,
+//	  INDEX (a) WHERE b != 100 AND c < 1000,
+//	  INDEX (c) WHERE a > 5 AND a % 2 = 0
+//	)
+//
+// Then multiplePartialIndexesReferencing will return true if called with (b) or
+// (a, b) or (b, c) or (a, b, c) but will return false if called with (a) or (c)
+// or (a, c).
+func (c *CustomFuncs) multiplePartialIndexesReferencing(
+	scanPrivate *memo.ScanPrivate, cols opt.ColSet,
+) bool {
+	md := c.e.mem.Metadata()
+	tabMeta := md.TableMeta(scanPrivate.Table)
+
+	var prevPartialIndexPredCols opt.ColSet
+
+	// Iterate through all partial indexes of the table and return true if one of
+	// the columns is referenced again after being referenced by a previous
+	// partial index.
+	for i := 0; i < tabMeta.Table.IndexCount(); i++ {
+		index := tabMeta.Table.Index(i)
+		if _, isPartialIndex := index.Predicate(); isPartialIndex {
+			p, ok := tabMeta.PartialIndexPredicate(i)
+			if !ok {
+				// A partial index predicate expression was not built for the
+				// partial index. See Builder.buildScan for details on when this
+				// can occur.
+				continue
+			}
+			pred := *p.(*memo.FiltersExpr)
+			partialIndexPredCols := pred.OuterCols().Intersection(cols)
+			// If one of the columns has now been referenced a second time, return
+			// true.
+			if partialIndexPredCols.Intersects(prevPartialIndexPredCols) {
+				return true
+			}
+			prevPartialIndexPredCols.UnionWith(partialIndexPredCols)
+		}
+	}
+	return false
+}
+
 // MakeSetPrivate constructs a new SetPrivate with given left, right, and out
 // columns.
 func (c *CustomFuncs) MakeSetPrivate(left, right, out opt.ColSet) *memo.SetPrivate {

pkg/sql/opt/xform/testdata/rules/select

@@ -12165,6 +12165,66 @@ memo (optimized, ~38KB, required=[presentation: a:1] [ordering: +2])
  ├── G51: (filters G39)
  └── G52: (scalar-list G48)
 
+# Regression test for #157073: explore splitting a disjunction on the same
+# column if there are multiple partial indexes with predicates on that column
+# that could be used.
+
+exec-ddl
+CREATE TABLE t157073 (
+  a INT NOT NULL PRIMARY KEY,
+  b INT NOT NULL,
+  c INT NOT NULL,
+  INDEX (c) WHERE b < 0,
+  INDEX (c) WHERE b >= 9990
+)
+----
+
+opt expect=SplitDisjunction
+SELECT c, a FROM t157073 WHERE c = 5 AND (b < 0 OR b >= 9990)
+----
+project
+ ├── columns: c:3!null a:1!null
+ ├── key: (1)
+ ├── fd: ()-->(3)
+ └── distinct-on
+      ├── columns: a:1!null b:2!null c:3!null
+      ├── grouping columns: a:1!null
+      ├── internal-ordering: +1
+      ├── key: (1)
+      ├── fd: ()-->(3), (1)-->(2)
+      ├── union-all
+      │    ├── columns: a:1!null b:2!null c:3!null
+      │    ├── left columns: a:6 b:7 c:8
+      │    ├── right columns: a:11 b:12 c:13
+      │    ├── ordering: +1
+      │    ├── index-join t157073
+      │    │    ├── columns: a:6!null b:7!null c:8!null
+      │    │    ├── key: (6)
+      │    │    ├── fd: ()-->(8), (6)-->(7)
+      │    │    ├── ordering: +6 opt(8) [actual: +6]
+      │    │    └── scan t157073@t157073_c_idx,partial
+      │    │         ├── columns: a:6!null c:8!null
+      │    │         ├── constraint: /8/6: [/5 - /5]
+      │    │         ├── key: (6)
+      │    │         ├── fd: ()-->(8)
+      │    │         └── ordering: +6 opt(8) [actual: +6]
+      │    └── index-join t157073
+      │         ├── columns: a:11!null b:12!null c:13!null
+      │         ├── key: (11)
+      │         ├── fd: ()-->(13), (11)-->(12)
+      │         ├── ordering: +11 opt(13) [actual: +11]
+      │         └── scan t157073@t157073_c_idx1,partial
+      │              ├── columns: a:11!null c:13!null
+      │              ├── constraint: /13/11: [/5 - /5]
+      │              ├── key: (11)
+      │              ├── fd: ()-->(13)
+      │              └── ordering: +11 opt(13) [actual: +11]
+      └── aggregations
+           ├── const-agg [as=b:2, outer=(2)]
+           │    └── b:2
+           └── const-agg [as=c:3, outer=(3)]
+                └── c:3
+
 # --------------------------------------------------
 # SplitDisjunctionAddKey
 # --------------------------------------------------