opt, sql: support UPDATE ... FROM statements

[ridwanmsharif]

Addresses #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.

Release note: None

[cockroach-teamcity]

This change is 

[jordanlewis]

Awesome!

[rytaft]

Nice work!

Reviewed 21 of 21 files at r1.
Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @justinj, @RaduBerinde, and @ridwanmsharif)

---

pkg/sql/logictest/testdata/logic_test/update_from, line 1 at r1 (raw file):

# LogicTest: local-opt

what about also adding fakedist-opt?

---

pkg/sql/logictest/testdata/logic_test/update_from, line 24 at r1 (raw file):

statement ok
INSERT INTO new_abc VALUES (1, 2, 3), (2, 3, 4)

[nit] might help to use different values for b and c so it's easier to distinguish from the original table

---

pkg/sql/logictest/testdata/logic_test/update_from, line 40 at r1 (raw file):

statement ok
UPDATE abc SET b = other.b, c = other.c FROM new_abc AS other WHERE abc.a = other.a

what's the correct behavior in this case? It just uses the first matching row?

---

pkg/sql/opt/exec/execbuilder/mutation.go, line 109 at r1 (raw file):

	// TODO(andyk): Using ensureColumns here can result in an extra Render.
	// Upgrade execution engine to not require this.
	colList := make(opt.ColList, 0, len(upd.FetchCols)+len(upd.UpdateCols)+len(upd.CheckCols))

should you also add capacity for PassthroughCols?

---

pkg/sql/opt/ops/mutation.opt, line 125 at r1 (raw file):

    # PassthroughCols are columns that the mutation needs to passthrough from
    # its input. Its similar to the passthrough columns in projections. This

[nit] Its similar -> It's similar

---

pkg/sql/opt/optbuilder/delete.go, line 93 at r1 (raw file):

	mb.outScope.expr = mb.b.factory.ConstructDelete(mb.outScope.expr, mb.checks, private)

	mb.buildReturning(returning, nil)

nil -> nil /* fromCols */

---

pkg/sql/opt/optbuilder/insert.go, line 623 at r1 (raw file):

	mb.outScope.expr = mb.b.factory.ConstructInsert(mb.outScope.expr, mb.checks, private)

	mb.buildReturning(returning, nil)

nil -> nil /* fromCols */

---

pkg/sql/opt/optbuilder/insert.go, line 874 at r1 (raw file):

	mb.outScope.expr = mb.b.factory.ConstructUpsert(mb.outScope.expr, mb.checks, private)

	mb.buildReturning(returning, nil)

nil -> nil /* fromCols */

---

pkg/sql/opt/optbuilder/mutation_builder.go, line 192 at r1 (raw file):

