opt: allow IN subquery to be converted to lookup join

This change adds a rule that handles a case which prevents Exists subqueries from becoming lookup joins. Fixes cockroachdb#43198. Release note (performance improvement): certain queries containing `<tuple> IN (<subquery>)` conditions may run significantly faster.
adityamaru · Apr 19, 2021 · e6fd187 · e6fd187
1 parent 8f1f382
commit e6fd187
Show file tree

Hide file tree

Showing 5 changed files with 322 additions and 78 deletions.
diff --git a/pkg/sql/opt/norm/rules/scalar.opt b/pkg/sql/opt/norm/rules/scalar.opt
@@ -126,6 +126,60 @@ $input
 =>
 (Exists $input $subqueryPrivate)
 
+# InlineExistsSelectTuple splits a tuple equality filter into multiple
+# (per-column) equalities, in the case where the tuple on one side is being
+# projected.
+#
+# We are specifically handling the case when this is under Exists because we
+# don't have to keep the same output columns for the Select. This case is
+# important because it is produced for an IN subquery:
+#
+#   SELECT * FROM ab WHERE (a, b) IN (SELECT c, d FROM cd)
+#
+# Without this rule, we would not be able to produce a lookup join plan for such
+# a query.
+#
+[InlineExistsSelectTuple, Normalize]
+(Exists
+    (Select
+        (Project
+            $input:*
+            [
+                ...
+                (ProjectionsItem $tuple:(Tuple) $tupleCol:*)
+                ...
+            ]
+        )
+        $filters:[
+            ...
+            $item:(FiltersItem
+                (Eq
+                    # CommuteVar ensures that the variable is on the left.
+                    (Variable
+                        $varCol:* &
+                            (EqualsColumn $varCol $tupleCol)
+                    )
+                    $rhs:(Tuple) &
+                        (TuplesHaveSameLength $tuple $rhs)
+                )
+            )
+            ...
+        ]
+    )
+    $subqueryPrivate:*
+)
+=>
+(Exists
+    (Select
+        $input
+        (ConcatFilters
+            (RemoveFiltersItem $filters $item)
+            (SplitTupleEq $tuple $rhs)
+        )
+    )
+    $subqueryPrivate
+)
+
 # IntroduceExistsLimit inserts a LIMIT 1 "under" Exists so as to save resources
 # to make the EXISTS determination.
 #

