opt, sql: support UPDATE ... FROM statements

Addresses cockroachdb#7841.

This change adds support for `UPDATE ... FROM` statements.
The FROM clause tables are joined together with the target
table and is used as input for the update. Furthermore, the
RETURNING clause can reference any table in the FROM clause.

TODO:
 - [ ] Prune the passthrough columns in the returning clause.

Release note: None

docs/generated/sql/bnf/stmt_block.bnf

@@ -193,7 +193,7 @@ truncate_stmt ::=
 	'TRUNCATE' opt_table relation_expr_list opt_drop_behavior
 
 update_stmt ::=
-	opt_with_clause 'UPDATE' table_name_expr_opt_alias_idx 'SET' set_clause_list opt_where_clause opt_sort_clause opt_limit_clause returning_clause
+	opt_with_clause 'UPDATE' table_name_expr_opt_alias_idx 'SET' set_clause_list opt_from_list opt_where_clause opt_sort_clause opt_limit_clause returning_clause
 
 upsert_stmt ::=
 	opt_with_clause 'UPSERT' 'INTO' insert_target insert_rest returning_clause
@@ -555,6 +555,10 @@ opt_drop_behavior ::=
 set_clause_list ::=
 	( set_clause ) ( ( ',' set_clause ) )*
 
+opt_from_list ::=
+	'FROM' from_list
+	| 
+
 db_object_name ::=
 	simple_db_object_name
 	| complex_db_object_name
@@ -1185,6 +1189,9 @@ set_clause ::=
 	single_set_clause
 	| multiple_set_clause
 
+from_list ::=
+	( table_ref ) ( ( ',' table_ref ) )*
+
 simple_db_object_name ::=
 	db_object_name_component
 
@@ -1545,6 +1552,16 @@ single_set_clause ::=
 multiple_set_clause ::=
 	'(' insert_column_list ')' '=' in_expr
 
+table_ref ::=
+	relation_expr opt_index_flags opt_ordinality opt_alias_clause
+	| select_with_parens opt_ordinality opt_alias_clause
+	| 'LATERAL' select_with_parens opt_ordinality opt_alias_clause
+	| joined_table
+	| '(' joined_table ')' opt_ordinality alias_clause
+	| func_table opt_ordinality opt_alias_clause
+	| 'LATERAL' func_table opt_ordinality opt_alias_clause
+	| '[' row_source_extension_stmt ']' opt_ordinality opt_alias_clause
+
 cockroachdb_extra_type_func_name_keyword ::=
 	'FAMILY'
 
@@ -1847,16 +1864,6 @@ distinct_clause ::=
 distinct_on_clause ::=
 	'DISTINCT' 'ON' '(' expr_list ')'
 
-table_ref ::=
-	relation_expr opt_index_flags opt_ordinality opt_alias_clause
-	| select_with_parens opt_ordinality opt_alias_clause
-	| 'LATERAL' select_with_parens opt_ordinality opt_alias_clause
-	| joined_table
-	| '(' joined_table ')' opt_ordinality alias_clause
-	| func_table opt_ordinality opt_alias_clause
-	| 'LATERAL' func_table opt_ordinality opt_alias_clause
-	| '[' row_source_extension_stmt ']' opt_ordinality opt_alias_clause
-
 all_or_distinct ::=
 	'ALL'
 	| 'DISTINCT'
@@ -1865,6 +1872,39 @@ all_or_distinct ::=
 var_list ::=
 	( var_value ) ( ( ',' var_value ) )*
 
+opt_ordinality ::=
+	'WITH' 'ORDINALITY'
+	| 
+
+opt_alias_clause ::=
+	alias_clause
+	| 
+
+joined_table ::=
+	'(' joined_table ')'
+	| table_ref 'CROSS' opt_join_hint 'JOIN' table_ref
+	| table_ref join_type opt_join_hint 'JOIN' table_ref join_qual
+	| table_ref 'JOIN' table_ref join_qual
+	| table_ref 'NATURAL' join_type opt_join_hint 'JOIN' table_ref
+	| table_ref 'NATURAL' 'JOIN' table_ref
+
+alias_clause ::=
+	'AS' table_alias_name opt_column_list
+	| table_alias_name opt_column_list
+
+func_table ::=
+	func_expr_windowless
+	| 'ROWS' 'FROM' '(' rowsfrom_list ')'
+
+row_source_extension_stmt ::=
+	delete_stmt
+	| explain_stmt
+	| insert_stmt
+	| select_stmt
+	| show_stmt
+	| update_stmt
+	| upsert_stmt
+
 user_priority ::=
 	'LOW'
 	| 'NORMAL'
