diff --git a/pkg/sql/logictest/testdata/logic_test/update_from b/pkg/sql/logictest/testdata/logic_test/update_from index c13cc9bc84ba..6028ca4aa528 100644 --- a/pkg/sql/logictest/testdata/logic_test/update_from +++ b/pkg/sql/logictest/testdata/logic_test/update_from @@ -168,7 +168,6 @@ a b c a b a c 1 200 300 1 200 1 300 2 300 400 2 300 2 400 - # 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 diff --git a/pkg/sql/opt/exec/execbuilder/testdata/update_from b/pkg/sql/opt/exec/execbuilder/testdata/update_from index fe2c44f4440b..807f44ad00c2 100644 --- a/pkg/sql/opt/exec/execbuilder/testdata/update_from +++ b/pkg/sql/opt/exec/execbuilder/testdata/update_from @@ -62,24 +62,23 @@ WHERE RETURNING abc.a, abc.b AS new_b, old.b as old_b, abc.c as new_c, old.c as old_c ---- -render · · - └── render · · - └── run · · - └── update · · - │ table abc - │ set b, c - │ strategy updater - └── render · · - └── merge-join · · - │ type inner - │ equality (a) = (a) - │ mergeJoinOrder +"(a=a)" - ├── scan · · - │ table abc@primary - │ spans ALL - └── scan · · -· table abc@primary -· spans ALL +render · · + └── run · · + └── update · · + │ table abc + │ set b, c + │ strategy updater + └── render · · + └── merge-join · · + │ type inner + │ equality (a) = (a) + │ mergeJoinOrder +"(a=a)" + ├── scan · · + │ table abc@primary + │ spans ALL + └── scan · · +· table abc@primary +· spans ALL # Check if RETURNING * returns everything query TTTTT @@ -184,28 +183,27 @@ WHERE RETURNING * ---- -render · · - └── run · · - └── update · · - │ table abc - │ set b, c - │ strategy updater - └── render · · - └── distinct · · - │ distinct on a - └── hash-join · · - │ type inner - │ equality (a) = (a) - ├── scan · · - │ table ac@primary - │ spans ALL - └── hash-join · · - │ type inner - │ equality (a) = (a) - │ right cols are key · - ├── scan · · - │ table ab@primary - │ spans ALL - └── scan · · -· table abc@primary -· spans ALL +run · · + └── update · · + │ table abc + │ set b, c + │ strategy updater + └── render · · + └── distinct · · + │ distinct on a + └── hash-join · · + │ type inner + │ equality (a) = (a) + ├── scan · · + │ table ac@primary + │ spans ALL + └── hash-join · · + │ type inner + │ equality (a) = (a) + │ right cols are key · + ├── scan · · + │ table ab@primary + │ spans ALL + └── scan · · +· table abc@primary +· spans ALL diff --git a/pkg/sql/opt/norm/prune_cols.go b/pkg/sql/opt/norm/prune_cols.go index 0c41e9f96533..9b54d89424bc 100644 --- a/pkg/sql/opt/norm/prune_cols.go +++ b/pkg/sql/opt/norm/prune_cols.go @@ -546,6 +546,12 @@ func (c *CustomFuncs) CanPruneMutationReturnCols( } } + for _, passthroughCol := range private.PassthroughCols { + if passthroughCol != 0 && !needed.Contains(passthroughCol) { + return true + } + } + return false } @@ -558,14 +564,24 @@ func (c *CustomFuncs) PruneMutationReturnCols( ) *memo.MutationPrivate { newPrivate := *private newReturnCols := make(opt.ColList, len(private.ReturnCols)) + newPassthroughCols := make(opt.ColList, 0, len(private.PassthroughCols)) tabID := c.mem.Metadata().TableMeta(private.Table).MetaID + // Prune away the ReturnCols that are unused. for i := range private.ReturnCols { if needed.Contains(tabID.ColumnID(i)) { newReturnCols[i] = private.ReturnCols[i] } } + // Prune away the PassthroughCols that are unused. + for _, passthroughCol := range private.PassthroughCols { + if passthroughCol != 0 && needed.Contains(passthroughCol) { + newPassthroughCols = append(newPassthroughCols, passthroughCol) + } + } + newPrivate.ReturnCols = newReturnCols + newPrivate.PassthroughCols = newPassthroughCols return &newPrivate } diff --git a/pkg/sql/opt/norm/testdata/rules/prune_cols b/pkg/sql/opt/norm/testdata/rules/prune_cols index d2c938f82151..0197af6113fa 100644 --- a/pkg/sql/opt/norm/testdata/rules/prune_cols +++ b/pkg/sql/opt/norm/testdata/rules/prune_cols @@ -2640,3 +2640,42 @@ project └── projections ├── CASE WHEN rowid IS NULL THEN column12 ELSE d END [type=int, outer=(12,17,21)] └── CASE WHEN rowid IS NULL THEN column13 ELSE rowid END [type=int, outer=(13,21)] + +# Make sure the passthrough columns of an UPDATE ... FROM query are pruned. +opt +UPDATE abcde +SET + b=family.b, c = family.c +FROM + family +WHERE + abcde.a=family.a +RETURNING + abcde.a, family.b, family.c +---- +update abcde + ├── columns: a:1(int!null) b:12(int) c:13(int) + ├── fetch columns: abcde.a:6(int) abcde.b:7(int) abcde.c:8(int) abcde.d:9(int) abcde.e:10(int) + ├── update-mapping: + │ ├── "family".b:12 => abcde.b:2 + │ └── "family".c:13 => abcde.c:3 + ├── side-effects, mutations + ├── key: (1) + ├── fd: (1)-->(12,13) + └── inner-join (merge) + ├── columns: abcde.a:6(int!null) abcde.b:7(int) abcde.c:8(int) abcde.d:9(int) abcde.e:10(int) "family".a:11(int!null) "family".b:12(int) "family".c:13(int) + ├── left ordering: +6 + ├── right ordering: +11 + ├── key: (11) + ├── fd: (6)-->(7-10), (7,8)~~>(6,9,10), (11)-->(12,13), (6)==(11), (11)==(6) + ├── scan abcde + │ ├── columns: abcde.a:6(int!null) abcde.b:7(int) abcde.c:8(int) abcde.d:9(int) abcde.e:10(int) + │ ├── key: (6) + │ ├── fd: (6)-->(7-10), (7,8)~~>(6,9,10) + │ └── ordering: +6 + ├── scan "family" + │ ├── columns: "family".a:11(int!null) "family".b:12(int) "family".c:13(int) + │ ├── key: (11) + │ ├── fd: (11)-->(12,13) + │ └── ordering: +11 + └── filters (true) diff --git a/pkg/sql/opt/optbuilder/testdata/update_from b/pkg/sql/opt/optbuilder/testdata/update_from index 6da36cbff868..c370ddfcf3e6 100644 --- a/pkg/sql/opt/optbuilder/testdata/update_from +++ b/pkg/sql/opt/optbuilder/testdata/update_from @@ -86,34 +86,32 @@ WHERE RETURNING abc.a, abc.b AS new_b, old.b as old_b, abc.c as new_c, old.c as old_c ---- -project +update abc ├── columns: a:1(int!null) new_b:2(int) old_b:8(int) new_c:3(int) old_c:9(int) - └── update abc - ├── columns: abc.a:1(int!null) abc.b:2(int) abc.c:3(int) old.a:7(int) old.b:8(int) old.c:9(int) - ├── fetch columns: abc.a:4(int) abc.b:5(int) abc.c:6(int) - ├── update-mapping: - │ ├── column10:10 => abc.b:2 - │ └── column11:11 => abc.c:3 - └── project - ├── columns: column10:10(int) column11:11(int) abc.a:4(int!null) abc.b:5(int) abc.c:6(int) old.a:7(int!null) old.b:8(int) old.c:9(int) - ├── inner-join (merge) - │ ├── columns: abc.a:4(int!null) abc.b:5(int) abc.c:6(int) old.a:7(int!null) old.b:8(int) old.c:9(int) - │ ├── left ordering: +4 - │ ├── right ordering: +7 - │ ├── scan abc - │ │ ├── columns: abc.a:4(int!null) abc.b:5(int) abc.c:6(int) - │ │ └── ordering: +4 - │ ├── scan old - │ │ ├── columns: old.a:7(int!null) old.b:8(int) old.c:9(int) - │ │ └── ordering: +7 - │ └── filters (true) - └── projections - ├── plus [type=int] - │ ├── variable: old.b [type=int] - │ └── const: 1 [type=int] - └── plus [type=int] - ├── variable: old.c [type=int] - └── const: 2 [type=int] + ├── fetch columns: abc.a:4(int) abc.b:5(int) abc.c:6(int) + ├── update-mapping: + │ ├── column10:10 => abc.b:2 + │ └── column11:11 => abc.c:3 + └── project + ├── columns: column10:10(int) column11:11(int) abc.a:4(int!null) abc.b:5(int) abc.c:6(int) old.b:8(int) old.c:9(int) + ├── inner-join (merge) + │ ├── columns: abc.a:4(int!null) abc.b:5(int) abc.c:6(int) old.a:7(int!null) old.b:8(int) old.c:9(int) + │ ├── left ordering: +4 + │ ├── right ordering: +7 + │ ├── scan abc + │ │ ├── columns: abc.a:4(int!null) abc.b:5(int) abc.c:6(int) + │ │ └── ordering: +4 + │ ├── scan old + │ │ ├── columns: old.a:7(int!null) old.b:8(int) old.c:9(int) + │ │ └── ordering: +7 + │ └── filters (true) + └── projections + ├── plus [type=int] + │ ├── variable: old.b [type=int] + │ └── const: 1 [type=int] + └── plus [type=int] + ├── variable: old.c [type=int] + └── const: 2 [type=int] # Check if RETURNING * returns everything opt @@ -294,50 +292,46 @@ WHERE RETURNING * ---- -project +update abc ├── columns: a:1(int!null) b:2(int) c:3(int) a:7(int) b:8(int) a:10(int) c:11(int) - └── update abc - ├── columns: abc.a:1(int!null) abc.b:2(int) abc.c:3(int) ab.a:7(int) ab.b:8(int) ab.rowid:9(int) ac.a:10(int) ac.c:11(int) - ├── fetch columns: abc.a:4(int) abc.b:5(int) abc.c:6(int) - ├── update-mapping: - │ ├── ab.b:8 => abc.b:2 - │ └── ac.c:11 => abc.c:3 - └── distinct-on - ├── columns: abc.a:4(int!null) abc.b:5(int) abc.c:6(int) ab.a:7(int) ab.b:8(int) ab.rowid:9(int) ac.a:10(int) ac.c:11(int) - ├── grouping columns: abc.a:4(int!null) - ├── inner-join (hash) - │ ├── columns: abc.a:4(int!null) abc.b:5(int) abc.c:6(int) ab.a:7(int!null) ab.b:8(int) ab.rowid:9(int!null) ac.a:10(int!null) ac.c:11(int) - │ ├── scan ac - │ │ └── columns: ac.a:10(int) ac.c:11(int) - │ ├── inner-join (hash) - │ │ ├── columns: abc.a:4(int!null) abc.b:5(int) abc.c:6(int) ab.a:7(int!null) ab.b:8(int) ab.rowid:9(int!null) - │ │ ├── scan ab - │ │ │ └── columns: ab.a:7(int) ab.b:8(int) ab.rowid:9(int!null) - │ │ ├── scan abc - │ │ │ └── columns: abc.a:4(int!null) abc.b:5(int) abc.c:6(int) - │ │ └── filters - │ │ └── eq [type=bool] - │ │ ├── variable: abc.a [type=int] - │ │ └── variable: ab.a [type=int] - │ └── filters - │ └── eq [type=bool] - │ ├── variable: ab.a [type=int] - │ └── variable: ac.a [type=int] - └── aggregations - ├── first-agg [type=int] - │ └── variable: abc.b [type=int] - ├── first-agg [type=int] - │ └── variable: abc.c [type=int] - ├── first-agg [type=int] - │ └── variable: ab.a [type=int] - ├── first-agg [type=int] - │ └── variable: ab.b [type=int] - ├── first-agg [type=int] - │ └── variable: ab.rowid [type=int] - ├── first-agg [type=int] - │ └── variable: ac.a [type=int] - └── first-agg [type=int] - └── variable: ac.c [type=int] + ├── fetch columns: abc.a:4(int) abc.b:5(int) abc.c:6(int) + ├── update-mapping: + │ ├── ab.b:8 => abc.b:2 + │ └── ac.c:11 => abc.c:3 + └── distinct-on + ├── columns: abc.a:4(int!null) abc.b:5(int) abc.c:6(int) ab.a:7(int) ab.b:8(int) ac.a:10(int) ac.c:11(int) + ├── grouping columns: abc.a:4(int!null) + ├── inner-join (hash) + │ ├── columns: abc.a:4(int!null) abc.b:5(int) abc.c:6(int) ab.a:7(int!null) ab.b:8(int) ac.a:10(int!null) ac.c:11(int) + │ ├── scan ac + │ │ └── columns: ac.a:10(int) ac.c:11(int) + │ ├── inner-join (hash) + │ │ ├── columns: abc.a:4(int!null) abc.b:5(int) abc.c:6(int) ab.a:7(int!null) ab.b:8(int) + │ │ ├── scan ab + │ │ │ └── columns: ab.a:7(int) ab.b:8(int) + │ │ ├── scan abc + │ │ │ └── columns: abc.a:4(int!null) abc.b:5(int) abc.c:6(int) + │ │ └── filters + │ │ └── eq [type=bool] + │ │ ├── variable: abc.a [type=int] + │ │ └── variable: ab.a [type=int] + │ └── filters + │ └── eq [type=bool] + │ ├── variable: ab.a [type=int] + │ └── variable: ac.a [type=int] + └── aggregations + ├── first-agg [type=int] + │ └── variable: abc.b [type=int] + ├── first-agg [type=int] + │ └── variable: abc.c [type=int] + ├── first-agg [type=int] + │ └── variable: ab.a [type=int] + ├── first-agg [type=int] + │ └── variable: ab.b [type=int] + ├── first-agg [type=int] + │ └── variable: ac.a [type=int] + └── first-agg [type=int] + └── variable: ac.c [type=int] # Make sure UPDATE ... FROM can return hidden columns. opt