diff --git a/pkg/sql/opt/norm/scalar_funcs.go b/pkg/sql/opt/norm/scalar_funcs.go
@@ -305,3 +305,27 @@ func (c *CustomFuncs) IsTupleOfVars(t *memo.TupleExpr, cols opt.ColList) bool {
 func (c *CustomFuncs) VarsAreSame(left, right *memo.VariableExpr) bool {
 	return left.Col == right.Col
 }
+
+// EqualsColumn returns true if the two column IDs are the same.
+func (c *CustomFuncs) EqualsColumn(left, right opt.ColumnID) bool {
+	return left == right
+}
+
+// TuplesHaveSameLength returns true if two tuples have the same number of
+// elements.
+func (c *CustomFuncs) TuplesHaveSameLength(a, b *memo.TupleExpr) bool {
+	return len(a.Elems) == len(b.Elems)
+}
+
+// SplitTupleEq splits an equality condition between two tuples into multiple
+// equalities, one for each tuple column.
+func (c *CustomFuncs) SplitTupleEq(lhs, rhs *memo.TupleExpr) memo.FiltersExpr {
+	if len(lhs.Elems) != len(rhs.Elems) {
+		panic(errors.AssertionFailedf("unequal tuple lengths"))
+	}
+	res := make(memo.FiltersExpr, len(lhs.Elems))
+	for i := range res {
+		res[i] = c.f.ConstructFiltersItem(c.f.ConstructEq(lhs.Elems[i], rhs.Elems[i]))
+	}
+	return res
+}
diff --git a/pkg/sql/opt/norm/testdata/rules/assign_placeholders b/pkg/sql/opt/norm/testdata/rules/assign_placeholders
@@ -60,76 +60,112 @@ select
  └── filters
       └── (a:1, b:2) IN ((1, 2),) [outer=(1,2), constraints=(/1/2: [/1/2 - /1/2]; /2: [/2 - /2]; tight), fd=()-->(1,2)]
 
-# TODO(radu): hoisting of the subquery results in projecting the tuple, which
-# gets in the way of simplifying the expression. In particular, we can't
-# convert this to a lookup join (see next testcase), which is unfortunate.
+# The normalized expression above can be explored into a constrained scan.
+opt
+SELECT * FROM abcd WHERE (a, b) IN (
+  SELECT unnest('{1}'::INT[]),
+         unnest('{2}'::INT[])
+)
+----
+scan abcd
+ ├── columns: a:1!null b:2!null c:3!null d:4
+ ├── constraint: /1/2/3: [/1/2 - /1/2]
+ ├── key: (3)
+ └── fd: ()-->(1,2), (3)-->(4)
+
 assign-placeholders-norm query-args=('{1}','{2}')
 SELECT * FROM abcd WHERE (a, b) IN (
   SELECT unnest($1:::STRING::INT[]),
          unnest($2:::STRING::INT[])
 )
 ----
-project
+select
  ├── columns: a:1!null b:2!null c:3!null d:4
- ├── immutable
+ ├── key: (3)
+ ├── fd: ()-->(1,2), (3)-->(4)
+ ├── scan abcd
+ │    ├── columns: a:1!null b:2!null c:3!null d:4
+ │    ├── key: (1-3)
+ │    └── fd: (1-3)-->(4)
+ └── filters
+      ├── a:1 = 1 [outer=(1), constraints=(/1: [/1 - /1]; tight), fd=()-->(1)]
+      └── b:2 = 2 [outer=(2), constraints=(/2: [/2 - /2]; tight), fd=()-->(2)]
+
+# We want this query to be optimized into a constrained scan, just like the
+# no-placeholders variant above.
+assign-placeholders-opt query-args=('{1}','{2}')
+SELECT * FROM abcd WHERE (a, b) IN (
+  SELECT unnest($1:::STRING::INT[]),
+         unnest($2:::STRING::INT[])
+)
+----
+scan abcd
+ ├── columns: a:1!null b:2!null c:3!null d:4
+ ├── constraint: /1/2/3: [/1/2 - /1/2]
+ ├── key: (3)
+ └── fd: ()-->(1,2), (3)-->(4)
+
+# Note: \x2c is a comma; we can't use a comma directly because of the
+# datadriven parser.
+assign-placeholders-norm query-args=('{1\x2c 2}','{3\x2c 4}')
+SELECT * FROM abcd WHERE (a, b) IN (
+  SELECT unnest($1:::STRING::INT[]),
+         unnest($2:::STRING::INT[])
+)
+----
+semi-join (hash)
+ ├── columns: a:1!null b:2!null c:3!null d:4
+ ├── stable
  ├── key: (1-3)
  ├── fd: (1-3)-->(4)
- └── semi-join (hash)
-      ├── columns: a:1!null b:2!null c:3!null d:4 column9:9!null
-      ├── immutable
-      ├── key: (1-3)
-      ├── fd: (1-3)-->(4), (1,2)-->(9)
-      ├── project
-      │    ├── columns: column9:9!null a:1!null b:2!null c:3!null d:4
-      │    ├── key: (1-3)
-      │    ├── fd: (1-3)-->(4), (1,2)-->(9)
-      │    ├── scan abcd
-      │    │    ├── columns: a:1!null b:2!null c:3!null d:4
-      │    │    ├── key: (1-3)
-      │    │    └── fd: (1-3)-->(4)
-      │    └── projections
-      │         └── (a:1, b:2) [as=column9:9, outer=(1,2)]
-      ├── values
-      │    ├── columns: column8:8
-      │    ├── cardinality: [1 - 1]
-      │    ├── key: ()
-      │    ├── fd: ()-->(8)
-      │    └── ((1, 2),)
-      └── filters
-           └── column9:9 = column8:8 [outer=(8,9), immutable, constraints=(/8: (/NULL - ]; /9: (/NULL - ]), fd=(8)==(9), (9)==(8)]
+ ├── scan abcd
+ │    ├── columns: a:1!null b:2!null c:3!null d:4
+ │    ├── key: (1-3)
+ │    └── fd: (1-3)-->(4)
+ ├── project-set
+ │    ├── columns: unnest:6 unnest:7
+ │    ├── stable
+ │    ├── values
+ │    │    ├── cardinality: [1 - 1]
+ │    │    ├── key: ()
+ │    │    └── ()
+ │    └── zip
+ │         ├── unnest(e'{1\\x2c 2}'::INT8[]) [stable]
+ │         └── unnest(e'{3\\x2c 4}'::INT8[]) [stable]
+ └── filters
+      ├── unnest:6 = a:1 [outer=(1,6), constraints=(/1: (/NULL - ]; /6: (/NULL - ]), fd=(1)==(6), (6)==(1)]
+      └── unnest:7 = b:2 [outer=(2,7), constraints=(/2: (/NULL - ]; /7: (/NULL - ]), fd=(2)==(7), (7)==(2)]
 
-# TODO(radu): the end result involves a full table scan.
-assign-placeholders-opt query-args=('{1}','{2}')
+assign-placeholders-opt query-args=('{1\x2c 2}','{3\x2c 4}')
 SELECT * FROM abcd WHERE (a, b) IN (
   SELECT unnest($1:::STRING::INT[]),
          unnest($2:::STRING::INT[])
 )
 ----
 project
  ├── columns: a:1!null b:2!null c:3!null d:4
- ├── immutable
+ ├── stable
  ├── key: (1-3)
  ├── fd: (1-3)-->(4)
- └── semi-join (hash)
-      ├── columns: a:1!null b:2!null c:3!null d:4 column9:9!null
-      ├── immutable
-      ├── key: (1-3)
-      ├── fd: (1-3)-->(4), (1,2)-->(9)
-      ├── project
-      │    ├── columns: column9:9!null a:1!null b:2!null c:3!null d:4
-      │    ├── key: (1-3)
-      │    ├── fd: (1-3)-->(4), (1,2)-->(9)
-      │    ├── scan abcd
-      │    │    ├── columns: a:1!null b:2!null c:3!null d:4
-      │    │    ├── key: (1-3)
-      │    │    └── fd: (1-3)-->(4)
-      │    └── projections
-      │         └── (a:1, b:2) [as=column9:9, outer=(1,2)]
-      ├── values
-      │    ├── columns: column8:8
-      │    ├── cardinality: [1 - 1]
-      │    ├── key: ()
-      │    ├── fd: ()-->(8)
-      │    └── ((1, 2),)
-      └── filters
-           └── column9:9 = column8:8 [outer=(8,9), immutable, constraints=(/8: (/NULL - ]; /9: (/NULL - ]), fd=(8)==(9), (9)==(8)]
+ └── inner-join (lookup abcd)
+      ├── columns: a:1!null b:2!null c:3!null d:4 unnest:6!null unnest:7!null
+      ├── key columns: [6 7] = [1 2]
+      ├── stable
+      ├── key: (3,6,7)
+      ├── fd: (1-3)-->(4), (1)==(6), (6)==(1), (2)==(7), (7)==(2)
+      ├── distinct-on
+      │    ├── columns: unnest:6 unnest:7
+      │    ├── grouping columns: unnest:6 unnest:7
+      │    ├── stable
+      │    ├── key: (6,7)
+      │    └── project-set
+      │         ├── columns: unnest:6 unnest:7
+      │         ├── stable
+      │         ├── values
+      │         │    ├── cardinality: [1 - 1]
+      │         │    ├── key: ()
+      │         │    └── ()
+      │         └── zip
+      │              ├── unnest(e'{1\\x2c 2}'::INT8[]) [stable]
+      │              └── unnest(e'{3\\x2c 4}'::INT8[]) [stable]
+      └── filters (true)