Skip to content

Commit 3eab8d4

Browse files
RaduBerindeRaduBerinde
committed
opt: remember subquery ASTs for explain
Store the original `*tree.Subquery` in `Any`, `Exists`, `Subquery` and use it to populate the subquery during execbuild. This causes the original subquery to show up in `EXPLAIN`. Fixes #29350. Release note: None
1 parent 7f98419 commit 3eab8d4

File tree

14 files changed

+156
-66
lines changed

14 files changed

+156
-66
lines changed

‎pkg/sql/opt/exec/execbuilder/scalar_builder.go

+14-6
Original file line numberDiff line numberDiff line change
@@ -360,15 +360,16 @@ func (b *Builder) buildAny(ctx *buildScalarCtx, ev memo.ExprView) (tree.TypedExp
360360
types.Types[val] = ev.Metadata().ColumnType(opt.ColumnID(key))
361361
})
362362

363-
subqueryExpr := b.addSubquery(exec.SubqueryAnyRows, types, plan.root)
363+
def := ev.Private().(*memo.SubqueryDef)
364+
subqueryExpr := b.addSubquery(exec.SubqueryAnyRows, types, plan.root, def.OriginalExpr)
364365

365366
// Build the scalar value that is compared against each row.
366367
scalar, err := b.buildScalar(ctx, ev.Child(1))
367368
if err != nil {
368369
return nil, err
369370
}
370371

371-
cmp := opt.ComparisonOpReverseMap[ev.Private().(opt.Operator)]
372+
cmp := opt.ComparisonOpReverseMap[def.Cmp]
372373
return tree.NewTypedComparisonExprWithSubOp(tree.Any, cmp, scalar, subqueryExpr), nil
373374
}
374375

@@ -388,7 +389,8 @@ func (b *Builder) buildExistsSubquery(
388389
return nil, err
389390
}
390391

391-
return b.addSubquery(exec.SubqueryExists, types.Bool, root), nil
392+
def := ev.Private().(*memo.SubqueryDef)
393+
return b.addSubquery(exec.SubqueryExists, types.Bool, root, def.OriginalExpr), nil
392394
}
393395

394396
func (b *Builder) buildSubquery(ctx *buildScalarCtx, ev memo.ExprView) (tree.TypedExpr, error) {
@@ -417,13 +419,19 @@ func (b *Builder) buildSubquery(ctx *buildScalarCtx, ev memo.ExprView) (tree.Typ
417419
return nil, err
418420
}
419421

420-
return b.addSubquery(exec.SubqueryOneRow, ev.Logical().Scalar.Type, root), nil
422+
def := ev.Private().(*memo.SubqueryDef)
423+
return b.addSubquery(exec.SubqueryOneRow, ev.Logical().Scalar.Type, root, def.OriginalExpr), nil
421424
}
422425

