opt: add GenerateUnionSelects exploration rule for disjunction

This commit adds a new exploration rule that can produce better query
plans for disjunctions (e.g. a = 1 OR b = 2). The rule transforms some
Select + Scan expressions with a disjunction filter into a Union of two
Select expressions, each with one side of the disjuction as a filter.
This can result in faster query plans in cases where two indexes cover
each side of the disjunction.

This rule only applies for Scan expressions that contain a strict key.

Fixes cockroachdb#2142

Release note (performance improvement): The query optimizer now produces
faster query plans for some disjunctions (OR expressions) by utilizing
multiple indexes.

pkg/sql/opt/xform/custom_funcs.go

@@ -2268,6 +2268,89 @@ func (c *CustomFuncs) MakeOrderingChoiceFromColumn(
 	return oc
 }
 
+// DuplicateScan constructs a new ScanPrivate that is identical to the input,
+// but has new table and column IDs.
+//
+// DuplicateScan assumes it is being called on a canonical scan. Non-canonical
+// scan properties like constraints are not copied to the new ScanPrivate.
+func (c *CustomFuncs) DuplicateScan(sp *memo.ScanPrivate) *memo.ScanPrivate {
+	md := c.e.mem.Metadata()
+	tabMD := md.TableMeta(sp.Table)
+	dupTabID := md.AddTable(tabMD.Table, &tabMD.Alias)
+
+	var dupTabColIDs opt.ColSet
+	cols := sp.Cols
+	for i, ok := cols.Next(0); ok; i, ok = cols.Next(i + 1) {
+		ord := tabMD.MetaID.ColumnOrdinal(i)
+		dupColID := dupTabID.ColumnID(ord)
+		dupTabColIDs.Add(dupColID)
+	}
+
+	return &memo.ScanPrivate{
+		Table: dupTabID,
+		Cols:  dupTabColIDs,
+	}
+}
+
+// MapFilterCols returns a new FiltersExpr with all the src column IDs in the input
+// expression replaced with column IDs in dst.
+func (c *CustomFuncs) MapFilterCols(
+	filters memo.FiltersExpr, src, dst *memo.ScanPrivate,
+) memo.FiltersExpr {
+	srcCols := opt.ColSetToList(src.Cols)
+	dstCols := opt.ColSetToList(dst.Cols)
+
+	// Map each column in src to a column in dst.
+	var colMap util.FastIntMap
+	for i, srcCol := range srcCols {
+		colMap.Set(int(srcCol), int(dstCols[i]))
+	}
+
+	// Map the columns of each filter in the FilterExpr.
+	newFilters := make([]memo.FiltersItem, 0, len(filters))
+	for _, f := range filters {
+		expr := c.mapFiltersItemCols(f, colMap)
+		newFilters = append(newFilters, c.e.f.ConstructFiltersItem(expr))
+	}
+
+	return newFilters
+}
+
+// mapFiltersItemCols replaces the column ID of each VariableExpr in the input
+// expression with new column IDs. The function resursively traverses the
+// children of the expression tree looking for columns that need to be replaced.
+// Each occurrence of the keys in colMap will be replaced with their
+// corresponding values in colMap.
+func (c *CustomFuncs) mapFiltersItemCols(
+	filter memo.FiltersItem, colMap util.FastIntMap,
+) opt.ScalarExpr {
+	var replace norm.ReplaceFunc
+	replace = func(nd opt.Expr) opt.Expr {
+		switch t := nd.(type) {
+		case *memo.VariableExpr:
+			dstCol, ok := colMap.Get(int(t.Col))
+			if !ok {
+				// TODO: Is this case even possible?
+				return nd
+			}
+			return c.e.f.ConstructVariable(opt.ColumnID(dstCol))
+		}
+		return c.e.f.Replace(nd, replace)
+	}
+
+	return replace(filter.Condition).(opt.ScalarExpr)
+}
+
+// ConstructSetPrivate constructs a new SetPrivate with columnn sets from the
+// left, right, and output of the operation.
+func (c *CustomFuncs) ConstructSetPrivate(left, right, out *memo.ScanPrivate) *memo.SetPrivate {
+	return &memo.SetPrivate{
+		LeftCols:  opt.ColSetToList(left.Cols),
+		RightCols: opt.ColSetToList(right.Cols),
+		OutCols:   opt.ColSetToList(out.Cols),
+	}
+}
+
 // scanIndexIter is a helper struct that supports iteration over the indexes
 // of a Scan operator table. For example:
 //

pkg/sql/opt/xform/rules/select.opt

@@ -12,8 +12,8 @@
 # and examples.
 [GenerateConstrainedScans, Explore]
 (Select
-  (Scan $scanPrivate:* & (IsCanonicalScan $scanPrivate))
-  $filters:*
+    (Scan $scanPrivate:* & (IsCanonicalScan $scanPrivate))
+    $filters:*
 )
 =>
 (GenerateConstrainedScans $scanPrivate $filters)
@@ -22,8 +22,45 @@
 # be serviced by an inverted index.
 [GenerateInvertedIndexScans, Explore]
 (Select
-  (Scan $scanPrivate:* & (IsCanonicalScan $scanPrivate) & (HasInvertedIndexes $scanPrivate))
-  $filters:*
+    (Scan $scanPrivate:* & (IsCanonicalScan $scanPrivate) & (HasInvertedIndexes $scanPrivate))
+    $filters:*
 )
 =>
 (GenerateInvertedIndexScans $scanPrivate $filters)
+
+# GenerateUnionSelects splits disjunctions (Or expressions) into a Union of two
+# Select expressions, the first containing the left sub-expression of the Or
+# expression and the second containing the right sub-expression. All other
+# filter items in the original expression are preserved in the new Select
+# expressions.
+#
+# This can produce better query plans in cases where indexes cover both sides of
+# the Or expression. The execution plan can use both indexes to satisfy both
+# sides of the disjunction and union the results together.
+#
+# Note that this rule only matches Selects with canonical scans. Therefore scan
+# constraints do not need to be duplicated in the left and right scans of the
+# union.
+#
+# Also note that this rule only matches Selects that have strict keys. Without
+# a strict key, the Union expression would deduplicate rows with equal values
+# but differing keys, leading to incorrect results.
+[GenerateUnionSelects, Explore]
+(Select
+    $input:(Scan
+        $scanPrivate:* & (IsCanonicalScan $scanPrivate)
+    ) & (HasStrictKey $input)
+    $filters:[ ... $item:(FiltersItem (Or $left:* $right:*)) ... ]
+)
+=>
+(Union
+    (Select
+        (Scan $leftScan:(DuplicateScan $scanPrivate))
+        (MapFilterCols (ReplaceFiltersItem $filters $item $left) $scanPrivate $leftScan)
+    )
+    (Select
+        (Scan $rightScan:(DuplicateScan $scanPrivate))
+        (MapFilterCols (ReplaceFiltersItem $filters $item $right) $scanPrivate $rightScan)
+    )
+    (ConstructSetPrivate $leftScan $rightScan $scanPrivate)
+)