Skip to content

Commit d155065

Browse files
RaduBerindeRaduBerinde
authored
Merge pull request #20708 from RaduBerinde/index-constraints-tuples
opt: support tuples, make index constraints for IN
2 parents 7bed6f5 + 3efb5cd commit d155065

File tree

6 files changed

+199
-29
lines changed

6 files changed

+199
-29
lines changed

‎pkg/sql/opt/build.go‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,20 @@ func buildScalar(buildCtx *buildContext, pexpr tree.TypedExpr) *expr {
153153
case *tree.IndexedVar:
154154
initVariableExpr(e, t.Idx)
155155

156+
case *tree.Tuple:
157+
children := make([]*expr, len(t.Exprs))
158+
for i, e := range t.Exprs {
159+
children[i] = buildScalar(buildCtx, e.(tree.TypedExpr))
160+
}
161+
initTupleExpr(e, children)
162+
163+
case *tree.DTuple:
164+
children := make([]*expr, len(t.D))
165+
for i, d := range t.D {
166+
children[i] = buildScalar(buildCtx, d)
167+
}
168+
initTupleExpr(e, children)
169+
156170
case tree.Datum:
157171
initConstExpr(e, t)
158172

‎pkg/sql/opt/index_constraints.go‎

Lines changed: 41 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -359,20 +359,32 @@ func (c *indexConstraintCalc) intersectSpanSet(a *LogicalSpan, spanSet LogicalSp
359359
return res
360360
}
361361

362-
// makeSpansForExpr creates spans for index columns starting at <depth>
363-
// from the given expression.
364-
func (c *indexConstraintCalc) makeSpansForExpr(depth int, e *expr) (LogicalSpans, bool) {
365-
// Check if the operator is supported.
366-
switch e.op {
367-
case eqOp, ltOp, gtOp, leOp, geOp, neOp:
368-
// We support comparisons when the left-hand side is an indexed var for the
369-
// current column and the right-hand side is a constant.
370-
if !isIndexedVar(e.children[0], depth) || e.children[1].op != constOp {
371-
return nil, false
362+
// makeSpansForSingleColumn creates spans for a single index column from a
363+
// simple comparison expression. The arguments are the operator and right
364+
// operand.
365+
func (c *indexConstraintCalc) makeSpansForSingleColumn(
366+
op operator, val *expr,
367+
) (LogicalSpans, bool) {
368+
if op == inOp && isTupleOfConstants(val) {
369+
// We assume that the values of the tuple are already ordered and distinct.
370+
spans := make(LogicalSpans, len(val.children))
371+
for i, v := range val.children {
372+
datum := v.private.(tree.Datum)
373+
spans[i].Start = LogicalKey{Vals: tree.Datums{datum}, Inclusive: true}
374+
spans[i].End = spans[i].Start
372375
}
373-
datum := e.children[1].private.(tree.Datum)
376+
return spans, true
377+
}
378+
// The rest of the supported expressions must have a constant scalar on the
379+
// right-hand side.
380+
if val.op != constOp {
381+
return nil, false
382+
}
383+
datum := val.private.(tree.Datum)
384+
switch op {
385+
case eqOp, ltOp, gtOp, leOp, geOp:
374386
sp := MakeFullSpan()
375-
switch e.op {
387+
switch op {
376388
case eqOp:
377389
sp.Start = LogicalKey{Vals: tree.Datums{datum}, Inclusive: true}
378390
sp.End = LogicalKey{Vals: tree.Datums{datum}, Inclusive: true}
@@ -386,19 +398,31 @@ func (c *indexConstraintCalc) makeSpansForExpr(depth int, e *expr) (LogicalSpans
386398
sp.End = LogicalKey{Vals: tree.Datums{datum}, Inclusive: true}
387399
case geOp:
388400
sp.Start = LogicalKey{Vals: tree.Datums{datum}, Inclusive: true}
389-
case neOp:
390-
sp.End = LogicalKey{Vals: tree.Datums{datum}, Inclusive: false}
391-
sp2 := MakeFullSpan()
392-
sp2.Start = LogicalKey{Vals: tree.Datums{datum}, Inclusive: false}
393-
return LogicalSpans{sp, sp2}, true
394401
}
395402
return LogicalSpans{sp}, true
396403

404+
case neOp:
405+
spans := LogicalSpans{MakeFullSpan(), MakeFullSpan()}
406+
spans[0].End = LogicalKey{Vals: tree.Datums{datum}, Inclusive: false}
407+
spans[1].Start = LogicalKey{Vals: tree.Datums{datum}, Inclusive: false}
408+
return spans, true
409+
397410
default:
398411
return nil, false
399412
}
400413
}
401414

415+
// makeSpansForExpr creates spans for index columns starting at <depth>
416+
// from the given expression.
417+
func (c *indexConstraintCalc) makeSpansForExpr(depth int, e *expr) (LogicalSpans, bool) {
418+
// Check for an operation where the left-hand side is an
419+
// indexed var for this column.
420+
if isIndexedVar(e.children[0], depth) {
421+
return c.makeSpansForSingleColumn(e.op, e.children[1])
422+
}
423+
return nil, false
424+
}
425+
402426
// calcDepth calculates constraints for the sequence of index columns starting
403427
// at <depth> (i.e. the <depth>-th index column, <depth+1>-th index column, etc.).
404428
//

‎pkg/sql/opt/opt_test.go‎

Lines changed: 38 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -18,19 +18,35 @@ package opt
1818
// is used for optimizer-specific testcases.
1919
//
2020
// Each testfile contains testcases of the form
21-
// <command>
21+
// <command>[,<command>...] [<index-var-types> ...]
2222
// <SQL statement or expression>
2323
// ----
2424
// <expected results>
2525
//
2626
// The supported commands are:
2727
//
28-
// - build-scalar [<index-var-types> ...]
28+
// - build-scalar
2929
//
3030
// Builds an expression tree from a scalar SQL expression and outputs a
3131
// representation of the tree. The expression can refer to external variables
3232
// using @1, @2, etc. in which case the types of the variables must be passed
3333
// on the command line.
34+
//
35+
// - legacy-normalize
36+
//
37+
// Runs the TypedExpr normalization code and rebuilds the scalar expression.
38+
// If present, must follow build-scalar.
39+
//
40+
// - normalize
41+
//
42+
// Normalizes the expression. If present, must follow build-scalar or
43+
// legacy-normalize.
44+
//
45+
// - index-constraints
46+
//
47+
// Creates index constraints on the assumption that the index is formed by
48+
// the index var columns (as specified by <index-var-types>).
49+
// If present, build-scalar must have been an earlier command.
3450

3551
import (
3652
"bufio"
@@ -229,7 +245,18 @@ func TestOpt(t *testing.T) {
229245
runTest(t, path, func(d *testdata) string {
230246
var e *expr
231247
var types []types.T
248+
var typedExpr tree.TypedExpr
232249

250+
buildScalarFn := func() {
251+
defer func() {
252+
if r := recover(); r != nil {
253+
d.fatalf(t, "buildScalar: %v", r)
254+
}
255+
}()
256+
e = buildScalar(&buildContext{}, typedExpr)
257+
}
258+
259+
evalCtx := tree.MakeTestingEvalContext()
233260
for _, cmd := range strings.Split(d.cmd, ",") {
234261
switch cmd {
235262
case "build-scalar":
@@ -238,27 +265,26 @@ func TestOpt(t *testing.T) {
238265
if err != nil {
239266
d.fatalf(t, "%v", err)
240267
}
241-
typedExpr, err := parseScalarExpr(d.sql, types)
268+
typedExpr, err = parseScalarExpr(d.sql, types)
242269
if err != nil {
243270
d.fatalf(t, "%v", err)
244271
}
245272

246-
e = func() *expr {
247-
defer func() {
248-
if r := recover(); r != nil {
249-
d.fatalf(t, "buildScalar: %v", r)
250-
}
251-
}()
252-
return buildScalar(&buildContext{}, typedExpr)
253-
}()
273+
buildScalarFn()
274+
case "legacy-normalize":
275+
// Apply the TypedExpr normalization and rebuild the expression.
276+
typedExpr, err = evalCtx.NormalizeExpr(typedExpr)
277+
if err != nil {
278+
d.fatalf(t, "%v", err)
279+
}
280+
buildScalarFn()
254281
case "normalize":
255282
normalizeScalar(e)
256283
case "index-constraints":
257284
if e == nil {
258285
d.fatalf(t, "no expression for index-constraints")
259286
}
260287

261-
evalCtx := tree.MakeTestingEvalContext()
262288
spans := MakeIndexConstraints(e, types, &evalCtx)
263289
var buf bytes.Buffer
264290
for _, sp := range spans {

‎pkg/sql/opt/scalar.go‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,27 @@ func isIndexedVar(e *expr, index int) bool {
144144
return e.op == variableOp && e.private.(int) == index
145145
}
146146

147+
func initTupleExpr(e *expr, children []*expr) {
148+
// In general, the order matters in a tuple so we use an "ordered list"
149+
// operator. In some cases (IN) the order doesn't matter; we could convert
150+
// those to listOp during normalization, but there doesn't seem to be a
151+
// benefit at this time.
152+
e.op = orderedListOp
153+
e.children = children
154+
}
155+
156+
func isTupleOfConstants(e *expr) bool {
157+
if e.op != orderedListOp {
158+
return false
159+
}
160+
for _, c := range e.children {
161+
if c.op != constOp {
162+
return false
163+
}
164+
}
165+
return true
166+
}
167+
147168
// Applies a set of normalization rules to a scalar expression.
148169
//
149170
// For now, we expect to build exprs from TypedExprs which have gone through a

‎pkg/sql/opt/testdata/build-scalar‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,3 +131,38 @@ or (type: bool)
131131
└── gt (type: bool)
132132
├── variable (2) (type: int)
133133
└── const (2) (type: int)
134+
135+
build-scalar int int
136+
(@1, @2) = (1, 2)
137+
----
138+
eq (type: bool)
139+
├── ordered-list (type: tuple{int, int})
140+
│ ├── variable (0) (type: int)
141+
│ └── variable (1) (type: int)
142+
└── ordered-list (type: tuple{int, int})
143+
├── const (1) (type: int)
144+
└── const (2) (type: int)
145+
146+
build-scalar int
147+
@1 IN (1, 2)
148+
----
149+
in (type: bool)
150+
├── variable (0) (type: int)
151+
└── ordered-list (type: tuple{int, int})
152+
├── const (1) (type: int)
153+
└── const (2) (type: int)
154+
155+
build-scalar int int
156+
(@1, @2) IN ((1, 2), (3, 4))
157+
----
158+
in (type: bool)
159+
├── ordered-list (type: tuple{int, int})
160+
│ ├── variable (0) (type: int)
161+
│ └── variable (1) (type: int)
162+
└── ordered-list (type: tuple{tuple{int, int}, tuple{int, int}})
163+
├── ordered-list (type: tuple{int, int})
164+
│ ├── const (1) (type: int)
165+
│ └── const (2) (type: int)
166+
└── ordered-list (type: tuple{int, int})
167+
├── const (3) (type: int)
168+
└── const (4) (type: int)

‎pkg/sql/opt/testdata/index-constraints‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,3 +113,53 @@ build-scalar,normalize,index-constraints decimal decimal
113113
@1 > 1.5 AND @2 > 2
114114
----
115115
(/1.5 - ]
116+
117+
build-scalar,normalize,index-constraints int
118+
@1 IN (1, 2, 3)
119+
----
120+
[/1 - /1]
121+
[/2 - /2]
122+
[/3 - /3]
123+
124+
build-scalar,legacy-normalize,index-constraints int
125+
@1 IN (1, 5, 1, 4)
126+
----
127+
[/1 - /1]
128+
[/4 - /4]
129+
[/5 - /5]
130+
131+
build-scalar,normalize,index-constraints int int
132+
@1 = 1 AND @2 IN (1, 2, 3)
133+
----
134+
[/1/1 - /1/1]
135+
[/1/2 - /1/2]
136+
[/1/3 - /1/3]
137+
138+
build-scalar,normalize,index-constraints int int
139+
@1 IN (1, 2) AND @2 IN (1, 2, 3)
140+
----
141+
[/1/1 - /1/1]
142+
[/1/2 - /1/2]
143+
[/1/3 - /1/3]
144+
[/2/1 - /2/1]
145+
[/2/2 - /2/2]
146+
[/2/3 - /2/3]
147+
148+
build-scalar,normalize,index-constraints int int
149+
@1 >= 2 AND @1 <= 4 AND @2 IN (1, 2, 3)
150+
----
151+
[/2/1 - /4/3]
152+
153+
build-scalar,normalize,index-constraints int int
154+
@1 IN (1, 2, 3) AND @2 = 4
155+
----
156+
[/1/4 - /1/4]
157+
[/2/4 - /2/4]
158+
[/3/4 - /3/4]
159+
160+
build-scalar,normalize,index-constraints int int
161+
@1 IN (1, 2, 3) AND @2 >= 2 AND @2 <= 4
162+
----
163+
[/1/2 - /1/4]
164+
[/2/2 - /2/4]
165+
[/3/2 - /3/4]

0 commit comments

Comments
 (0)