423426
// addSubquery adds an entry to b.subqueries and creates a tree.Subquery
424427
// expression node associated with it.
425-
func (b *Builder) addSubquery(mode exec.SubqueryMode, typ types.T, root exec.Node) *tree.Subquery {
426-
exprNode := &tree.Subquery{Exists: mode == exec.SubqueryExists}
428+
func (b *Builder) addSubquery(
429+
mode exec.SubqueryMode, typ types.T, root exec.Node, originalExpr *tree.Subquery,
430+
) *tree.Subquery {
431+
exprNode := &tree.Subquery{
432+
Select: originalExpr.Select,
433+
Exists: mode == exec.SubqueryExists,
434+
}
427435
exprNode.SetType(typ)
428436
b.subqueries = append(b.subqueries, exec.Subquery{
429437
ExprNode: exprNode,

‎pkg/sql/opt/exec/execbuilder/testdata/srfs

+19-19
Original file line numberDiff line numberDiff line change
@@ -163,25 +163,25 @@ render · ·
163163
query TTTTT
164164
EXPLAIN (VERBOSE) SELECT generate_series((SELECT unnest(ARRAY[x, y]) FROM xy), z) FROM xz
165165
----
166-
root · · (generate_series) ·
167-
├── render · · (generate_series) ·
168-
│ │ render 0 generate_series · ·
169-
│ └── project set · · (z, generate_series) ·
170-
│ │ render 0 generate_series(@S1, @1) · ·
171-
│ └── scan · · (z) ·
172-
│ table xz@primary · ·
173-
│ spans ALL · ·
174-
└── subquery · · (generate_series) ·
175-
│ id @S1 · ·
176-
│ original sql <unknown> · ·
177-
│ exec mode one row · ·
178-
└── render · · (unnest) ·
179-
│ render 0 unnest · ·
180-
└── project set · · (x, y, unnest) ·
181-
│ render 0 unnest(ARRAY[@1, @2]) · ·
182-
└── scan · · (x, y) ·
183-
· table xy@primary · ·
184-
· spans ALL · ·
166+
root · · (generate_series) ·
167+
├── render · · (generate_series) ·
168+
│ │ render 0 generate_series · ·
169+
│ └── project set · · (z, generate_series) ·
170+
│ │ render 0 generate_series(@S1, @1) · ·
171+
│ └── scan · · (z) ·
172+
│ table xz@primary · ·
173+
│ spans ALL · ·
174+
└── subquery · · (generate_series) ·
175+
│ id @S1 · ·
176+
│ original sql (SELECT unnest(ARRAY[x, y]) FROM xy) · ·
177+
│ exec mode one row · ·
178+
└── render · · (unnest) ·
179+
│ render 0 unnest · ·
180+
└── project set · · (x, y, unnest) ·
181+
│ render 0 unnest(ARRAY[@1, @2]) · ·
182+
└── scan · · (x, y) ·
183+
· table xy@primary · ·
184+
· spans ALL · ·
185185

186186
# Regression test for #24676.
187187
statement ok

‎pkg/sql/opt/exec/execbuilder/testdata/subquery

+27-27
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ root · ·
3636
│ └── emptyrow · ·
3737
└── subquery · ·
3838
│ id @S1
39-
│ original sql EXISTS <unknown>
39+
│ original sql EXISTS (SELECT a FROM abc)
4040
│ exec mode exists
4141
└── scan · ·
4242
· table abc@primary
@@ -45,30 +45,30 @@ root · ·
4545
query TTTTT
4646
EXPLAIN (VERBOSE) SELECT * FROM abc WHERE a = (SELECT max(a) FROM abc WHERE EXISTS(SELECT * FROM abc WHERE c=a+3))
4747
----
48-
root · · (a, b, c) ·
49-
├── scan · · (a, b, c) ·
50-
│ table abc@primary · ·
51-
│ spans ALL · ·
52-
│ filter a = @S2 · ·
53-
├── subquery · · (a, b, c) ·
54-
│ │ id @S1 · ·
55-
│ │ original sql EXISTS <unknown> · ·
56-
│ │ exec mode exists · ·
57-
│ └── scan · · (a, b, c) ·
58-
│ table abc@primary · ·
59-
│ spans ALL · ·
60-
│ filter c = (a + 3) · ·
61-
└── subquery · · (a, b, c) ·
62-
│ id @S2 · ·
63-
│ original sql <unknown> · ·
64-
│ exec mode one row · ·
65-
└── group · · (agg0) ·
66-
│ aggregate 0 max(a) · ·
67-
│ scalar · · ·
68-
└── scan · · (a) ·
69-
· table abc@primary · ·
70-
· spans ALL · ·
71-
· filter @S1 · ·
48+
root · · (a, b, c) ·
49+
├── scan · · (a, b, c) ·
50+
│ table abc@primary · ·
51+
│ spans ALL · ·
52+
│ filter a = @S2 · ·
53+
├── subquery · · (a, b, c) ·
54+
│ │ id @S1 · ·
55+
│ │ original sql EXISTS (SELECT * FROM abc WHERE c = (a + 3)) · ·
56+
│ │ exec mode exists · ·
57+
│ └── scan · · (a, b, c) ·
58+
│ table abc@primary · ·
59+
│ spans ALL · ·
60+
│ filter c = (a + 3) · ·
61+
└── subquery · · (a, b, c) ·
62+
│ id @S2 · ·
63+
│ original sql (SELECT max(a) FROM abc WHERE EXISTS (SELECT * FROM abc WHERE c = (a + 3))) · ·
64+
│ exec mode one row · ·
65+
└── group · · (agg0) ·
66+
│ aggregate 0 max(a) · ·
67+
│ scalar · · ·
68+
└── scan · · (a) ·
69+
· table abc@primary · ·
70+
· spans ALL · ·
71+
· filter @S1 · ·
7272

7373
# IN expression transformed into semi-join.
7474
query TTTTT
@@ -103,7 +103,7 @@ root · ·
103103
│ size 1 column, 2 rows
104104
└── subquery · ·
105105
│ id @S1
106-
│ original sql <unknown>
106+
│ original sql (SELECT 2)
107107
│ exec mode one row
108108
└── render · ·
109109
└── emptyrow · ·
@@ -228,7 +228,7 @@ root · · ("array") ·
228228
│ └── emptyrow · · () ·
229229
└── subquery · · ("array") ·
230230
│ id @S1 · ·
231-
│ original sql <unknown> · ·
231+
│ original sql (SELECT x FROM b) · ·
232232
│ exec mode one row · ·
233233
└── group · · (agg0) ·
234234
│ aggregate 0 array_agg(x) · ·

‎pkg/sql/opt/memo/expr_view_format.go

+8
Original file line numberDiff line numberDiff line change
@@ -387,6 +387,14 @@ func (ev ExprView) formatScalarPrivate(f *ExprFmtCtx, private interface{}) {
387387
// columns for their containing op (Project or GroupBy), so no need to
388388
// print again.
389389
private = nil
390+
391+
case opt.AnyOp:
392+
// We don't want to show the OriginalExpr; just show Cmp.
393+
private = private.(*SubqueryDef).Cmp
394+
395+
case opt.SubqueryOp, opt.ExistsOp:
396+
// We don't want to show the OriginalExpr.
397+
private = nil
390398
}
391399

392400
if private != nil {

‎pkg/sql/opt/memo/private_defs.go

+8
Original file line numberDiff line numberDiff line change
@@ -459,3 +459,11 @@ func (m *MergeOnDef) CanProvideOrdering(required *props.OrderingChoice) bool {
459459
return false
460460
}
461461
}
462+
463+
// SubqueryDef contains information related to a subquery (Subquery, Any,
464+
// Exists).
465+
type SubqueryDef struct {
466+
OriginalExpr *tree.Subquery
467+
// Cmp is only used for AnyOp.
468+
Cmp opt.Operator
469+
}

‎pkg/sql/opt/memo/private_storage.go

+12
Original file line numberDiff line numberDiff line change
@@ -413,6 +413,18 @@ func (ps *privateStorage) internRowNumberDef(def *RowNumberDef) PrivateID {
413413
return ps.addValue(privateKey{iface: typ, str: ps.keyBuf.String()}, def)
414414
}
415415

416+
func (ps *privateStorage) internSubqueryDef(def *SubqueryDef) PrivateID {
417+
ps.keyBuf.Reset()
418+
ps.keyBuf.writeUvarint(uint64(uintptr(unsafe.Pointer(def.OriginalExpr))))
419+
ps.keyBuf.writeUvarint(uint64(def.Cmp))
420+
421+
typ := (*SubqueryDef)(nil)
422+
if id, ok := ps.privatesMap[privateKey{iface: typ, str: ps.keyBuf.String()}]; ok {
423+
return id
424+
}
425+
return ps.addValue(privateKey{iface: typ, str: ps.keyBuf.String()}, def)
426+
}
427+
416428
// internSetOpColMap adds the given value to storage and returns an id that can
417429
// later be used to retrieve the value by calling the lookup method. If the
418430
// value has been previously added to storage, then internSetOpColMap always

‎pkg/sql/opt/memo/private_storage_test.go

+37
Original file line numberDiff line numberDiff line change
@@ -564,6 +564,39 @@ func TestInternFuncOpDef(t *testing.T) {
564564
test(funcDef2, funcDef3, false)
565565
}
566566

567+
func TestInternSubqueryDef(t *testing.T) {
568+
var ps privateStorage
569+
ps.init()
570+
571+
test := func(left, right *SubqueryDef, expected bool) {
572+
t.Helper()
573+
leftID := ps.internSubqueryDef(left)
574+
rightID := ps.internSubqueryDef(right)
575+
if (leftID == rightID) != expected {
576+
t.Errorf("%v == %v, expected %v, got %v", left, right, expected, !expected)
577+
}
578+
}
579+
580+
expr1 := &tree.Subquery{}
581+
expr2 := &tree.Subquery{}
582+
583+
defs := []SubqueryDef{
584+
{},
585+
{OriginalExpr: expr1},
586+
{OriginalExpr: expr2},
587+
{Cmp: opt.LeOp},
588+
{Cmp: opt.GtOp},
589+
{OriginalExpr: expr1, Cmp: opt.GtOp},
590+
{OriginalExpr: expr1, Cmp: opt.LeOp},
591+
{OriginalExpr: expr2, Cmp: opt.GtOp},
592+
}
593+
for i := range defs {
594+
for j := range defs {
595+
test(&defs[i], &defs[j], i == j)
596+
}
597+
}
598+
}
599+
567600
func TestInternSetOpColMap(t *testing.T) {
568601
var ps privateStorage
569602
ps.init()
@@ -883,6 +916,7 @@ func TestPrivateStorageAllocations(t *testing.T) {
883916
}
884917
indexJoinDef := &IndexJoinDef{Table: 1, Cols: colSet}
885918
lookupJoinDef := &LookupJoinDef{Table: 1, Index: 2, KeyCols: colList, LookupCols: colSet}
919+
subqueryDef := &SubqueryDef{OriginalExpr: &tree.Subquery{}, Cmp: opt.LtOp}
886920
setOpColMap := &SetOpColMap{Left: colList, Right: colList, Out: colList}
887921
datum := tree.NewDInt(1)
888922
typ := types.Int
@@ -903,6 +937,7 @@ func TestPrivateStorageAllocations(t *testing.T) {
903937
ps.internMergeOnDef(mergeOnDef)
904938
ps.internIndexJoinDef(indexJoinDef)
905939
ps.internLookupJoinDef(lookupJoinDef)
940+
ps.internSubqueryDef(subqueryDef)
906941
ps.internSetOpColMap(setOpColMap)
907942
ps.internDatum(datum)
908943
ps.internType(typ)
@@ -943,6 +978,7 @@ func BenchmarkPrivateStorage(b *testing.B) {
943978
}
944979
indexJoinDef := &IndexJoinDef{Table: 1, Cols: colSet}
945980
lookupJoinDef := &LookupJoinDef{Table: 1, Index: 2, KeyCols: colList, LookupCols: colSet}
981+
subqueryDef := &SubqueryDef{OriginalExpr: &tree.Subquery{}, Cmp: opt.LtOp}
946982
setOpColMap := &SetOpColMap{Left: colList, Right: colList, Out: colList}
947983
datum := tree.NewDInt(1)
948984
typ := types.Int
@@ -965,6 +1001,7 @@ func BenchmarkPrivateStorage(b *testing.B) {
9651001
ps.internMergeOnDef(mergeOnDef)
9661002
ps.internIndexJoinDef(indexJoinDef)
9671003
ps.internLookupJoinDef(lookupJoinDef)
1004+
ps.internSubqueryDef(subqueryDef)
9681005
ps.internSetOpColMap(setOpColMap)
9691006
ps.internDatum(datum)
9701007
ps.internType(typ)

‎pkg/sql/opt/norm/decorrelate.go

+4-3
Original file line numberDiff line numberDiff line change
@@ -611,10 +611,11 @@ func (c *CustomFuncs) AddColsToGroupByDef(
611611
// expression with the first (and only) column of the input rowset, using the
612612
// given comparison operator.
613613
func (c *CustomFuncs) ConstructAnyCondition(
614-
input, scalar memo.GroupID, cmp memo.PrivateID,
614+
input, scalar memo.GroupID, subqueryDef memo.PrivateID,
615615
) memo.GroupID {
616616
inputVar := c.referenceSingleColumn(input)
617-
return c.ConstructBinary(c.f.mem.LookupPrivate(cmp).(opt.Operator), scalar, inputVar)
617+
def := c.f.mem.LookupPrivate(subqueryDef).(*memo.SubqueryDef)
618+
return c.ConstructBinary(def.Cmp, scalar, inputVar)
618619
}
619620

620621
// ConstructBinary builds a dynamic binary expression, given the binary
@@ -731,7 +732,7 @@ func (r *subqueryHoister) hoistAll(root memo.GroupID) memo.GroupID {
731732
case opt.AnyOp:
732733
input := ev.ChildGroup(0)
733734
scalar := ev.ChildGroup(1)
734-
cmp := ev.Private().(opt.Operator)
735+
cmp := ev.Private().(*memo.SubqueryDef).Cmp
735736
subquery = r.constructGroupByAny(scalar, cmp, input)
736737
}
737738

‎pkg/sql/opt/norm/rules/decorrelate.opt

+6-4
Original file line numberDiff line numberDiff line change
@@ -674,7 +674,7 @@
674674
# addition, the Exists can be transformed into a semi-join.
675675
[NormalizeAnyFilter, Normalize]
676676
(Filters
677-
$list:[ ... $any:(Any $input:* $scalar:* $cmp:*) ... ]
677+
$list:[ ... $any:(Any $input:* $scalar:* $subqueryDef:*) ... ]
678678
)
679679
=>
680680
(Filters
@@ -684,8 +684,9 @@
684684
(Exists
685685
(Select
686686
$input
687-
(Filters [ (ConstructAnyCondition $input $scalar $cmp) ])
687+
(Filters [ (ConstructAnyCondition $input $scalar $subqueryDef) ])
688688
)
689+
$subqueryDef
689690
)
690691
)
691692
)
@@ -700,7 +701,7 @@
700701
# Citations: [5] (section 3.5)
701702
[NormalizeNotAnyFilter, Normalize]
702703
(Filters
703-
$list:[ ... $notany:(Not (Any $input:* $scalar:* $cmp:*)) ... ]
704+
$list:[ ... $notany:(Not (Any $input:* $scalar:* $subqueryDef:*)) ... ]
704705
)
705706
=>
706707
(Filters
@@ -712,9 +713,10 @@
712713
(Select
713714
$input
714715
(Filters
715-
[ (IsNot (ConstructAnyCondition $input $scalar $cmp) (False)) ]
716+
[ (IsNot (ConstructAnyCondition $input $scalar $subqueryDef) (False)) ]
716717
)
717718
)
719+
$subqueryDef
718720
)
719721
)
720722
)

‎pkg/sql/opt/norm/rules/scalar.opt

+2-2
Original file line numberDiff line numberDiff line change
@@ -143,15 +143,15 @@
143143
# Project operator never changes the row cardinality of its input, and row
144144
# cardinality is the only thing that Exists cares about, so Project is a no-op.
145145
[EliminateExistsProject, Normalize]
146-
(Exists (Project $input:*)) => (Exists $input)
146+
(Exists (Project $input:*) $def:*) => (Exists $input $def)
147147

148148
# EliminateExistsGroupBy discards a non-scalar GroupBy input to the Exists
149149
# operator. While non-scalar GroupBy (or DistinctOn) can change row cardinality,
150150
# it always returns a non-empty set if its input is non-empty. Similarly, if its
151151
# input is empty, then it returns the empty set. Therefore, it's a no-op for
152152
# Exists.
153153
[EliminateExistsGroupBy, Normalize]
154-
(Exists (GroupBy | DistinctOn $input:*)) => (Exists $input)
154+
(Exists (GroupBy | DistinctOn $input:*) $def:*) => (Exists $input $def)
155155

156156
# NormalizeJSONFieldAccess transforms field access into a containment with a
157157
# simpler LHS. This allows inverted index constraints to be generated in some

0 commit comments

Comments
 (0)