sql: Add support for JOIN LATERAL syntax

As correlated subqueries are already supported, only minor scoping changes are required to fully support Postgres' JOIN LATERAL syntax. h/t @justinj for the support Release note (sql change): Adds support for JOIN LATERAL syntax.
cockroachdb · Oct 7, 2019 · 6ee4fee · 6ee4fee
1 parent 4ff8aad
commit 6ee4fee
Show file tree

Hide file tree

Showing 3 changed files with 120 additions and 2 deletions.
diff --git a/pkg/sql/logictest/testdata/logic_test/subquery_correlated b/pkg/sql/logictest/testdata/logic_test/subquery_correlated
@@ -1170,3 +1170,48 @@ LIMIT 5;
 {}
 {}
 {}
+
+# Customers, their billing address, and all orders not going to their billing address
+query ITT rowsort
+SELECT
+    c_id, bill, states
+FROM
+    c
+    JOIN LATERAL (
+            SELECT
+                COALESCE(array_agg(o.ship), '{}') AS states
+            FROM
+                o
+            WHERE
+                o.c_id = c.c_id AND o.ship != c.bill
+        ) ON true;
+----
+1  CA    {}
+3  MA    {}
+4  TX    {WY}
+5  NULL  {}
+6  FL    {WA}
+2  TX    {CA}
+
+# Customers that have billing addresses and all orders not going to their billing address
+query IT rowsort
+SELECT
+    c_id, states
+FROM
+    c
+    LEFT JOIN LATERAL (
+            SELECT
+                COALESCE(array_agg(o.ship), '{}') AS states
+            FROM
+                o
+            WHERE
+                o.c_id = c.c_id AND o.ship != c.bill
+        ) ON true
+WHERE
+    bill IS NOT NULL;
+----
+1  {}
+3  {}
+2  {CA}
+4  {WY}
+6  {WA}
diff --git a/pkg/sql/opt/optbuilder/join.go b/pkg/sql/opt/optbuilder/join.go
@@ -31,7 +31,18 @@ import (
 // return values.
 func (b *Builder) buildJoin(join *tree.JoinTableExpr, inScope *scope) (outScope *scope) {
 	leftScope := b.buildDataSource(join.Left, nil /* indexFlags */, inScope)
-	rightScope := b.buildDataSource(join.Right, nil /* indexFlags */, inScope)
+
+	isLateral := false
+	inScopeRight := inScope
+	// If this is a lateral join, use leftScope as inScope for the right side.
+	// The right side scope of a LATERAL join includes the columns produced by
+	// the left side.
+	if t, ok := join.Right.(*tree.AliasedTableExpr); ok && t.Lateral {
+		isLateral = true
+		inScopeRight = leftScope
+	}
+
+	rightScope := b.buildDataSource(join.Right, nil /* indexFlags */, inScopeRight)
 
 	// Check that the same table name is not used on both sides.
 	b.validateJoinTableNames(leftScope, rightScope)
@@ -103,7 +114,7 @@ func (b *Builder) buildJoin(join *tree.JoinTableExpr, inScope *scope) (outScope
 		left := leftScope.expr.(memo.RelExpr)
 		right := rightScope.expr.(memo.RelExpr)
 		outScope.expr = b.constructJoin(
-			joinType, left, right, filters, &memo.JoinPrivate{Flags: flags},
+			joinType, left, right, filters, &memo.JoinPrivate{Flags: flags}, isLateral,
 		)
 		return outScope
 
@@ -167,20 +178,35 @@ func (b *Builder) findJoinColsToValidate(scope *scope) util.FastIntSet {
 	return ords
 }
 
+var invalidLateralJoin = pgerror.New(pgcode.Syntax, "The combining JOIN type must be INNER or LEFT for a LATERAL reference")
+
 func (b *Builder) constructJoin(
 	joinType sqlbase.JoinType,
 	left, right memo.RelExpr,
 	on memo.FiltersExpr,
 	private *memo.JoinPrivate,
+	isLateral bool,
 ) memo.RelExpr {
 	switch joinType {
 	case sqlbase.InnerJoin:
+		if isLateral {
+			return b.factory.ConstructInnerJoinApply(left, right, on, private)
+		}
 		return b.factory.ConstructInnerJoin(left, right, on, private)
 	case sqlbase.LeftOuterJoin:
+		if isLateral {
+			return b.factory.ConstructLeftJoinApply(left, right, on, private)
+		}
 		return b.factory.ConstructLeftJoin(left, right, on, private)
 	case sqlbase.RightOuterJoin:
+		if isLateral {
+			panic(invalidLateralJoin)
+		}
 		return b.factory.ConstructRightJoin(left, right, on, private)
 	case sqlbase.FullOuterJoin:
+		if isLateral {
+			panic(invalidLateralJoin)
+		}
 		return b.factory.ConstructFullJoin(left, right, on, private)
 	default:
 		panic(pgerror.Newf(pgcode.FeatureNotSupported,
@@ -351,6 +377,7 @@ func (jb *usingJoinBuilder) finishBuild() {
 		jb.rightScope.expr.(memo.RelExpr),
 		jb.filters,
 		&memo.JoinPrivate{Flags: jb.joinFlags},
+		false /* isLateral */,
 	)
 
 	if !jb.ifNullCols.Empty() {

diff --git a/pkg/sql/opt/optbuilder/testdata/lateral b/pkg/sql/opt/optbuilder/testdata/lateral
@@ -287,3 +287,49 @@ inner-join-apply
  │                                       ├── variable: j2.j [type=jsonb]
  │                                       └── const: 'members' [type=string]
  └── filters (true)
+
+build
+SELECT * FROM x JOIN LATERAL (SELECT * FROM y WHERE b = x.a) ON true
+----
+inner-join-apply
+ ├── columns: a:1(int!null) b:2(int!null)
+ ├── scan x
+ │    └── columns: a:1(int!null)
+ ├── select
+ │    ├── columns: b:2(int!null)
+ │    ├── scan y
+ │    │    └── columns: b:2(int!null)
+ │    └── filters
+ │         └── eq [type=bool]
+ │              ├── variable: b [type=int]
+ │              └── variable: a [type=int]
+ └── filters
+      └── true [type=bool]
+
+build
+SELECT * FROM x LEFT JOIN LATERAL (SELECT * FROM y WHERE b = x.a) ON true
+----
+left-join-apply
+ ├── columns: a:1(int!null) b:2(int)
+ ├── scan x
+ │    └── columns: a:1(int!null)
+ ├── select
+ │    ├── columns: b:2(int!null)
+ │    ├── scan y
+ │    │    └── columns: b:2(int!null)
+ │    └── filters
+ │         └── eq [type=bool]
+ │              ├── variable: b [type=int]
+ │              └── variable: a [type=int]
+ └── filters
+      └── true [type=bool]
+
+build
+SELECT * FROM x RIGHT JOIN LATERAL (SELECT * FROM y WHERE b = x.a) ON true
+----
+error (42601): The combining JOIN type must be INNER or LEFT for a LATERAL reference
+
+build
+SELECT * FROM x FULL OUTER JOIN LATERAL (SELECT * FROM y WHERE b = x.a) ON true
+----
+error (42601): The combining JOIN type must be INNER or LEFT for a LATERAL reference