@@ -2060,44 +2100,31 @@ character_base ::=
 	| 'VARCHAR'
 	| 'STRING'
 
-from_list ::=
-	( table_ref ) ( ( ',' table_ref ) )*
-
 window_definition_list ::=
 	( window_definition ) ( ( ',' window_definition ) )*
 
-opt_ordinality ::=
-	'WITH' 'ORDINALITY'
-	| 
-
-opt_alias_clause ::=
-	alias_clause
+opt_join_hint ::=
+	'HASH'
+	| 'MERGE'
+	| 'LOOKUP'
 	| 
 
-joined_table ::=
-	'(' joined_table ')'
-	| table_ref 'CROSS' opt_join_hint 'JOIN' table_ref
-	| table_ref join_type opt_join_hint 'JOIN' table_ref join_qual
-	| table_ref 'JOIN' table_ref join_qual
-	| table_ref 'NATURAL' join_type opt_join_hint 'JOIN' table_ref
-	| table_ref 'NATURAL' 'JOIN' table_ref
+join_type ::=
+	'FULL' join_outer
+	| 'LEFT' join_outer
+	| 'RIGHT' join_outer
+	| 'INNER'
 
-alias_clause ::=
-	'AS' table_alias_name opt_column_list
-	| table_alias_name opt_column_list
+join_qual ::=
+	'USING' '(' name_list ')'
+	| 'ON' a_expr
 
-func_table ::=
-	func_expr_windowless
-	| 'ROWS' 'FROM' '(' rowsfrom_list ')'
+func_expr_windowless ::=
+	func_application
+	| func_expr_common_subexpr
 
-row_source_extension_stmt ::=
-	delete_stmt
-	| explain_stmt
-	| insert_stmt
-	| select_stmt
-	| show_stmt
-	| update_stmt
-	| upsert_stmt
+rowsfrom_list ::=
+	( rowsfrom_item ) ( ( ',' rowsfrom_item ) )*
 
 opt_column ::=
 	'COLUMN'
@@ -2229,28 +2256,12 @@ char_aliases ::=
 window_definition ::=
 	window_name 'AS' window_specification
 
-opt_join_hint ::=
-	'HASH'
-	| 'MERGE'
-	| 'LOOKUP'
+join_outer ::=
+	'OUTER'
 	| 
 
-join_type ::=
-	'FULL' join_outer
-	| 'LEFT' join_outer
-	| 'RIGHT' join_outer
-	| 'INNER'
-
-join_qual ::=
-	'USING' '(' name_list ')'
-	| 'ON' a_expr
-
-func_expr_windowless ::=
-	func_application
-	| func_expr_common_subexpr
-
-rowsfrom_list ::=
-	( rowsfrom_item ) ( ( ',' rowsfrom_item ) )*
+rowsfrom_item ::=
+	func_expr_windowless
 
 create_as_col_qualification_elem ::=
 	'PRIMARY' 'KEY'
@@ -2316,13 +2327,6 @@ trim_list ::=
 	| 'FROM' expr_list
 	| expr_list
 
-join_outer ::=
-	'OUTER'
-	| 
-
-rowsfrom_item ::=
-	func_expr_windowless
-
 create_as_param ::=
 	column_name
 

docs/generated/sql/bnf/update_stmt.bnf

@@ -1,2 +1,2 @@
 update_stmt ::=