// clause are returned.
// TODO(andyk): Do needed column analysis to project fewer columns if possible.
func (mb *mutationBuilder) buildInputForUpdate(

Can you extract some of the duplicated code into helper functions?

---

pkg/sql/opt/optbuilder/mutation_builder.go, line 294 at r1 (raw file):

	// possible to remap columns, as in this example:
	//
	//   UPDATE abc SET a=b

this example should be changed to DELETE

---

pkg/sql/opt/xform/optimizer.go, line 217 at r1 (raw file):

	// Validate there are no dangling references.
	if !root.Relational().OuterCols.Empty() && false {

I assume this was for debugging?

[RaduBerinde]

Very nice! Thanks for taking care of this so quickly! overall, just some minor comments.

Reviewable status: complete! 1 of 0 LGTMs obtained (waiting on @justinj, @ridwanmsharif, and @rytaft)

---

pkg/sql/logictest/testdata/logic_test/update_from, line 24 at r1 (raw file):

Previously, rytaft (Rebecca Taft) wrote…

[nit] might help to use different values for b and c so it's easier to distinguish from the original table

Also, we should have some rows in abc that don't match and aren't updated.

---

pkg/sql/logictest/testdata/logic_test/update_from, line 40 at r1 (raw file):

Previously, rytaft (Rebecca Taft) wrote…

what's the correct behavior in this case? It just uses the first matching row?

I'd add "an arbitrary row from each group is chosen (this is consistent with Postgres)" to the comment above

---

pkg/sql/logictest/testdata/logic_test/update_from, line 2 at r2 (raw file):

# LogicTest: local-opt

Could you also add some corresponding execbuilder tests that show the EXPLAIN for the statements in this test?

---

pkg/sql/logictest/testdata/logic_test/update_from, line 27 at r2 (raw file):

statement ok
UPDATE abc SET b = other.b, c = other.c FROM new_abc AS other WHERE abc.a = other.a

AS other is not necessary, and makes it harder to read

---

pkg/sql/logictest/testdata/logic_test/update_from, line 50 at r2 (raw file):

# 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

[nit] break this up into multiple lines

---

pkg/sql/logictest/testdata/logic_test/update_from, line 57 at r2 (raw file):

# Check if RETURNING * returns everything
query IIIIII

[nit] add colnames so the output is more readable

---

pkg/sql/logictest/testdata/logic_test/update_from, line 61 at r2 (raw file):

----
1  4  7  1  3  5
2  5  8  2  4  6

I'd add a test or two with multiple tables in the FROM list

---

pkg/sql/opt/exec/factory.go, line 323 at r2 (raw file):

	// columns in the same order as they appear in the table schema, with the
	// fetch columns first and the update columns second. The rowsNeeded parameter
	// is true if a RETURNING clause needs the updated row(s) as output.

Add an explanation of passthrough

---

pkg/sql/opt/optbuilder/mutation_builder.go, line 201 at r2 (raw file):

	//

	// FROM

this comment is confusing. I think it made sense for DELETE but not UPDATE

---

pkg/sql/opt/optbuilder/mutation_builder.go, line 753 at r2 (raw file):

// buildReturning wraps the input expression with a Project operator that
// projects the given RETURNING expressions.
func (mb *mutationBuilder) buildReturning(returning tree.ReturningExprs, fromCols []scopeColumn) {

[nit] document fromCols. Wondering if it makes more sense to name them extraAccessibleCols here

---

pkg/sql/opt/optbuilder/testdata/update_from, line 4 at r2 (raw file):

CREATE TABLE abc (a int primary key, b int, c int)
----

We need to test that the expressions inside the FROM list can't refer to columns from the table.

---

pkg/sql/opt/optbuilder/testdata/update_from, line 77 at r2 (raw file):

# Check if UPDATE FROM works well with RETURNING expressions that reference the FROM tables.
opt
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

[nit] break this up into multiple lines

[justinj]

This is awesome, great job!!

high-level nit: I think it would be good if the change to support passthrough cols was a separate commit from the change to add UPDATE FROM, since they're conceptually pretty different, even though one depends on the other (if this is too much trouble, don't worry about it)

Reviewable status: complete! 0 of 0 LGTMs obtained (and 1 stale) (waiting on @justinj, @RaduBerinde, @ridwanmsharif, and @rytaft)

---

pkg/sql/logictest/testdata/logic_test/update_from, line 24 at r1 (raw file):

Previously, RaduBerinde wrote…

Also, we should have some rows in abc that don't match and aren't updated.

I'm a big fan of making the different columns have different orders of magnitude, so it's easy to tell them apart (ex: (1, 20, 300))

---

pkg/sql/logictest/testdata/logic_test/update_from, line 40 at r2 (raw file):

statement ok
UPDATE abc SET b = other.b, c = other.c FROM new_abc AS other WHERE abc.a = other.a

is LATERAL valid here? If yes, should it be tested?

---

pkg/sql/opt/optbuilder/mutation_builder.go, line 189 at r2 (raw file):

//
// All columns from the table to update are added to fetchColList.
// If a FROM clause is defined, the columns in the FROM

nit: this comment could be a little more descriptive, like where the columns come from, how the data source is constructed, etc.

---

pkg/sql/opt/optbuilder/mutation_builder.go, line 213 at r2 (raw file):

	numCols := len(mb.outScope.cols)

	// If there is a from clause present, we must join all the tables together with the

nit: FROM

---

pkg/sql/opt/optbuilder/update.go, line 106 at r2 (raw file):

	//
	// All columns from the update table will be projected.
	fromCols := mb.buildInputForUpdate(inScope, upd.From, upd.Where, upd.Limit, upd.OrderBy)

can this be accessed via mb.outScope.cols?

---

pkg/sql/opt/optbuilder/testdata/update_from, line 77 at r2 (raw file):

Previously, RaduBerinde wrote…

[nit] break this up into multiple lines

I generally use vim to edit these test files, and I have this in my vimrc to sqlfmt highlighted queries:

noremap ,s :!cockroach sqlfmt --print-width 80 --use-spaces<cr><cr>

---

pkg/sql/parser/sql.y, line 5526 at r2 (raw file):

update_stmt:
  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

did you confirm with PG that this is the right place for this? i've gotten these ordered wrong in the past

---

pkg/sql/sem/tree/pretty.go, line 1027 at r2 (raw file):

		p.row("UPDATE", p.Doc(node.Table)),
		p.row("SET", p.Doc(&node.Exprs)))
	if len(node.From) > 0 {

Are there tests for this? see sem/tree/testdata/pretty

[ridwanmsharif]

TFTRs!

Reviewable status: complete! 0 of 0 LGTMs obtained (and 1 stale) (waiting on @justinj, @RaduBerinde, and @rytaft)

---

pkg/sql/logictest/testdata/logic_test/update_from, line 1 at r1 (raw file):

Previously, rytaft (Rebecca Taft) wrote…

what about also adding fakedist-opt?

Done.

---

pkg/sql/logictest/testdata/logic_test/update_from, line 24 at r1 (raw file):

Previously, justinj (Justin Jaffray) wrote…

I'm a big fan of making the different columns have different orders of magnitude, so it's easy to tell them apart (ex: (1, 20, 300))

Done.

---

pkg/sql/logictest/testdata/logic_test/update_from, line 40 at r1 (raw file):

Previously, RaduBerinde wrote…

I'd add "an arbitrary row from each group is chosen (this is consistent with Postgres)" to the comment above

Done.

---

pkg/sql/logictest/testdata/logic_test/update_from, line 27 at r2 (raw file):

Previously, RaduBerinde wrote…

AS other is not necessary, and makes it harder to read

Done.

---

pkg/sql/logictest/testdata/logic_test/update_from, line 40 at r2 (raw file):

Previously, justinj (Justin Jaffray) wrote…

is LATERAL valid here? If yes, should it be tested?

Done.

---

pkg/sql/logictest/testdata/logic_test/update_from, line 50 at r2 (raw file):

Previously, RaduBerinde wrote…

[nit] break this up into multiple lines

Done.

---

pkg/sql/logictest/testdata/logic_test/update_from, line 57 at r2 (raw file):

Previously, RaduBerinde wrote…

[nit] add colnames so the output is more readable

Done.

---

pkg/sql/logictest/testdata/logic_test/update_from, line 61 at r2 (raw file):

Previously, RaduBerinde wrote…

I'd add a test or two with multiple tables in the FROM list

Done.

---

pkg/sql/opt/exec/factory.go, line 323 at r2 (raw file):

Previously, RaduBerinde wrote…

Add an explanation of passthrough

Done.

---

pkg/sql/opt/exec/execbuilder/mutation.go, line 109 at r1 (raw file):

Previously, rytaft (Rebecca Taft) wrote…

should you also add capacity for PassthroughCols?

Done.

---

pkg/sql/opt/ops/mutation.opt, line 125 at r1 (raw file):

Previously, rytaft (Rebecca Taft) wrote…

[nit] Its similar -> It's similar

Done.

---

pkg/sql/opt/optbuilder/delete.go, line 93 at r1 (raw file):

Previously, rytaft (Rebecca Taft) wrote…

nil -> nil /* fromCols */

Done.

---

pkg/sql/opt/optbuilder/insert.go, line 623 at r1 (raw file):

Previously, rytaft (Rebecca Taft) wrote…

nil -> nil /* fromCols */

Done.

---

pkg/sql/opt/optbuilder/insert.go, line 874 at r1 (raw file):

Previously, rytaft (Rebecca Taft) wrote…

nil -> nil /* fromCols */

Done.

---

pkg/sql/opt/optbuilder/mutation_builder.go, line 192 at r1 (raw file):

Previously, rytaft (Rebecca Taft) wrote…

Can you extract some of the duplicated code into helper functions?

It becomes a little ugly when I do that. Like the only part that's common is the buildWhere part and the SELECT + ORDER part. buildInputForUpdate then applies a distinct on on top of that where as buildInputForDelete doesn't.

---

pkg/sql/opt/optbuilder/mutation_builder.go, line 294 at r1 (raw file):

Previously, rytaft (Rebecca Taft) wrote…

this example should be changed to DELETE

Done.

---

pkg/sql/opt/optbuilder/mutation_builder.go, line 189 at r2 (raw file):

Previously, justinj (Justin Jaffray) wrote…

nit: this comment could be a little more descriptive, like where the columns come from, how the data source is constructed, etc.

Done.

---

pkg/sql/opt/optbuilder/mutation_builder.go, line 213 at r2 (raw file):

Previously, justinj (Justin Jaffray) wrote…

nit: FROM

Done.

---

pkg/sql/opt/optbuilder/mutation_builder.go, line 753 at r2 (raw file):

Previously, RaduBerinde wrote…

[nit] document fromCols. Wondering if it makes more sense to name them extraAccessibleCols here

That's good. Done. Naming is hard.

---

pkg/sql/opt/optbuilder/update.go, line 106 at r2 (raw file):

Previously, justinj (Justin Jaffray) wrote…

can this be accessed via mb.outScope.cols?

Well yes but we want just the FROM columns to be extracted and so we do it before they are merged with mb.outScope.cols.

---

pkg/sql/opt/optbuilder/testdata/update_from, line 4 at r2 (raw file):

Previously, RaduBerinde wrote…

We need to test that the expressions inside the FROM list can't refer to columns from the table.

Done.

---

pkg/sql/opt/optbuilder/testdata/update_from, line 77 at r2 (raw file):

Previously, justinj (Justin Jaffray) wrote…

I generally use vim to edit these test files, and I have this in my vimrc to sqlfmt highlighted queries:

noremap ,s :!cockroach sqlfmt --print-width 80 --use-spaces<cr><cr>

Done.

---

pkg/sql/opt/xform/optimizer.go, line 217 at r1 (raw file):

Previously, rytaft (Rebecca Taft) wrote…

I assume this was for debugging?

omg yes thank you. I should've caught that. I was subconsciously relying on some linter to remind me to get rid of that.

---

pkg/sql/parser/sql.y, line 5526 at r2 (raw file):

Previously, justinj (Justin Jaffray) wrote…

did you confirm with PG that this is the right place for this? i've gotten these ordered wrong in the past

Yep. Right place.

---

pkg/sql/sem/tree/pretty.go, line 1027 at r2 (raw file):

Previously, justinj (Justin Jaffray) wrote…

Are there tests for this? see sem/tree/testdata/pretty

Done.

[RaduBerinde]

Reviewable status: complete! 1 of 0 LGTMs obtained (waiting on @justinj, @ridwanmsharif, and @rytaft)

---

pkg/sql/logictest/testdata/logic_test/update_from, line 57 at r2 (raw file):

Previously, ridwanmsharif (Ridwan Sharif) wrote…

Done.

I think it needs to be colnames,rowsort

---

pkg/sql/logictest/testdata/logic_test/update_from, line 87 at r3 (raw file):

# Check if RETURNING * returns everything
query IIIIII colnames rowsort

colnames,rowsort

---

pkg/sql/logictest/testdata/logic_test/update_from, line 149 at r3 (raw file):

# Make sure UPDATE ... FROM works with LATERAL.
query IIIIIII colnames rowsort

colnames,rowsort

---

pkg/sql/logictest/testdata/logic_test/update_from, line 168 at r3 (raw file):

# Make sure the FROM clause cannot reference the target table.
statement error no data source matches prefix: abc
UPDATE abc SET a = other.a FROM (SELECT abc.a from abc as x) as other WHERE abc.a=other.a

[nit] AS

[rytaft]

What are these new files that start with 22?

Reviewed 1 of 1 files at r2, 19 of 20 files at r3, 1 of 1 files at r4.
Reviewable status: complete! 0 of 0 LGTMs obtained (and 1 stale) (waiting on @justinj and @ridwanmsharif)

---

pkg/sql/logictest/testdata/logic_test/update_from, line 118 at r4 (raw file):

# Update values of table from values expression
statement ok
UPDATE abc SET b = other.b, c = other.c FROM (values (1, 2, 3), (2, 3, 4)) as other ("a", "b", "c") WHERE abc.a = other.a

[nit] what did abc contain before this statement was executed? Might be good to just print out the value of SELECT * FROM abc before this.

---

pkg/sql/opt/optbuilder/mutation_builder.go, line 192 at r1 (raw file):

Previously, ridwanmsharif (Ridwan Sharif) wrote…

It becomes a little ugly when I do that. Like the only part that's common is the buildWhere part and the SELECT + ORDER part. buildInputForUpdate then applies a distinct on on top of that where as buildInputForDelete doesn't.

Ok!

---

pkg/sql/opt/optbuilder/mutation_builder.go, line 294 at r1 (raw file):

Previously, ridwanmsharif (Ridwan Sharif) wrote…

Done.

This isn't correct SQL anymore

---

pkg/sql/opt/optbuilder/mutation_builder.go, line 200 at r4 (raw file):

// predictable, consistent with the POSTGRES implementation).
// buildInputForUpdate returns the columns of the FROM tables so they
// can me made accessible to other parts of the query (RETURNING clause).

me -> be

---

pkg/sql/opt/optbuilder/update.go, line 329 at r4 (raw file):

	private := mb.makeMutationPrivate(returning != nil)
	for _, col := range fromCols {
		if col.id != 0 && !col.hidden {

Is it not possible for the returning clause to refer to a hidden column explicitly?

[andy-kimball]

Reviewable status: complete! 0 of 0 LGTMs obtained (and 1 stale) (waiting on @justinj and @ridwanmsharif)

---

pkg/sql/opt/optbuilder/mutation_builder.go, line 202 at r4 (raw file):

// can me made accessible to other parts of the query (RETURNING clause).
// TODO(andyk): Do needed column analysis to project fewer columns if possible.
func (mb *mutationBuilder) buildInputForUpdate(

I think you should instead make a mutationBuilder.fromCols field, since that's how mutation builder works in every other case. Why should fromCols be passed around as params/return values, when tab, tabID, alias, outScope, fetchOrds, etc. all be fields on mutationBuilder? I don't think it's any more special than any of those.

[andy-kimball]

Reviewable status: complete! 0 of 0 LGTMs obtained (and 1 stale) (waiting on @justinj and @ridwanmsharif)

---

pkg/sql/opt/ops/mutation.opt, line 129 at r4 (raw file):

    # 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

This feels awkward to me to have both ReturnCols and PassthroughCols. Why isn't it just ReturnCols? These mutation operators are already complex and hard to reason about it, and segregating these two collections makes them even tougher.

[ridwanmsharif]

TFTRs! Also the files added were the test files for the pretty formatting of UPDATE...FROM queries.

Reviewable status: complete! 0 of 0 LGTMs obtained (and 1 stale) (waiting on @andy-kimball, @justinj, and @rytaft)

---

pkg/sql/logictest/testdata/logic_test/update_from, line 118 at r4 (raw file):

Previously, rytaft (Rebecca Taft) wrote…

[nit] what did abc contain before this statement was executed? Might be good to just print out the value of SELECT * FROM abc before this.

Done.

---

pkg/sql/opt/ops/mutation.opt, line 129 at r4 (raw file):

Previously, andy-kimball (Andy Kimball) wrote…

This feels awkward to me to have both ReturnCols and PassthroughCols. Why isn't it just ReturnCols? These mutation operators are already complex and hard to reason about it, and segregating these two collections makes them even tougher.

There are a few places where we require the ReturnCols to have exactly the number of columns as columns in the table (we have 0 values for columns not returned). Also we need to communicate with exec, the passthrough columns from the input we want to fetch separately. We could have the passthrough columns be part of the ReturnCols but it would make the them even more complex to reason about (the first |table columns| columns being part of the table and the rest from the passthrough set). Keeping them separate seems cleaner and is also similar to what we do with projections, and it also makes the pruning logic easier to understand.

---

pkg/sql/opt/optbuilder/mutation_builder.go, line 200 at r4 (raw file):

Previously, rytaft (Rebecca Taft) wrote…

me -> be

Done.

---

pkg/sql/opt/optbuilder/mutation_builder.go, line 202 at r4 (raw file):

Previously, andy-kimball (Andy Kimball) wrote…

I think you should instead make a mutationBuilder.fromCols field, since that's how mutation builder works in every other case. Why should fromCols be passed around as params/return values, when tab, tabID, alias, outScope, fetchOrds, etc. all be fields on mutationBuilder? I don't think it's any more special than any of those.

Done.

---

pkg/sql/opt/optbuilder/update.go, line 329 at r4 (raw file):

Previously, rytaft (Rebecca Taft) wrote…

Is it not possible for the returning clause to refer to a hidden column explicitly?

It wasn't possible before because of the distinct on. I added a change to allow this.

[ridwanmsharif]

Reviewable status: complete! 0 of 0 LGTMs obtained (and 1 stale) (waiting on @andy-kimball, @justinj, and @rytaft)

---

pkg/sql/opt/optbuilder/distinct.go, line 115 at r7 (raw file):

	// Add the output columns.
	for i := range inScope.cols {
		if !inScope.cols[i].hidden {

I don't see why we excluded hidden columns before? Am I missing something here?

[rytaft]

nice work!

Reviewed 7 of 9 files at r5, 2 of 4 files at r6, 5 of 5 files at r7.
Reviewable status: complete! 1 of 0 LGTMs obtained (and 1 stale) (waiting on @andy-kimball, @justinj, and @RaduBerinde)

---

pkg/sql/opt/optbuilder/distinct.go, line 115 at r7 (raw file):

Previously, ridwanmsharif (Ridwan Sharif) wrote…

I don't see why we excluded hidden columns before? Am I missing something here?

I'm not sure -- I think @RaduBerinde wrote this code

[ridwanmsharif]

TFTR!

bors r+

Reviewable status: complete! 1 of 0 LGTMs obtained (and 1 stale) (waiting on @andy-kimball, @justinj, and @RaduBerinde)

---

pkg/sql/opt/optbuilder/distinct.go, line 115 at r7 (raw file):

Previously, rytaft (Rebecca Taft) wrote…

I'm not sure -- I think @RaduBerinde wrote this code

We spoke offline, it doesn't seem that there's a reason we couldn't include the hidden columns here.

[craig]

Build succeeded

• GitHub CI (Cockroach)