diff --git a/pkg/sql/opt/build.go b/pkg/sql/opt/build.go index 63b7daa16f05..bbd16ed544dd 100644 --- a/pkg/sql/opt/build.go +++ b/pkg/sql/opt/build.go @@ -153,6 +153,20 @@ func buildScalar(buildCtx *buildContext, pexpr tree.TypedExpr) *expr { case *tree.IndexedVar: initVariableExpr(e, t.Idx) + case *tree.Tuple: + children := make([]*expr, len(t.Exprs)) + for i, e := range t.Exprs { + children[i] = buildScalar(buildCtx, e.(tree.TypedExpr)) + } + initTupleExpr(e, children) + + case *tree.DTuple: + children := make([]*expr, len(t.D)) + for i, d := range t.D { + children[i] = buildScalar(buildCtx, d) + } + initTupleExpr(e, children) + case tree.Datum: initConstExpr(e, t) diff --git a/pkg/sql/opt/index_constraints.go b/pkg/sql/opt/index_constraints.go index 3d23e7da80a5..0a3dce005a7c 100644 --- a/pkg/sql/opt/index_constraints.go +++ b/pkg/sql/opt/index_constraints.go @@ -359,20 +359,32 @@ func (c *indexConstraintCalc) intersectSpanSet(a *LogicalSpan, spanSet LogicalSp return res } -// makeSpansForExpr creates spans for index columns starting at -// from the given expression. -func (c *indexConstraintCalc) makeSpansForExpr(depth int, e *expr) (LogicalSpans, bool) { - // Check if the operator is supported. - switch e.op { - case eqOp, ltOp, gtOp, leOp, geOp, neOp: - // We support comparisons when the left-hand side is an indexed var for the - // current column and the right-hand side is a constant. - if !isIndexedVar(e.children[0], depth) || e.children[1].op != constOp { - return nil, false +// makeSpansForSingleColumn creates spans for a single index column from a +// simple comparison expression. The arguments are the operator and right +// operand. +func (c *indexConstraintCalc) makeSpansForSingleColumn( + op operator, val *expr, +) (LogicalSpans, bool) { + if op == inOp && isTupleOfConstants(val) { + // We assume that the values of the tuple are already ordered and distinct. + spans := make(LogicalSpans, len(val.children)) + for i, v := range val.children { + datum := v.private.(tree.Datum) + spans[i].Start = LogicalKey{Vals: tree.Datums{datum}, Inclusive: true} + spans[i].End = spans[i].Start } - datum := e.children[1].private.(tree.Datum) + return spans, true + } + // The rest of the supported expressions must have a constant scalar on the + // right-hand side. + if val.op != constOp { + return nil, false + } + datum := val.private.(tree.Datum) + switch op { + case eqOp, ltOp, gtOp, leOp, geOp: sp := MakeFullSpan() - switch e.op { + switch op { case eqOp: sp.Start = LogicalKey{Vals: tree.Datums{datum}, Inclusive: true} sp.End = LogicalKey{Vals: tree.Datums{datum}, Inclusive: true} @@ -386,19 +398,31 @@ func (c *indexConstraintCalc) makeSpansForExpr(depth int, e *expr) (LogicalSpans sp.End = LogicalKey{Vals: tree.Datums{datum}, Inclusive: true} case geOp: sp.Start = LogicalKey{Vals: tree.Datums{datum}, Inclusive: true} - case neOp: - sp.End = LogicalKey{Vals: tree.Datums{datum}, Inclusive: false} - sp2 := MakeFullSpan() - sp2.Start = LogicalKey{Vals: tree.Datums{datum}, Inclusive: false} - return LogicalSpans{sp, sp2}, true } return LogicalSpans{sp}, true + case neOp: + spans := LogicalSpans{MakeFullSpan(), MakeFullSpan()} + spans[0].End = LogicalKey{Vals: tree.Datums{datum}, Inclusive: false} + spans[1].Start = LogicalKey{Vals: tree.Datums{datum}, Inclusive: false} + return spans, true + default: return nil, false } } +// makeSpansForExpr creates spans for index columns starting at +// from the given expression. +func (c *indexConstraintCalc) makeSpansForExpr(depth int, e *expr) (LogicalSpans, bool) { + // Check for an operation where the left-hand side is an + // indexed var for this column. + if isIndexedVar(e.children[0], depth) { + return c.makeSpansForSingleColumn(e.op, e.children[1]) + } + return nil, false +} + // calcDepth calculates constraints for the sequence of index columns starting // at (i.e. the -th index column, -th index column, etc.). // diff --git a/pkg/sql/opt/opt_test.go b/pkg/sql/opt/opt_test.go index 5c41c9a62b16..332c67cb0995 100644 --- a/pkg/sql/opt/opt_test.go +++ b/pkg/sql/opt/opt_test.go @@ -18,19 +18,35 @@ package opt // is used for optimizer-specific testcases. // // Each testfile contains testcases of the form -// +// [,...] [ ...] // // ---- // // // The supported commands are: // -// - build-scalar [ ...] +// - build-scalar // // Builds an expression tree from a scalar SQL expression and outputs a // representation of the tree. The expression can refer to external variables // using @1, @2, etc. in which case the types of the variables must be passed // on the command line. +// +// - legacy-normalize +// +// Runs the TypedExpr normalization code and rebuilds the scalar expression. +// If present, must follow build-scalar. +// +// - normalize +// +// Normalizes the expression. If present, must follow build-scalar or +// legacy-normalize. +// +// - index-constraints +// +// Creates index constraints on the assumption that the index is formed by +// the index var columns (as specified by ). +// If present, build-scalar must have been an earlier command. import ( "bufio" @@ -229,7 +245,18 @@ func TestOpt(t *testing.T) { runTest(t, path, func(d *testdata) string { var e *expr var types []types.T + var typedExpr tree.TypedExpr + buildScalarFn := func() { + defer func() { + if r := recover(); r != nil { + d.fatalf(t, "buildScalar: %v", r) + } + }() + e = buildScalar(&buildContext{}, typedExpr) + } + + evalCtx := tree.MakeTestingEvalContext() for _, cmd := range strings.Split(d.cmd, ",") { switch cmd { case "build-scalar": @@ -238,19 +265,19 @@ func TestOpt(t *testing.T) { if err != nil { d.fatalf(t, "%v", err) } - typedExpr, err := parseScalarExpr(d.sql, types) + typedExpr, err = parseScalarExpr(d.sql, types) if err != nil { d.fatalf(t, "%v", err) } - e = func() *expr { - defer func() { - if r := recover(); r != nil { - d.fatalf(t, "buildScalar: %v", r) - } - }() - return buildScalar(&buildContext{}, typedExpr) - }() + buildScalarFn() + case "legacy-normalize": + // Apply the TypedExpr normalization and rebuild the expression. + typedExpr, err = evalCtx.NormalizeExpr(typedExpr) + if err != nil { + d.fatalf(t, "%v", err) + } + buildScalarFn() case "normalize": normalizeScalar(e) case "index-constraints": @@ -258,7 +285,6 @@ func TestOpt(t *testing.T) { d.fatalf(t, "no expression for index-constraints") } - evalCtx := tree.MakeTestingEvalContext() spans := MakeIndexConstraints(e, types, &evalCtx) var buf bytes.Buffer for _, sp := range spans { diff --git a/pkg/sql/opt/scalar.go b/pkg/sql/opt/scalar.go index 287af8137656..3a08bbac6e26 100644 --- a/pkg/sql/opt/scalar.go +++ b/pkg/sql/opt/scalar.go @@ -144,6 +144,27 @@ func isIndexedVar(e *expr, index int) bool { return e.op == variableOp && e.private.(int) == index } +func initTupleExpr(e *expr, children []*expr) { + // In general, the order matters in a tuple so we use an "ordered list" + // operator. In some cases (IN) the order doesn't matter; we could convert + // those to listOp during normalization, but there doesn't seem to be a + // benefit at this time. + e.op = orderedListOp + e.children = children +} + +func isTupleOfConstants(e *expr) bool { + if e.op != orderedListOp { + return false + } + for _, c := range e.children { + if c.op != constOp { + return false + } + } + return true +} + // Applies a set of normalization rules to a scalar expression. // // For now, we expect to build exprs from TypedExprs which have gone through a diff --git a/pkg/sql/opt/testdata/build-scalar b/pkg/sql/opt/testdata/build-scalar index 3ac62531e61d..76c9bd6b6abb 100644 --- a/pkg/sql/opt/testdata/build-scalar +++ b/pkg/sql/opt/testdata/build-scalar @@ -131,3 +131,38 @@ or (type: bool) └── gt (type: bool) ├── variable (2) (type: int) └── const (2) (type: int) + +build-scalar int int +(@1, @2) = (1, 2) +---- +eq (type: bool) + ├── ordered-list (type: tuple{int, int}) + │ ├── variable (0) (type: int) + │ └── variable (1) (type: int) + └── ordered-list (type: tuple{int, int}) + ├── const (1) (type: int) + └── const (2) (type: int) + +build-scalar int +@1 IN (1, 2) +---- +in (type: bool) + ├── variable (0) (type: int) + └── ordered-list (type: tuple{int, int}) + ├── const (1) (type: int) + └── const (2) (type: int) + +build-scalar int int +(@1, @2) IN ((1, 2), (3, 4)) +---- +in (type: bool) + ├── ordered-list (type: tuple{int, int}) + │ ├── variable (0) (type: int) + │ └── variable (1) (type: int) + └── ordered-list (type: tuple{tuple{int, int}, tuple{int, int}}) + ├── ordered-list (type: tuple{int, int}) + │ ├── const (1) (type: int) + │ └── const (2) (type: int) + └── ordered-list (type: tuple{int, int}) + ├── const (3) (type: int) + └── const (4) (type: int) diff --git a/pkg/sql/opt/testdata/index-constraints b/pkg/sql/opt/testdata/index-constraints index b4bd2331f0df..3be2523155e1 100644 --- a/pkg/sql/opt/testdata/index-constraints +++ b/pkg/sql/opt/testdata/index-constraints @@ -113,3 +113,53 @@ build-scalar,normalize,index-constraints decimal decimal @1 > 1.5 AND @2 > 2 ---- (/1.5 - ] + +build-scalar,normalize,index-constraints int +@1 IN (1, 2, 3) +---- +[/1 - /1] +[/2 - /2] +[/3 - /3] + +build-scalar,legacy-normalize,index-constraints int +@1 IN (1, 5, 1, 4) +---- +[/1 - /1] +[/4 - /4] +[/5 - /5] + +build-scalar,normalize,index-constraints int int +@1 = 1 AND @2 IN (1, 2, 3) +---- +[/1/1 - /1/1] +[/1/2 - /1/2] +[/1/3 - /1/3] + +build-scalar,normalize,index-constraints int int +@1 IN (1, 2) AND @2 IN (1, 2, 3) +---- +[/1/1 - /1/1] +[/1/2 - /1/2] +[/1/3 - /1/3] +[/2/1 - /2/1] +[/2/2 - /2/2] +[/2/3 - /2/3] + +build-scalar,normalize,index-constraints int int +@1 >= 2 AND @1 <= 4 AND @2 IN (1, 2, 3) +---- +[/2/1 - /4/3] + +build-scalar,normalize,index-constraints int int +@1 IN (1, 2, 3) AND @2 = 4 +---- +[/1/4 - /1/4] +[/2/4 - /2/4] +[/3/4 - /3/4] + +build-scalar,normalize,index-constraints int int +@1 IN (1, 2, 3) AND @2 >= 2 AND @2 <= 4 +---- +[/1/2 - /1/4] +[/2/2 - /2/4] +[/3/2 - /3/4]