-	( ( 'WITH' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) ) |  ) 'UPDATE' ( ( table_name opt_index_flags ) | ( table_name opt_index_flags ) table_alias_name | ( table_name opt_index_flags ) 'AS' table_alias_name ) 'SET' ( ( ( ( column_name '=' a_expr ) | ( '(' ( ( ( column_name ) ) ( ( ',' ( column_name ) ) )* ) ')' '=' ( '(' select_stmt ')' | ( '(' ')' | '(' ( a_expr | a_expr ',' | a_expr ',' ( ( a_expr ) ( ( ',' a_expr ) )* ) ) ')' ) ) ) ) ) ( ( ',' ( ( column_name '=' a_expr ) | ( '(' ( ( ( column_name ) ) ( ( ',' ( column_name ) ) )* ) ')' '=' ( '(' select_stmt ')' | ( '(' ')' | '(' ( a_expr | a_expr ',' | a_expr ',' ( ( a_expr ) ( ( ',' a_expr ) )* ) ) ')' ) ) ) ) ) )* ) ( ( 'WHERE' a_expr ) |  ) ( sort_clause |  ) ( limit_clause |  ) ( 'RETURNING' target_list | 'RETURNING' 'NOTHING' |  )
+	( ( 'WITH' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) ) |  ) 'UPDATE' ( ( table_name opt_index_flags ) | ( table_name opt_index_flags ) table_alias_name | ( table_name opt_index_flags ) 'AS' table_alias_name ) 'SET' ( ( ( ( column_name '=' a_expr ) | ( '(' ( ( ( column_name ) ) ( ( ',' ( column_name ) ) )* ) ')' '=' ( '(' select_stmt ')' | ( '(' ')' | '(' ( a_expr | a_expr ',' | a_expr ',' ( ( a_expr ) ( ( ',' a_expr ) )* ) ) ')' ) ) ) ) ) ( ( ',' ( ( column_name '=' a_expr ) | ( '(' ( ( ( column_name ) ) ( ( ',' ( column_name ) ) )* ) ')' '=' ( '(' select_stmt ')' | ( '(' ')' | '(' ( a_expr | a_expr ',' | a_expr ',' ( ( a_expr ) ( ( ',' a_expr ) )* ) ) ')' ) ) ) ) ) )* ) opt_from_list ( ( 'WHERE' a_expr ) |  ) ( sort_clause |  ) ( limit_clause |  ) ( 'RETURNING' target_list | 'RETURNING' 'NOTHING' |  )

pkg/sql/logictest/testdata/logic_test/update_from

@@ -0,0 +1,61 @@
+# LogicTest: local-opt
+
+statement ok
+CREATE TABLE abc (a int primary key, b int, c int)
+
+statement ok
+INSERT INTO abc VALUES (1, 2, 3), (2, 3, 4)
+
+# Updating using self join.
+statement ok
+UPDATE abc SET b = other.b + 1, c = other.c + 1 FROM abc AS other WHERE abc.a = other.a
+
+query III
+SELECT * FROM abc
+----
+1  3  4
+2  4  5
+
+# Update from another table.
+statement ok
+CREATE TABLE new_abc (a int, b int, c int)
+
+statement ok
+INSERT INTO new_abc VALUES (1, 2, 3), (2, 3, 4)
+
+statement ok
+UPDATE abc SET b = other.b, c = other.c FROM new_abc AS other WHERE abc.a = other.a
+
+query III
+SELECT * FROM abc
+----
+1  2  3
+2  3  4
+
+# Multiple matching values for a given row.
+statement ok
+INSERT INTO new_abc VALUES (1, 1, 1)
+
+statement ok
+UPDATE abc SET b = other.b, c = other.c FROM new_abc AS other WHERE abc.a = other.a
+
+query III
+SELECT * FROM abc
+----
+1  2  3
+2  3  4
+
+# Returning old values.
+query IIIII colnames
+UPDATE abc SET b = old.b + 1, c = old.c + 2 FROM abc AS old WHERE abc.a = old.a RETURNING abc.a, abc.b AS new_b, old.b as old_b, abc.c as new_c, old.c as old_c
+----
+a  new_b  old_b  new_c  old_c
+1  3      2      5      3
+2  4      3      6      4
+
+# Check if RETURNING * returns everything
+query IIIIII
+UPDATE abc SET b = old.b + 1, c = old.c + 2 FROM abc AS old WHERE abc.a = old.a RETURNING *
+----
+1  4  7  1  3  5
+2  5  8  2  4  6

pkg/sql/opt/bench/stub_factory.go

@@ -236,6 +236,7 @@ func (f *stubFactory) ConstructUpdate(
 	updateCols exec.ColumnOrdinalSet,
 	returnCols exec.ColumnOrdinalSet,
 	checks exec.CheckOrdinalSet,
+	passthrough sqlbase.ResultColumns,
 ) (exec.Node, error) {
 	return struct{}{}, nil
 }

pkg/sql/opt/exec/execbuilder/mutation.go

@@ -22,6 +22,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
 	"github.com/cockroachdb/cockroach/pkg/util"
 	"github.com/cockroachdb/errors"
 )
