opt: generate lookup joins on partitioned indexes

pkg/sql/opt/constraint/constraint_set.go

@@ -312,20 +312,26 @@ func (s *Set) ExtractValueForConstCol(evalCtx *tree.EvalContext, col opt.ColumnI
 	return nil
 }
 
-// IsSingleColumnConstValue returns true if the Set contains a single constraint
-// on a single column which allows for a single constant value. On success,
-// returns the column and the constant value.
-func (s *Set) IsSingleColumnConstValue(
+// HasSingleColumnConstValues returns true if the Set contains a single
+// constraint on a single column which allows for one or more non-ranging
+// constant values. On success, returns the column and the constant value.
+func (s *Set) HasSingleColumnConstValues(
 	evalCtx *tree.EvalContext,
-) (col opt.ColumnID, constValue tree.Datum, ok bool) {
+) (col opt.ColumnID, constValues tree.Datums, ok bool) {
 	if s.Length() != 1 {
 		return 0, nil, false
 	}
 	c := s.Constraint(0)
-	if c.Columns.Count() != 1 || c.ExactPrefix(evalCtx) != 1 {
+	if c.Columns.Count() != 1 || c.Prefix(evalCtx) != 1 {
 		return 0, nil, false
 	}
-	return c.Columns.Get(0).ID(), c.Spans.Get(0).StartKey().Value(0), true
+	numSpans := c.Spans.Count()
+	constValues = make(tree.Datums, numSpans)
+	for i := range constValues {
+		val := c.Spans.Get(i).StartKey().Value(0)
+		constValues[i] = val
+	}
+	return c.Columns.Get(0).ID(), constValues, true
 }
 
 // allocConstraint allocates space for a new constraint in the set and returns

pkg/sql/opt/constraint/constraint_set_test.go

@@ -11,6 +11,7 @@
 package constraint
 
 import (
+	"reflect"
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/sql/opt"
@@ -347,32 +348,34 @@ func TestExtractConstCols(t *testing.T) {
 	}
 }
 
-func TestIsSingleColumnConstValue(t *testing.T) {
+func TestHasSingleColumnConstValues(t *testing.T) {
 	type testCase struct {
 		constraints []string
 		col         opt.ColumnID
-		val         int
+		vals        []int
 	}
 	cases := []testCase{
-		{[]string{`/1: [/10 - /10]`}, 1, 10},
-		{[]string{`/-1: [/10 - /10]`}, 1, 10},
-		{[]string{`/1: [/10 - /11]`}, 0, 0},
-		{[]string{`/1: [/10 - /10] [/11 - /11]`}, 0, 0},
-		{[]string{`/1/2: [/10/2 - /10/4]`}, 0, 0},
-		{[]string{`/1/2: [/10/2 - /10/2]`}, 0, 0},
+		{[]string{`/1: [/10 - /10]`}, 1, []int{10}},
+		{[]string{`/-1: [/10 - /10]`}, 1, []int{10}},
+		{[]string{`/1: [/10 - /11]`}, 0, nil},
+		{[]string{`/1: [/10 - /10] [/11 - /11]`}, 1, []int{10, 11}},
+		{[]string{`/1: [/10 - /10] [/11 - /11] [/12 - /12]`}, 1, []int{10, 11, 12}},
+		{[]string{`/1: [/10 - /10] [/11 - /11] [/12 - /13]`}, 0, nil},
+		{[]string{`/1/2: [/10/2 - /10/4]`}, 0, nil},
+		{[]string{`/1/2: [/10/2 - /10/2]`}, 0, nil},
 		{
 			[]string{
 				`/1: [/10 - /10]`,
 				`/2: [/8 - /8]`,
 			},
-			0, 0,
+			0, nil,
 		},
 		{
 			[]string{
 				`/1: [/10 - /10]`,
 				`/1/2: [/10/8 - /10/8]`,
 			},
-			0, 0,
+			0, nil,
 		},
 	}
 	evalCtx := tree.NewTestingEvalContext(nil)
@@ -382,13 +385,13 @@ func TestIsSingleColumnConstValue(t *testing.T) {
 			constraint := ParseConstraint(evalCtx, constraint)
 			cs = cs.Intersect(evalCtx, SingleConstraint(&constraint))
 		}
-		col, val, ok := cs.IsSingleColumnConstValue(evalCtx)
-		intVal := 0
-		if ok {
-			intVal = int(*val.(*tree.DInt))
+		col, vals, _ := cs.HasSingleColumnConstValues(evalCtx)
+		var intVals []int
+		for _, val := range vals {
+			intVals = append(intVals, int(*val.(*tree.DInt)))
 		}
-		if tc.col != col || tc.val != intVal {
-			t.Errorf("%s: expected %d,%d got %d,%d", cs, tc.col, tc.val, col, intVal)
+		if tc.col != col || !reflect.DeepEqual(tc.vals, intVals) {
+			t.Errorf("%s: expected %d,%d got %d,%d", cs, tc.col, tc.vals, col, intVals)
 		}
 	}
 }

pkg/sql/opt/exec/execbuilder/testdata/lookup_join

@@ -1736,13 +1736,13 @@ vectorized: true
     │ estimated row count: 10 (missing stats)
     │
     └── • lookup join (semi)
-        │ columns: ("project_const_col_@15", pk, col0, col3)
+        │ columns: ("lookup_join_const_col_@15", pk, col0, col3)
         │ table: tab4@tab4_col3_col4_key
-        │ equality: (col0, project_const_col_@15) = (col3,col4)
+        │ equality: (col0, lookup_join_const_col_@15) = (col3,col4)
         │ equality cols are key
         │
         └── • render
-            │ columns: ("project_const_col_@15", pk, col0, col3)
+            │ columns: ("lookup_join_const_col_@15", pk, col0, col3)
             │ estimated row count: 10 (missing stats)
             │ render 0: 495.6
             │ render 1: pk

pkg/sql/opt/memo/testdata/stats/lookup-join

@@ -406,14 +406,14 @@ inner-join (lookup wxyz)
  │    ├── stats: [rows=19.8, distinct(1)=19.8, null(1)=0, distinct(6)=19.8, null(6)=0, distinct(7)=1, null(7)=0, distinct(11)=1, null(11)=0]
  │    ├── fd: ()-->(7), (9)-->(6,8), (1)==(6), (6)==(1)
  │    ├── project
- │    │    ├── columns: "project_const_col_@7":11(int!null) m:1(int) n:2(int)
+ │    │    ├── columns: "lookup_join_const_col_@7":11(int!null) m:1(int) n:2(int)
  │    │    ├── stats: [rows=40, distinct(1)=40, null(1)=0, distinct(11)=1, null(11)=0]
  │    │    ├── fd: ()-->(11)
  │    │    ├── scan medium
  │    │    │    ├── columns: m:1(int) n:2(int)
  │    │    │    └── stats: [rows=40, distinct(1)=40, null(1)=0]
  │    │    └── projections
- │    │         └── 10 [as="project_const_col_@7":11, type=int]
+ │    │         └── 10 [as="lookup_join_const_col_@7":11, type=int]
  │    └── filters (true)
  └── filters (true)
 

pkg/sql/opt/xform/general_funcs.go

@@ -154,3 +154,77 @@ func (c *CustomFuncs) initIdxConstraintForIndex(
 	)
 	return ic
 }
+
+// computedColFilters generates all filters that can be derived from the list of
+// computed column expressions from the given table. A computed column can be
+// used as a filter when it has a constant value. That is true when:
+//
+//   1. All other columns it references are constant, because other filters in
+//      the query constrain them to be so.
+//   2. All functions in the computed column expression can be folded into
+//      constants (i.e. they do not have problematic side effects).
+//
+// Note that computed columns can depend on other computed columns; in general
+// the dependencies form an acyclic directed graph. computedColFilters will
+// return filters for all constant computed columns, regardless of the order of
+// their dependencies.
+//
+// As with checkConstraintFilters, computedColFilters do not really filter any
+// rows, they are rather facts or guarantees about the data. Treating them as
+// filters may allow some indexes to be constrained and used. Consider the
+// following example:
+//
+//   CREATE TABLE t (
+//     k INT NOT NULL,
+//     hash INT AS (k % 4) STORED,
+//     PRIMARY KEY (hash, k)
+//   )
+//
+//   SELECT * FROM t WHERE k = 5
+//
+// Notice that the filter provided explicitly wouldn't allow the optimizer to
+// seek using the primary index (it would have to fall back to a table scan).
+// However, column "hash" can be proven to have the constant value of 1, since
+// it's dependent on column "k", which has the constant value of 5. This enables
+// usage of the primary index:
+//
+//     scan t
+//      ├── columns: k:1(int!null) hash:2(int!null)
+//      ├── constraint: /2/1: [/1/5 - /1/5]
+//      ├── key: (2)
+//      └── fd: ()-->(1)
+//
+// The values of both columns in that index are known, enabling a single value
+// constraint to be generated.
+func (c *CustomFuncs) computedColFilters(
+	tabID opt.TableID, requiredFilters, optionalFilters memo.FiltersExpr,
+) memo.FiltersExpr {
+	tabMeta := c.e.mem.Metadata().TableMeta(tabID)
+	if len(tabMeta.ComputedCols) == 0 {
+		return nil
+	}
+
+	// Start with set of constant columns, as derived from the list of filter
+	// conditions.
+	constCols := make(map[opt.ColumnID]opt.ScalarExpr)
+	c.findConstantFilterCols(constCols, tabID, requiredFilters)
+	c.findConstantFilterCols(constCols, tabID, optionalFilters)
+	if len(constCols) == 0 {
+		// No constant values could be derived from filters, so assume that there
+		// are also no constant computed columns.
+		return nil
+	}
+
+	// Construct a new filter condition for each computed column that is
+	// constant (i.e. all of its variables are in the constCols set).
+	var computedColFilters memo.FiltersExpr
+	for colID := range tabMeta.ComputedCols {
+		if c.tryFoldComputedCol(tabMeta, colID, constCols) {
+			constVal := constCols[colID]
+			// Note: Eq is not correct here because of NULLs.
+			eqOp := c.e.f.ConstructIs(c.e.f.ConstructVariable(colID), constVal)
+			computedColFilters = append(computedColFilters, c.e.f.ConstructFiltersItem(eqOp))
+		}
+	}
+	return computedColFilters
+}