@@ -108,8 +109,16 @@ func (b *Builder) buildUpdate(upd *memo.UpdateExpr) (execPlan, error) {
 	colList := make(opt.ColList, 0, len(upd.FetchCols)+len(upd.UpdateCols)+len(upd.CheckCols))
 	colList = appendColsWhenPresent(colList, upd.FetchCols)
 	colList = appendColsWhenPresent(colList, upd.UpdateCols)
-	colList = appendColsWhenPresent(colList, upd.CheckCols)
 
+	// The RETURNING clause of the Update can refer to the columns
+	// in any of the FROM tables. As a result, the Update may need
+	// to passthrough those columns so the projection above can use
+	// them.
+	if upd.NeedResults() {
+		colList = appendColsWhenPresent(colList, upd.PassthroughCols)
+	}
+
+	colList = appendColsWhenPresent(colList, upd.CheckCols)
 	input, err := b.buildMutationInput(upd.Input, colList, &upd.MutationPrivate)
 	if err != nil {
 		return execPlan{}, err
@@ -122,13 +131,24 @@ func (b *Builder) buildUpdate(upd *memo.UpdateExpr) (execPlan, error) {
 	updateColOrds := ordinalSetFromColList(upd.UpdateCols)
 	returnColOrds := ordinalSetFromColList(upd.ReturnCols)
 	checkOrds := ordinalSetFromColList(upd.CheckCols)
+
+	// Construct the result columns for the passthrough set.
+	var passthroughCols sqlbase.ResultColumns
+	if upd.NeedResults() {
+		for _, passthroughCol := range upd.PassthroughCols {
+			colMeta := b.mem.Metadata().ColumnMeta(passthroughCol)
+			passthroughCols = append(passthroughCols, sqlbase.ResultColumn{Name: colMeta.Alias, Typ: colMeta.Type})
+		}
+	}
+
 	node, err := b.factory.ConstructUpdate(
 		input.root,
 		tab,
 		fetchColOrds,
 		updateColOrds,
 		returnColOrds,
 		checkOrds,
+		passthroughCols,
 	)
 	if err != nil {
 		return execPlan{}, err
@@ -354,6 +374,16 @@ func mutationOutputColMap(mutation memo.RelExpr) opt.ColMap {
 			ord++
 		}
 	}
+
+	// The output columns of the mutation will also include all
+	// columns it allowed to pass through.
+	for _, colID := range private.PassthroughCols {
+		if colID != 0 {
+			colMap.Set(int(colID), ord)
+			ord++
+		}
+	}
+
 	return colMap
 }
 

pkg/sql/opt/exec/factory.go

@@ -328,6 +328,7 @@ type Factory interface {
 		updateCols ColumnOrdinalSet,
 		returnCols ColumnOrdinalSet,
 		checks CheckOrdinalSet,
+		passthrough sqlbase.ResultColumns,
 	) (Node, error)
 
 	// ConstructUpsert creates a node that implements an INSERT..ON CONFLICT or

pkg/sql/opt/memo/logical_props_builder.go

@@ -1195,6 +1195,14 @@ func (b *logicalPropsBuilder) buildMutationProps(mutation RelExpr, rel *props.Re
 		}
 	}
 
+	// The output columns of the mutation will also include all
+	// columns it allowed to pass through.
+	for _, col := range private.PassthroughCols {
+		if col != 0 {
+			rel.OutputCols.Add(col)
+		}
+	}
+
 	// Not Null Columns
 	// ----------------
 	// A column should be marked as not-null if the target table column is not

pkg/sql/opt/norm/prune_cols.go

@@ -59,6 +59,7 @@ func (c *CustomFuncs) NeededMutationCols(private *memo.MutationPrivate) opt.ColS
 	addCols(private.UpdateCols)
 	addCols(private.CheckCols)
 	addCols(private.ReturnCols)
+	addCols(private.PassthroughCols)
 	if private.CanaryCol != 0 {
 		cols.Add(private.CanaryCol)
 	}

pkg/sql/opt/ops/mutation.opt

@@ -121,6 +121,13 @@ define MutationPrivate {
     # then ReturnCols is nil.
     ReturnCols ColList
 
+    # PassthroughCols are columns that the mutation needs to passthrough from
+    # its input. Its similar to the passthrough columns in projections. This
+    # is useful for `UPDATE .. FROM` mutations where the `RETURNING` clause
+    # references columns from tables in the `FROM` clause. When this happens
+    # the update will need to pass through those refenced columns from its input.
+    PassthroughCols ColList
+
     # Mutation operators can act similarly to a With operator: they buffer their
     # input, making it accessible to FK queries. If this is not required, WithID
     # is zero.