From 7d2a5be3127d121033ae9dcfb8f278d69affcb03 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Villemain?= <Cedric.Villemain@Data-Bene.io>
Date: Mon, 19 Aug 2024 19:20:00 +0200
Subject: [PATCH 1/6] Add more test for target list indirections

Adding similar indirection test to existing ones but using a sublink.
i.e. INSERT INTO .... SELECT and UPDATE ... SELECT are added.

Queries of the form:

```
... SET (a, b) = (SELECT '1', '2')
... SET (b, a) = (SELECT '2', '1')
```

Should do the same thing, but currently the order of the attributes in
`SET (...)` as rewriten for pushdown is only based on the physical ordering
of the attributes in the relation. This leads to several subtle problems
including situation where a DROPped then reADDed attributes will change its
placement in the attribute list.

There are maybe more tests to add in other situation where a SET
(MULTIEXPR) is possible, though some alternatives are not supported yet
by citus, for example: `(INSERT .. ON CONFLICT SET (a,b).....`
---
 .../regress/expected/distributed_types.out    |  1 +
 .../regress/expected/multi_modifications.out  | 30 +++++++++++++++++++
 .../expected/multi_mx_modifications.out       | 21 +++++++++++++
 src/test/regress/sql/distributed_types.sql    |  4 +++
 src/test/regress/sql/multi_modifications.sql  | 18 +++++++++++
 .../regress/sql/multi_mx_modifications.sql    | 15 ++++++++++
 6 files changed, 89 insertions(+)

diff --git a/src/test/regress/expected/distributed_types.out b/src/test/regress/expected/distributed_types.out
index 703614d61..68b952aae 100644
--- a/src/test/regress/expected/distributed_types.out
+++ b/src/test/regress/expected/distributed_types.out
@@ -444,6 +444,7 @@ HINT:  Use the column name to insert or update the composite type as a single va
 UPDATE domain_indirection_test SET domain_array[0].if2 = (SELECT 5);
 ERROR:  inserting or modifying composite type fields is not supported
 HINT:  Use the column name to insert or update the composite type as a single value
+HINT:  Use the column name to insert or update the composite type as a single value
 -- below are supported as we don't do any field indirection
 INSERT INTO field_indirection_test_2 (ct2_col, int_col, ct1_col)
 VALUES ('(1, "text1", 2)', 3, '(4, 5)'), ('(6, "text2", 7)', 8, '(9, 10)');
diff --git a/src/test/regress/expected/multi_modifications.out b/src/test/regress/expected/multi_modifications.out
index 93f6c8c45..1a2c01a4d 100644
--- a/src/test/regress/expected/multi_modifications.out
+++ b/src/test/regress/expected/multi_modifications.out
@@ -330,6 +330,12 @@ UPDATE limit_orders SET (kind, limit_price) = ('buy', 999) WHERE id = 246 RETURN
  246 | GM     |        30 | Mon Jul 02 16:32:15 2007 | buy  |         999
 (1 row)
 
+UPDATE limit_orders SET (kind, limit_price) = (SELECT 'buy'::order_side, 999) WHERE id = 246 RETURNING *;
+ id  | symbol | bidder_id |        placed_at         | kind | limit_price
+---------------------------------------------------------------------
+ 246 | GM     |        30 | Mon Jul 02 16:32:15 2007 | buy  |         999
+(1 row)
+
 -- Test that on unique contraint violations, we fail fast
 \set VERBOSITY terse
 INSERT INTO limit_orders VALUES (275, 'ADR', 140, '2007-07-02 16:32:15', 'sell', 43.67);
@@ -435,6 +441,8 @@ UPDATE limit_orders SET limit_price = 0.00 FROM bidders
 -- should succeed with a CTE
 WITH deleted_orders AS (INSERT INTO limit_orders VALUES (399, 'PDR', 14, '2017-07-02 16:32:15', 'sell', 43))
 UPDATE limit_orders SET symbol = 'GM';
+WITH deleted_orders AS (INSERT INTO limit_orders SELECT 400, 'PDR', 14, '2017-07-02 16:32:15', 'sell', 43)
+UPDATE limit_orders SET symbol = 'GM';
 SELECT symbol, bidder_id FROM limit_orders WHERE id = 246;
  symbol | bidder_id
 ---------------------------------------------------------------------
@@ -927,6 +935,17 @@ SELECT * FROM summary_table ORDER BY id;
   2 |       400 | 450.0000000000000000 |       |
 (2 rows)
 
+--- TODO this one is a silent corruption:
+-- UPDATE summary_table SET (average_value, min_value) =
+-- 	(SELECT avg(value), min(value) FROM raw_table WHERE id = 2)
+-- WHERE id = 2;
+SELECT * FROM summary_table ORDER BY id;
+ id | min_value |    average_value     | count | uniques
+---------------------------------------------------------------------
+  1 |           | 200.0000000000000000 |       |
+  2 |       400 | 450.0000000000000000 |       |
+(2 rows)
+
 UPDATE summary_table SET min_value = 100
 	WHERE id IN (SELECT id FROM raw_table WHERE id = 1 and value > 100) AND id = 1;
 SELECT * FROM summary_table ORDER BY id;
@@ -1103,6 +1122,17 @@ SELECT * FROM reference_summary_table ORDER BY id;
   2 |       400 | 450.0000000000000000 |       |
 (2 rows)
 
+--- TODO this one is a silent corruption:
+-- UPDATE reference_summary_table SET (average_value, min_value) =
+-- 	(SELECT avg(value), min(value) FROM reference_raw_table WHERE id = 2)
+-- WHERE id = 2;
+SELECT * FROM reference_summary_table ORDER BY id;
+ id | min_value |    average_value     | count | uniques
+---------------------------------------------------------------------
+  1 |           | 200.0000000000000000 |       |
+  2 |       400 | 450.0000000000000000 |       |
+(2 rows)
+
 -- no need partition colum equalities on reference tables
 UPDATE reference_summary_table SET (count) =
 	(SELECT id AS inner_id FROM reference_raw_table WHERE value = 500)
diff --git a/src/test/regress/expected/multi_mx_modifications.out b/src/test/regress/expected/multi_mx_modifications.out
index 9e053d3f2..b401a67e3 100644
--- a/src/test/regress/expected/multi_mx_modifications.out
+++ b/src/test/regress/expected/multi_mx_modifications.out
@@ -91,6 +91,8 @@ SET client_min_messages TO DEFAULT;
 -- commands with non-constant partition values are unsupported
 INSERT INTO limit_orders_mx VALUES (random() * 100, 'ORCL', 152, '2011-08-25 11:50:45',
 								 'sell', 0.58);
+INSERT INTO limit_orders_mx SELECT random() * 100, 'ORCL', 152, '2011-08-25 11:50:45',
+								 'sell', 0.58;
 -- values for other columns are totally fine
 INSERT INTO limit_orders_mx VALUES (2036, 'GOOG', 5634, now(), 'buy', random());
 -- commands with mutable functions in their quals
@@ -103,6 +105,22 @@ DELETE FROM limit_orders_mx WHERE id = 246 AND placed_at = current_timestamp::ti
 INSERT INTO limit_orders_mx VALUES (2037, 'GOOG', 5634, now(), 'buy', random()),
                                    (2038, 'GOOG', 5634, now(), 'buy', random()),
                                    (2039, 'GOOG', 5634, now(), 'buy', random());
+SELECT count(distinct placed_at) FROM limit_orders_mx WHERE id BETWEEN 2037 AND 2039;
+ count
+---------------------------------------------------------------------
+     1
+(1 row)
+
+INSERT INTO limit_orders_mx SELECT * FROM (VALUES
+                                   (2040, 'GOOG', 5634, now(), 'buy'::citus_mx_test_schema.order_side_mx, random()),
+                                   (2041, 'GOOG', 5634, now(), 'buy'::citus_mx_test_schema.order_side_mx, random()),
+                                   (2042, 'GOOG', 5634, now(), 'buy'::citus_mx_test_schema.order_side_mx, random()));
+SELECT count(distinct placed_at) FROM limit_orders_mx WHERE id BETWEEN 2040 AND 2042;
+ count
+---------------------------------------------------------------------
+     1
+(1 row)
+
 -- connect back to the other node
 \c - - - :worker_1_port
 -- commands containing a CTE are supported
@@ -216,9 +234,12 @@ INSERT INTO limit_orders_mx VALUES (275, 'ADR', 140, '2007-07-02 16:32:15', 'sel
 ERROR:  duplicate key value violates unique constraint "limit_orders_mx_pkey_1220093"
 -- multi shard update is supported
 UPDATE limit_orders_mx SET limit_price = 0.00;
+UPDATE limit_orders_mx SET limit_price = (SELECT 0.00);
 -- attempting to change the partition key is unsupported
 UPDATE limit_orders_mx SET id = 0 WHERE id = 246;
 ERROR:  modifying the partition value of rows is not allowed
+UPDATE limit_orders_mx SET id = (SELECT 0) WHERE id = 246;
+ERROR:  modifying the partition value of rows is not allowed
 -- UPDATEs with a FROM clause are unsupported
 UPDATE limit_orders_mx SET limit_price = 0.00 FROM bidders
 					WHERE limit_orders_mx.id = 246 AND
diff --git a/src/test/regress/sql/distributed_types.sql b/src/test/regress/sql/distributed_types.sql
index 205faf9c7..ac44457da 100644
--- a/src/test/regress/sql/distributed_types.sql
+++ b/src/test/regress/sql/distributed_types.sql
@@ -270,10 +270,14 @@ SELECT create_distributed_table('domain_indirection_test', 'f1');
 -- not supported (field indirection to underlying composite type)
 INSERT INTO domain_indirection_test (f1,f3.if1, f3.if2) VALUES (0, 1, 2);
 INSERT INTO domain_indirection_test (f1,f3.if1, f3.if2) SELECT 0, 1, 2;
+-- INSERT INTO domain_indirection_test (f1,f3.if1, f3.if2) SELECT 0, 1, 2;
 INSERT INTO domain_indirection_test (f1,f3.if1) VALUES (0, 1);
 INSERT INTO domain_indirection_test (f1,f3.if1) SELECT 0, 1;
+-- INSERT INTO domain_indirection_test (f1,f3.if1) SELECT 0, 1;
+-- ERROR:  could not find a conversion path from type 23 to 17619
 UPDATE domain_indirection_test SET domain_array[0].if2 = 5;
 UPDATE domain_indirection_test SET domain_array[0].if2 = (SELECT 5);
+UPDATE domain_indirection_test SET domain_array[0].if2 = (SELECT 5);
 
 -- below are supported as we don't do any field indirection
 INSERT INTO field_indirection_test_2 (ct2_col, int_col, ct1_col)
diff --git a/src/test/regress/sql/multi_modifications.sql b/src/test/regress/sql/multi_modifications.sql
index 2a00e7992..eeab5536f 100644
--- a/src/test/regress/sql/multi_modifications.sql
+++ b/src/test/regress/sql/multi_modifications.sql
@@ -234,6 +234,7 @@ SELECT kind, limit_price FROM limit_orders WHERE id = 246;
 
 -- multi-column UPDATE with RETURNING
 UPDATE limit_orders SET (kind, limit_price) = ('buy', 999) WHERE id = 246 RETURNING *;
+UPDATE limit_orders SET (kind, limit_price) = (SELECT 'buy'::order_side, 999) WHERE id = 246 RETURNING *;
 
 -- Test that on unique contraint violations, we fail fast
 \set VERBOSITY terse
@@ -337,6 +338,9 @@ UPDATE limit_orders SET limit_price = 0.00 FROM bidders
 WITH deleted_orders AS (INSERT INTO limit_orders VALUES (399, 'PDR', 14, '2017-07-02 16:32:15', 'sell', 43))
 UPDATE limit_orders SET symbol = 'GM';
 
+WITH deleted_orders AS (INSERT INTO limit_orders SELECT 400, 'PDR', 14, '2017-07-02 16:32:15', 'sell', 43)
+UPDATE limit_orders SET symbol = 'GM';
+
 SELECT symbol, bidder_id FROM limit_orders WHERE id = 246;
 
 -- updates referencing just a var are supported
@@ -584,6 +588,13 @@ WHERE id = 2;
 
 SELECT * FROM summary_table ORDER BY id;
 
+--- TODO this one is a silent corruption:
+-- UPDATE summary_table SET (average_value, min_value) =
+-- 	(SELECT avg(value), min(value) FROM raw_table WHERE id = 2)
+-- WHERE id = 2;
+
+SELECT * FROM summary_table ORDER BY id;
+
 UPDATE summary_table SET min_value = 100
 	WHERE id IN (SELECT id FROM raw_table WHERE id = 1 and value > 100) AND id = 1;
 
@@ -712,6 +723,13 @@ WHERE id = 2;
 
 SELECT * FROM reference_summary_table ORDER BY id;
 
+--- TODO this one is a silent corruption:
+-- UPDATE reference_summary_table SET (average_value, min_value) =
+-- 	(SELECT avg(value), min(value) FROM reference_raw_table WHERE id = 2)
+-- WHERE id = 2;
+
+SELECT * FROM reference_summary_table ORDER BY id;
+
 -- no need partition colum equalities on reference tables
 UPDATE reference_summary_table SET (count) =
 	(SELECT id AS inner_id FROM reference_raw_table WHERE value = 500)
diff --git a/src/test/regress/sql/multi_mx_modifications.sql b/src/test/regress/sql/multi_mx_modifications.sql
index 852bf3a42..185fef451 100644
--- a/src/test/regress/sql/multi_mx_modifications.sql
+++ b/src/test/regress/sql/multi_mx_modifications.sql
@@ -67,6 +67,8 @@ SET client_min_messages TO DEFAULT;
 -- commands with non-constant partition values are unsupported
 INSERT INTO limit_orders_mx VALUES (random() * 100, 'ORCL', 152, '2011-08-25 11:50:45',
 								 'sell', 0.58);
+INSERT INTO limit_orders_mx SELECT random() * 100, 'ORCL', 152, '2011-08-25 11:50:45',
+								 'sell', 0.58;
 
 -- values for other columns are totally fine
 INSERT INTO limit_orders_mx VALUES (2036, 'GOOG', 5634, now(), 'buy', random());
@@ -83,6 +85,17 @@ INSERT INTO limit_orders_mx VALUES (2037, 'GOOG', 5634, now(), 'buy', random()),
                                    (2038, 'GOOG', 5634, now(), 'buy', random()),
                                    (2039, 'GOOG', 5634, now(), 'buy', random());
 
+SELECT count(distinct placed_at) FROM limit_orders_mx WHERE id BETWEEN 2037 AND 2039;
+
+INSERT INTO limit_orders_mx SELECT * FROM (VALUES
+                                   (2040, 'GOOG', 5634, now(), 'buy'::citus_mx_test_schema.order_side_mx, random()),
+                                   (2041, 'GOOG', 5634, now(), 'buy'::citus_mx_test_schema.order_side_mx, random()),
+                                   (2042, 'GOOG', 5634, now(), 'buy'::citus_mx_test_schema.order_side_mx, random()));
+
+SELECT count(distinct placed_at) FROM limit_orders_mx WHERE id BETWEEN 2040 AND 2042;
+
+
+
 -- connect back to the other node
 \c - - - :worker_1_port
 
@@ -153,9 +166,11 @@ INSERT INTO limit_orders_mx VALUES (275, 'ADR', 140, '2007-07-02 16:32:15', 'sel
 
 -- multi shard update is supported
 UPDATE limit_orders_mx SET limit_price = 0.00;
+UPDATE limit_orders_mx SET limit_price = (SELECT 0.00);
 
 -- attempting to change the partition key is unsupported
 UPDATE limit_orders_mx SET id = 0 WHERE id = 246;
+UPDATE limit_orders_mx SET id = (SELECT 0) WHERE id = 246;
 
 -- UPDATEs with a FROM clause are unsupported
 UPDATE limit_orders_mx SET limit_price = 0.00 FROM bidders

From dfc26f9a2a6be14b22254be73ad9a3eca0c7aa12 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Villemain?= <Cedric.Villemain@Data-Bene.io>
Date: Thu, 27 Feb 2025 05:05:12 +0100
Subject: [PATCH 2/6] Fix insert planner crash

 -- TODO crash Citus
--- INSERT INTO domain_indirection_test (f1,f3.if1, f3.if2) SELECT 0, 1, 2;
+INSERT INTO domain_indirection_test (f1,f3.if1, f3.if2) SELECT 0, 1, 2;
+ERROR:  cannot plan distributed INSERT INTO ... SELECT query
+HINT:  Do not use array references and field stores on the INSERT target list.
 INSERT INTO domain_indirection_test (f1,f3.if1) VALUES (0, 1);
 ERROR:  inserting or modifying composite type fields is not supported
 HINT:  Use the column name to insert or update the composite type as a single value
 -- TODO here: not the expected ERROR
 -- INSERT INTO domain_indirection_test (f1,f3.if1) SELECT 0, 1;
 -- ERROR:  could not find a conversion path from type 23 to 17619
  UPDATE domain_indirection_test SET domain_array[0].if2 = 5;
 ERROR:  inserting or modifying composite type fields is not supported
 HINT:  Use the column name to insert or update the composite type as a single value
 -- TODO here: not the expected ERROR
---
 src/backend/distributed/planner/insert_select_planner.c | 5 +++++
 src/test/regress/sql/distributed_types.sql              | 4 ----
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c
index db5c3d4ff..f8f0b8b8c 100644
--- a/src/backend/distributed/planner/insert_select_planner.c
+++ b/src/backend/distributed/planner/insert_select_planner.c
@@ -1082,6 +1082,11 @@ ReorderInsertSelectTargetLists(Query *originalQuery, RangeTblEntry *insertRte,
 		TargetEntry *newSubqueryTargetEntry = NULL;
 		AttrNumber originalAttrNo = get_attnum(insertRelationId,
 											   oldInsertTargetEntry->resname);
+		Node *expr;
+
+		/* we need to explore the underlying expression */
+		expr = (Node *) oldInsertTargetEntry->expr;
+		expr = strip_implicit_coercions(expr);
 
 		/* we need to explore the underlying expression */
 		Node *expr = strip_implicit_coercions((Node *) oldInsertTargetEntry->expr);
diff --git a/src/test/regress/sql/distributed_types.sql b/src/test/regress/sql/distributed_types.sql
index ac44457da..205faf9c7 100644
--- a/src/test/regress/sql/distributed_types.sql
+++ b/src/test/regress/sql/distributed_types.sql
@@ -270,14 +270,10 @@ SELECT create_distributed_table('domain_indirection_test', 'f1');
 -- not supported (field indirection to underlying composite type)
 INSERT INTO domain_indirection_test (f1,f3.if1, f3.if2) VALUES (0, 1, 2);
 INSERT INTO domain_indirection_test (f1,f3.if1, f3.if2) SELECT 0, 1, 2;
--- INSERT INTO domain_indirection_test (f1,f3.if1, f3.if2) SELECT 0, 1, 2;
 INSERT INTO domain_indirection_test (f1,f3.if1) VALUES (0, 1);
 INSERT INTO domain_indirection_test (f1,f3.if1) SELECT 0, 1;
--- INSERT INTO domain_indirection_test (f1,f3.if1) SELECT 0, 1;
--- ERROR:  could not find a conversion path from type 23 to 17619
 UPDATE domain_indirection_test SET domain_array[0].if2 = 5;
 UPDATE domain_indirection_test SET domain_array[0].if2 = (SELECT 5);
-UPDATE domain_indirection_test SET domain_array[0].if2 = (SELECT 5);
 
 -- below are supported as we don't do any field indirection
 INSERT INTO field_indirection_test_2 (ct2_col, int_col, ct1_col)

From 7ceb49068e1395ed0ac3a5fc47517f10b094239f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Villemain?= <Cedric.Villemain@Data-Bene.io>
Date: Wed, 26 Feb 2025 16:04:45 +0100
Subject: [PATCH 3/6] WIP new test

    Add more testing for UPDATE with indirection

    It'll probably be interesting to add more test related to indirection
    here.
---
 src/test/regress/expected/indirections.out | 358 +++++++++++++++++++++
 src/test/regress/sql/indirections.sql      | 206 ++++++++++++
 2 files changed, 564 insertions(+)
 create mode 100644 src/test/regress/expected/indirections.out
 create mode 100644 src/test/regress/sql/indirections.sql

diff --git a/src/test/regress/expected/indirections.out b/src/test/regress/expected/indirections.out
new file mode 100644
index 000000000..a9ef5078a
--- /dev/null
+++ b/src/test/regress/expected/indirections.out
@@ -0,0 +1,358 @@
+SET citus.shard_count TO 2;
+SET citus.next_shard_id TO 750000;
+SET citus.next_placement_id TO 750000;
+CREATE SCHEMA indirections;
+SET search_path TO indirections;
+-- specific tests related to get_update_query_targetlist_def
+-- we test only queries with sublinks, like:
+-- ( ... SET (...) = (SELECT ...))
+-- Reference tables
+CREATE TABLE test_ref_indirection (
+    id bigint primary key
+  , col_bool bool , col_date date , col_int integer , col_text text
+  );
+SELECT create_reference_table('indirections.test_ref_indirection');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE test_ref_indirection_new (
+    id bigint primary key
+  , col_bool bool , col_date date , col_int integer , col_text text
+  );
+SELECT create_reference_table('indirections.test_ref_indirection_new');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- Distributed tables
+CREATE TABLE test_dist_indirection (
+    id bigint primary key
+  , col_bool bool , col_date date , col_int integer , col_text text
+  );
+SELECT create_distributed_table('indirections.test_dist_indirection', 'id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE test_dist_indirection_new (
+    id bigint primary key
+  , col_bool bool , col_date date , col_int integer , col_text text
+  );
+SELECT create_distributed_table('indirections.test_dist_indirection_new', 'id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- Local tables required ?
+-- those should work:
+INSERT INTO test_ref_indirection (id, col_bool, col_date, col_int, col_text)
+    SELECT 1, true, '1970-01-01'::date, 1, 'one';
+INSERT INTO test_dist_indirection (id, col_bool, col_date, col_int, col_text)
+    SELECT 1, true, '1970-01-01'::date, 1, 'one';
+INSERT INTO test_ref_indirection (id, col_text, col_bool, col_date, col_int)
+    SELECT 2, 'two', false, '1970-01-01'::date, 2;
+INSERT INTO test_dist_indirection (id, col_text, col_bool, col_date, col_int)
+    SELECT 2, 'two', false, '1970-01-01'::date, 2;
+INSERT INTO test_ref_indirection SELECT 3, false, '1970-01-01'::date, 0, 'empty';
+INSERT INTO test_dist_indirection SELECT 3, false, '1970-01-01'::date, 0, 'empty';
+INSERT INTO test_ref_indirection SELECT 4, false, '1970-01-01'::date, 0, 'empty';
+INSERT INTO test_dist_indirection SELECT 4, false, '1970-01-01'::date, 0, 'empty';
+INSERT INTO test_ref_indirection_new SELECT * FROM test_ref_indirection;
+INSERT INTO test_dist_indirection_new SELECT * FROM test_dist_indirection;
+SELECT * FROM test_ref_indirection ORDER BY id;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1970 |       1 | one
+  2 | f        | 01-01-1970 |       2 | two
+  3 | f        | 01-01-1970 |       0 | empty
+  4 | f        | 01-01-1970 |       0 | empty
+(4 rows)
+
+SELECT * FROM test_dist_indirection ORDER BY id;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1970 |       1 | one
+  2 | f        | 01-01-1970 |       2 | two
+  3 | f        | 01-01-1970 |       0 | empty
+  4 | f        | 01-01-1970 |       0 | empty
+(4 rows)
+
+SELECT * FROM test_ref_indirection_new ORDER BY id;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1970 |       1 | one
+  2 | f        | 01-01-1970 |       2 | two
+  3 | f        | 01-01-1970 |       0 | empty
+  4 | f        | 01-01-1970 |       0 | empty
+(4 rows)
+
+SELECT * FROM test_dist_indirection_new ORDER BY id;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1970 |       1 | one
+  2 | f        | 01-01-1970 |       2 | two
+  3 | f        | 01-01-1970 |       0 | empty
+  4 | f        | 01-01-1970 |       0 | empty
+(4 rows)
+
+-- now UPDATEs
+UPDATE test_ref_indirection
+    SET (col_bool, col_date, col_int, col_text)
+        = (SELECT true, '1970-12-31'::date, 1, 'ok')
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 12-31-1970 |       1 | ok
+  2 | t        | 12-31-1970 |       1 | ok
+  3 | t        | 12-31-1970 |       1 | ok
+  4 | t        | 12-31-1970 |       1 | ok
+(4 rows)
+
+UPDATE test_dist_indirection
+    SET (col_bool, col_date, col_int, col_text)
+        = (SELECT true, '1970-12-31'::date, 1, 'ok')
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 12-31-1970 |       1 | ok
+  2 | t        | 12-31-1970 |       1 | ok
+  3 | t        | 12-31-1970 |       1 | ok
+  4 | t        | 12-31-1970 |       1 | ok
+(4 rows)
+
+UPDATE test_ref_indirection
+    SET (col_bool, col_date)  = (select true, '1970-06-06'::date)
+      , (col_int, col_text) = (select 1, 'still ok')
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 06-06-1970 |       1 | still ok
+  2 | t        | 06-06-1970 |       1 | still ok
+  3 | t        | 06-06-1970 |       1 | still ok
+  4 | t        | 06-06-1970 |       1 | still ok
+(4 rows)
+
+UPDATE test_dist_indirection
+    SET (col_bool, col_date)  = (select true, '1970-06-06'::date)
+      , (col_int, col_text) = (select 1, 'still ok')
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 06-06-1970 |       1 | still ok
+  2 | t        | 06-06-1970 |       1 | still ok
+  3 | t        | 06-06-1970 |       1 | still ok
+  4 | t        | 06-06-1970 |       1 | still ok
+(4 rows)
+
+SELECT * FROM test_ref_indirection ORDER BY id;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 06-06-1970 |       1 | still ok
+  2 | t        | 06-06-1970 |       1 | still ok
+  3 | t        | 06-06-1970 |       1 | still ok
+  4 | t        | 06-06-1970 |       1 | still ok
+(4 rows)
+
+SELECT * FROM test_dist_indirection ORDER BY id;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 06-06-1970 |       1 | still ok
+  2 | t        | 06-06-1970 |       1 | still ok
+  3 | t        | 06-06-1970 |       1 | still ok
+  4 | t        | 06-06-1970 |       1 | still ok
+(4 rows)
+
+-- but those should not:
+-- TODO wrong ERROR
+UPDATE test_ref_indirection
+    SET (col_date, col_text, col_int, col_bool)
+        = (SELECT '1970-12-31'::date, 'not ok', 2, false)
+RETURNING *;
+ERROR:  column "col_bool" is of type boolean but expression is of type date
+HINT:  You will need to rewrite or cast the expression.
+CONTEXT:  while executing command on localhost:xxxxx
+-- TODO wrong ERROR
+UPDATE test_dist_indirection
+    SET (col_date, col_text, col_int, col_bool)
+        = (SELECT '1970-12-31'::date, 'not ok', 2, false)
+RETURNING *;
+ERROR:  column "col_bool" is of type boolean but expression is of type date
+HINT:  You will need to rewrite or cast the expression.
+CONTEXT:  while executing command on localhost:xxxxx
+-- TODO wrong ERROR
+UPDATE test_ref_indirection
+    SET (col_int, col_date)  = (select 2, '1970-06-06'::date)
+      , (col_text, col_bool) = (select 'not ok', false)
+RETURNING *;
+ERROR:  column "col_bool" is of type boolean but expression is of type integer
+HINT:  You will need to rewrite or cast the expression.
+CONTEXT:  while executing command on localhost:xxxxx
+-- TODO wrong ERROR
+UPDATE test_dist_indirection
+    SET (col_int, col_date)  = (select 2, '1970-06-06'::date)
+      , (col_text, col_bool) = (select 'not ok', false)
+RETURNING *;
+ERROR:  column "col_bool" is of type boolean but expression is of type integer
+HINT:  You will need to rewrite or cast the expression.
+CONTEXT:  while executing command on localhost:xxxxx
+--
+-- more restrictive ones, just in case we miss a wrong value
+--
+-- those should work
+UPDATE test_ref_indirection
+    SET (col_bool, col_text) = (SELECT true, 'ok')
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 06-06-1970 |       1 | ok
+  2 | t        | 06-06-1970 |       1 | ok
+  3 | t        | 06-06-1970 |       1 | ok
+  4 | t        | 06-06-1970 |       1 | ok
+(4 rows)
+
+UPDATE test_dist_indirection
+    SET (col_bool, col_text) = (SELECT true, 'ok')
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 06-06-1970 |       1 | ok
+  2 | t        | 06-06-1970 |       1 | ok
+  3 | t        | 06-06-1970 |       1 | ok
+  4 | t        | 06-06-1970 |       1 | ok
+(4 rows)
+
+UPDATE test_ref_indirection
+    SET (col_bool, col_text) = (SELECT true, 'ok')
+WHERE id = 1
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 06-06-1970 |       1 | ok
+(1 row)
+
+UPDATE test_dist_indirection
+    SET (col_bool, col_text) = (SELECT true, 'ok')
+WHERE id = 1
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 06-06-1970 |       1 | ok
+(1 row)
+
+SELECT * FROM test_ref_indirection ORDER BY id;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 06-06-1970 |       1 | ok
+  2 | t        | 06-06-1970 |       1 | ok
+  3 | t        | 06-06-1970 |       1 | ok
+  4 | t        | 06-06-1970 |       1 | ok
+(4 rows)
+
+SELECT * FROM test_dist_indirection ORDER BY id;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 06-06-1970 |       1 | ok
+  2 | t        | 06-06-1970 |       1 | ok
+  3 | t        | 06-06-1970 |       1 | ok
+  4 | t        | 06-06-1970 |       1 | ok
+(4 rows)
+
+-- those should not
+-- TODO wrong ERROR
+UPDATE test_ref_indirection
+    SET (col_text, col_bool) = (SELECT 'not ok', false)
+RETURNING *;
+ERROR:  column "col_bool" is of type boolean but expression is of type text
+HINT:  You will need to rewrite or cast the expression.
+CONTEXT:  while executing command on localhost:xxxxx
+-- TODO wrong ERROR
+UPDATE test_dist_indirection
+    SET (col_text, col_bool) = (SELECT 'not ok', false)
+RETURNING *;
+ERROR:  column "col_bool" is of type boolean but expression is of type text
+HINT:  You will need to rewrite or cast the expression.
+CONTEXT:  while executing command on localhost:xxxxx
+-- TODO wrong ERROR
+UPDATE test_ref_indirection
+    SET (col_text, col_bool) = (SELECT 'not ok', false)
+WHERE id = 2
+RETURNING *;
+ERROR:  column "col_bool" is of type boolean but expression is of type text
+HINT:  You will need to rewrite or cast the expression.
+CONTEXT:  while executing command on localhost:xxxxx
+-- TODO wrong ERROR
+UPDATE test_dist_indirection
+    SET (col_text, col_bool) = (SELECT 'not ok', false)
+WHERE id = 2
+RETURNING *;
+ERROR:  column "col_bool" is of type boolean but expression is of type text
+HINT:  You will need to rewrite or cast the expression.
+CONTEXT:  while executing command on localhost:xxxxx
+-- several updates in CTE shoult not work
+-- TODO wrong ERROR
+with qq3 as (
+    update test_ref_indirection
+    SET (col_text, col_bool)
+      = (SELECT 'full', true)
+    where id = 3
+    returning *
+),
+qq4 as (
+    update test_ref_indirection
+    SET (col_text, col_bool)
+      = (SELECT 'fully', true)
+    where id = 4
+    returning *
+)
+select * from qq3 union all select * from qq4;
+ERROR:  column "col_bool" is of type boolean but expression is of type text
+HINT:  You will need to rewrite or cast the expression.
+CONTEXT:  while executing command on localhost:xxxxx
+-- TODO wrong ERROR
+with qq3 as (
+    update test_dist_indirection
+    SET (col_text, col_bool)
+      = (SELECT 'full', true)
+    where id = 3
+    returning *
+),
+qq4 as (
+    update test_dist_indirection
+    SET (col_text, col_bool)
+      = (SELECT 'fully', true)
+    where id = 4
+    returning *
+)
+select * from qq3 union all select * from qq4;
+ERROR:  column "col_bool" is of type boolean but expression is of type text
+HINT:  You will need to rewrite or cast the expression.
+CONTEXT:  while executing command on localhost:xxxxx
+DROP TABLE test_dist_indirection;
+DROP TABLE test_dist_indirection_new;
+DROP TABLE test_ref_indirection;
+DROP TABLE test_ref_indirection_new;
+-- https://github.com/citusdata/citus/issues/4092
+CREATE TABLE update_test (
+    a   INT DEFAULT 10,
+    b   INT,
+    c   TEXT
+);
+SELECT create_reference_table('indirections.update_test');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+UPDATE update_test
+SET (b,a) = (select a,b from update_test where b = 41 and c = 'car')
+WHERE a = 100 AND b = 20;
+-- https://github.com/citusdata/citus/pull/5692
+DROP SCHEMA indirections CASCADE;
+NOTICE:  drop cascades to 2 other objects
+DETAIL:  drop cascades to table update_test
+drop cascades to table update_test_750006
diff --git a/src/test/regress/sql/indirections.sql b/src/test/regress/sql/indirections.sql
new file mode 100644
index 000000000..fd8119ff5
--- /dev/null
+++ b/src/test/regress/sql/indirections.sql
@@ -0,0 +1,206 @@
+SET citus.shard_count TO 2;
+SET citus.next_shard_id TO 750000;
+SET citus.next_placement_id TO 750000;
+
+CREATE SCHEMA indirections;
+SET search_path TO indirections;
+
+-- specific tests related to get_update_query_targetlist_def
+-- we test only queries with sublinks, like:
+-- ( ... SET (...) = (SELECT ...))
+
+-- Reference tables
+CREATE TABLE test_ref_indirection (
+    id bigint primary key
+  , col_bool bool , col_date date , col_int integer , col_text text
+  );
+SELECT create_reference_table('indirections.test_ref_indirection');
+
+CREATE TABLE test_ref_indirection_new (
+    id bigint primary key
+  , col_bool bool , col_date date , col_int integer , col_text text
+  );
+SELECT create_reference_table('indirections.test_ref_indirection_new');
+
+-- Distributed tables
+CREATE TABLE test_dist_indirection (
+    id bigint primary key
+  , col_bool bool , col_date date , col_int integer , col_text text
+  );
+SELECT create_distributed_table('indirections.test_dist_indirection', 'id');
+
+CREATE TABLE test_dist_indirection_new (
+    id bigint primary key
+  , col_bool bool , col_date date , col_int integer , col_text text
+  );
+SELECT create_distributed_table('indirections.test_dist_indirection_new', 'id');
+
+-- Local tables required ?
+
+-- those should work:
+INSERT INTO test_ref_indirection (id, col_bool, col_date, col_int, col_text)
+    SELECT 1, true, '1970-01-01'::date, 1, 'one';
+INSERT INTO test_dist_indirection (id, col_bool, col_date, col_int, col_text)
+    SELECT 1, true, '1970-01-01'::date, 1, 'one';
+
+INSERT INTO test_ref_indirection (id, col_text, col_bool, col_date, col_int)
+    SELECT 2, 'two', false, '1970-01-01'::date, 2;
+INSERT INTO test_dist_indirection (id, col_text, col_bool, col_date, col_int)
+    SELECT 2, 'two', false, '1970-01-01'::date, 2;
+
+INSERT INTO test_ref_indirection SELECT 3, false, '1970-01-01'::date, 0, 'empty';
+INSERT INTO test_dist_indirection SELECT 3, false, '1970-01-01'::date, 0, 'empty';
+INSERT INTO test_ref_indirection SELECT 4, false, '1970-01-01'::date, 0, 'empty';
+INSERT INTO test_dist_indirection SELECT 4, false, '1970-01-01'::date, 0, 'empty';
+
+INSERT INTO test_ref_indirection_new SELECT * FROM test_ref_indirection;
+INSERT INTO test_dist_indirection_new SELECT * FROM test_dist_indirection;
+
+SELECT * FROM test_ref_indirection ORDER BY id;
+SELECT * FROM test_dist_indirection ORDER BY id;
+
+SELECT * FROM test_ref_indirection_new ORDER BY id;
+SELECT * FROM test_dist_indirection_new ORDER BY id;
+
+-- now UPDATEs
+UPDATE test_ref_indirection
+    SET (col_bool, col_date, col_int, col_text)
+        = (SELECT true, '1970-12-31'::date, 1, 'ok')
+RETURNING *;
+UPDATE test_dist_indirection
+    SET (col_bool, col_date, col_int, col_text)
+        = (SELECT true, '1970-12-31'::date, 1, 'ok')
+RETURNING *;
+
+UPDATE test_ref_indirection
+    SET (col_bool, col_date)  = (select true, '1970-06-06'::date)
+      , (col_int, col_text) = (select 1, 'still ok')
+RETURNING *;
+UPDATE test_dist_indirection
+    SET (col_bool, col_date)  = (select true, '1970-06-06'::date)
+      , (col_int, col_text) = (select 1, 'still ok')
+RETURNING *;
+
+SELECT * FROM test_ref_indirection ORDER BY id;
+SELECT * FROM test_dist_indirection ORDER BY id;
+
+-- but those should not:
+-- TODO wrong ERROR
+UPDATE test_ref_indirection
+    SET (col_date, col_text, col_int, col_bool)
+        = (SELECT '1970-12-31'::date, 'not ok', 2, false)
+RETURNING *;
+-- TODO wrong ERROR
+UPDATE test_dist_indirection
+    SET (col_date, col_text, col_int, col_bool)
+        = (SELECT '1970-12-31'::date, 'not ok', 2, false)
+RETURNING *;
+
+-- TODO wrong ERROR
+UPDATE test_ref_indirection
+    SET (col_int, col_date)  = (select 2, '1970-06-06'::date)
+      , (col_text, col_bool) = (select 'not ok', false)
+RETURNING *;
+-- TODO wrong ERROR
+UPDATE test_dist_indirection
+    SET (col_int, col_date)  = (select 2, '1970-06-06'::date)
+      , (col_text, col_bool) = (select 'not ok', false)
+RETURNING *;
+
+--
+-- more restrictive ones, just in case we miss a wrong value
+--
+-- those should work
+UPDATE test_ref_indirection
+    SET (col_bool, col_text) = (SELECT true, 'ok')
+RETURNING *;
+UPDATE test_dist_indirection
+    SET (col_bool, col_text) = (SELECT true, 'ok')
+RETURNING *;
+
+UPDATE test_ref_indirection
+    SET (col_bool, col_text) = (SELECT true, 'ok')
+WHERE id = 1
+RETURNING *;
+UPDATE test_dist_indirection
+    SET (col_bool, col_text) = (SELECT true, 'ok')
+WHERE id = 1
+RETURNING *;
+
+SELECT * FROM test_ref_indirection ORDER BY id;
+SELECT * FROM test_dist_indirection ORDER BY id;
+
+-- those should not
+-- TODO wrong ERROR
+UPDATE test_ref_indirection
+    SET (col_text, col_bool) = (SELECT 'not ok', false)
+RETURNING *;
+-- TODO wrong ERROR
+UPDATE test_dist_indirection
+    SET (col_text, col_bool) = (SELECT 'not ok', false)
+RETURNING *;
+
+-- TODO wrong ERROR
+UPDATE test_ref_indirection
+    SET (col_text, col_bool) = (SELECT 'not ok', false)
+WHERE id = 2
+RETURNING *;
+-- TODO wrong ERROR
+UPDATE test_dist_indirection
+    SET (col_text, col_bool) = (SELECT 'not ok', false)
+WHERE id = 2
+RETURNING *;
+
+-- several updates in CTE shoult not work
+-- TODO wrong ERROR
+with qq3 as (
+    update test_ref_indirection
+    SET (col_text, col_bool)
+      = (SELECT 'full', true)
+    where id = 3
+    returning *
+),
+qq4 as (
+    update test_ref_indirection
+    SET (col_text, col_bool)
+      = (SELECT 'fully', true)
+    where id = 4
+    returning *
+)
+select * from qq3 union all select * from qq4;
+-- TODO wrong ERROR
+with qq3 as (
+    update test_dist_indirection
+    SET (col_text, col_bool)
+      = (SELECT 'full', true)
+    where id = 3
+    returning *
+),
+qq4 as (
+    update test_dist_indirection
+    SET (col_text, col_bool)
+      = (SELECT 'fully', true)
+    where id = 4
+    returning *
+)
+select * from qq3 union all select * from qq4;
+
+DROP TABLE test_dist_indirection;
+DROP TABLE test_dist_indirection_new;
+DROP TABLE test_ref_indirection;
+DROP TABLE test_ref_indirection_new;
+
+-- https://github.com/citusdata/citus/issues/4092
+CREATE TABLE update_test (
+    a   INT DEFAULT 10,
+    b   INT,
+    c   TEXT
+);
+SELECT create_reference_table('indirections.update_test');
+UPDATE update_test
+SET (b,a) = (select a,b from update_test where b = 41 and c = 'car')
+WHERE a = 100 AND b = 20;
+
+-- https://github.com/citusdata/citus/pull/5692
+
+DROP SCHEMA indirections CASCADE;

From 2a6b9ffd39408287e34e503bacbcfff107eacdc3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Villemain?= <Cedric.Villemain@Data-Bene.io>
Date: Wed, 26 Feb 2025 10:09:47 +0100
Subject: [PATCH 4/6] Check more carrefuly UPDATE SET ()=(SELECT) query

Produce an ERROR when attribute in the target list have been reordered
by PostgreSQL rewritter.
---
 .../distributed/deparser/ruleutils_15.c       |  40 ++++
 .../distributed/deparser/ruleutils_16.c       |  40 ++++
 .../distributed/deparser/ruleutils_17.c       |  40 ++++
 src/test/regress/expected/indirections.out    | 176 +++++++++---------
 src/test/regress/sql/indirections.sql         |  31 ++-
 5 files changed, 223 insertions(+), 104 deletions(-)

diff --git a/src/backend/distributed/deparser/ruleutils_15.c b/src/backend/distributed/deparser/ruleutils_15.c
index 9004f7bbc..88db9662f 100644
--- a/src/backend/distributed/deparser/ruleutils_15.c
+++ b/src/backend/distributed/deparser/ruleutils_15.c
@@ -3509,6 +3509,9 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 	SubLink    *cur_ma_sublink;
 	List	   *ma_sublinks;
 
+	AttrNumber	previous_attnum = InvalidAttrNumber;
+	int 		paramid_increment = 0;
+
 	/*
 	 * Prepare to deal with MULTIEXPR assignments: collect the source SubLinks
 	 * into a list.  We expect them to appear, in ID order, in resjunk tlist
@@ -3631,11 +3634,48 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 		 */
 		if (cur_ma_sublink != NULL)
 		{
+			AttrNumber attnum = InvalidAttrNumber;
+			if (IsA(expr, Param))
+			{
+				Param *param = (Param *) expr;
+				attnum = param->paramid + paramid_increment;
+			}
+			else if (IsA(expr, FuncExpr))
+			{
+				FuncExpr *func = (FuncExpr *) expr;
+				ListCell *lc;
+
+				/* Iterate through the arguments of the FuncExpr */
+				foreach(lc, func->args)
+				{
+					Node *arg = (Node *) lfirst(lc);
+
+					/* Check if the argument is a PARAM node */
+					if (IsA(arg, Param))
+					{
+						Param *param = (Param *) arg;
+						attnum = param->paramid + paramid_increment;
+
+						break;  /* Exit loop once we find the PARAM node */
+					}
+				}
+			}
+
+			if (previous_attnum >= attnum)
+				ereport(ERROR, (errcode(ERRCODE_WRONG_OBJECT_TYPE),
+								errmsg(
+									"cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order"),
+								errhint("Sort the columns on the left side by physical order.")));
+
+			previous_attnum = attnum;
+
 			if (--remaining_ma_columns > 0)
 				continue;		/* not the last column of multiassignment */
+
 			appendStringInfoChar(buf, ')');
 			expr = (Node *) cur_ma_sublink;
 			cur_ma_sublink = NULL;
+			paramid_increment = previous_attnum;
 		}
 
 		appendStringInfoString(buf, " = ");
diff --git a/src/backend/distributed/deparser/ruleutils_16.c b/src/backend/distributed/deparser/ruleutils_16.c
index 65bbd1720..027fe3242 100644
--- a/src/backend/distributed/deparser/ruleutils_16.c
+++ b/src/backend/distributed/deparser/ruleutils_16.c
@@ -3525,6 +3525,9 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 	SubLink    *cur_ma_sublink;
 	List	   *ma_sublinks;
 
+	AttrNumber	previous_attnum = InvalidAttrNumber;
+	int 		paramid_increment = 0;
+
 	/*
 	 * Prepare to deal with MULTIEXPR assignments: collect the source SubLinks
 	 * into a list.  We expect them to appear, in ID order, in resjunk tlist
@@ -3647,11 +3650,48 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 		 */
 		if (cur_ma_sublink != NULL)
 		{
+			AttrNumber attnum = InvalidAttrNumber;
+			if (IsA(expr, Param))
+			{
+				Param *param = (Param *) expr;
+				attnum = param->paramid + paramid_increment;
+			}
+			else if (IsA(expr, FuncExpr))
+			{
+				FuncExpr *func = (FuncExpr *) expr;
+				ListCell *lc;
+
+				/* Iterate through the arguments of the FuncExpr */
+				foreach(lc, func->args)
+				{
+					Node *arg = (Node *) lfirst(lc);
+
+					/* Check if the argument is a PARAM node */
+					if (IsA(arg, Param))
+					{
+						Param *param = (Param *) arg;
+						attnum = param->paramid + paramid_increment;
+
+						break;  /* Exit loop once we find the PARAM node */
+					}
+				}
+			}
+
+			if (previous_attnum >= attnum)
+				ereport(ERROR, (errcode(ERRCODE_WRONG_OBJECT_TYPE),
+								errmsg(
+									"cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order"),
+								errhint("Sort the columns on the left side by physical order.")));
+
+			previous_attnum = attnum;
+
 			if (--remaining_ma_columns > 0)
 				continue;		/* not the last column of multiassignment */
+
 			appendStringInfoChar(buf, ')');
 			expr = (Node *) cur_ma_sublink;
 			cur_ma_sublink = NULL;
+			paramid_increment = previous_attnum;
 		}
 
 		appendStringInfoString(buf, " = ");
diff --git a/src/backend/distributed/deparser/ruleutils_17.c b/src/backend/distributed/deparser/ruleutils_17.c
index f0710e684..64ac077cb 100644
--- a/src/backend/distributed/deparser/ruleutils_17.c
+++ b/src/backend/distributed/deparser/ruleutils_17.c
@@ -3542,6 +3542,9 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 	SubLink    *cur_ma_sublink;
 	List	   *ma_sublinks;
 
+	AttrNumber	previous_attnum = InvalidAttrNumber;
+	int 		paramid_increment = 0;
+
 	/*
 	 * Prepare to deal with MULTIEXPR assignments: collect the source SubLinks
 	 * into a list.  We expect them to appear, in ID order, in resjunk tlist
@@ -3664,11 +3667,48 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 		 */
 		if (cur_ma_sublink != NULL)
 		{
+			AttrNumber attnum = InvalidAttrNumber;
+			if (IsA(expr, Param))
+			{
+				Param *param = (Param *) expr;
+				attnum = param->paramid + paramid_increment;
+			}
+			else if (IsA(expr, FuncExpr))
+			{
+				FuncExpr *func = (FuncExpr *) expr;
+				ListCell *lc;
+
+				/* Iterate through the arguments of the FuncExpr */
+				foreach(lc, func->args)
+				{
+					Node *arg = (Node *) lfirst(lc);
+
+					/* Check if the argument is a PARAM node */
+					if (IsA(arg, Param))
+					{
+						Param *param = (Param *) arg;
+						attnum = param->paramid + paramid_increment;
+
+						break;  /* Exit loop once we find the PARAM node */
+					}
+				}
+			}
+
+			if (previous_attnum >= attnum)
+				ereport(ERROR, (errcode(ERRCODE_WRONG_OBJECT_TYPE),
+								errmsg(
+									"cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order"),
+								errhint("Sort the columns on the left side by physical order.")));
+
+			previous_attnum = attnum;
+
 			if (--remaining_ma_columns > 0)
 				continue;		/* not the last column of multiassignment */
+
 			appendStringInfoChar(buf, ')');
 			expr = (Node *) cur_ma_sublink;
 			cur_ma_sublink = NULL;
+			paramid_increment = previous_attnum;
 		}
 
 		appendStringInfoString(buf, " = ");
diff --git a/src/test/regress/expected/indirections.out b/src/test/regress/expected/indirections.out
index a9ef5078a..20bd73d2b 100644
--- a/src/test/regress/expected/indirections.out
+++ b/src/test/regress/expected/indirections.out
@@ -103,103 +103,113 @@ SELECT * FROM test_dist_indirection_new ORDER BY id;
 -- now UPDATEs
 UPDATE test_ref_indirection
     SET (col_bool, col_date, col_int, col_text)
-        = (SELECT true, '1970-12-31'::date, 1, 'ok')
+        = (SELECT true, '1970-01-01'::date, 1, 'ok')
 RETURNING *;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 12-31-1970 |       1 | ok
-  2 | t        | 12-31-1970 |       1 | ok
-  3 | t        | 12-31-1970 |       1 | ok
-  4 | t        | 12-31-1970 |       1 | ok
+  1 | t        | 01-01-1970 |       1 | ok
+  2 | t        | 01-01-1970 |       1 | ok
+  3 | t        | 01-01-1970 |       1 | ok
+  4 | t        | 01-01-1970 |       1 | ok
 (4 rows)
 
 UPDATE test_dist_indirection
     SET (col_bool, col_date, col_int, col_text)
-        = (SELECT true, '1970-12-31'::date, 1, 'ok')
+        = (SELECT true, '1970-01-01'::date, 1, 'ok')
 RETURNING *;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 12-31-1970 |       1 | ok
-  2 | t        | 12-31-1970 |       1 | ok
-  3 | t        | 12-31-1970 |       1 | ok
-  4 | t        | 12-31-1970 |       1 | ok
+  1 | t        | 01-01-1970 |       1 | ok
+  2 | t        | 01-01-1970 |       1 | ok
+  3 | t        | 01-01-1970 |       1 | ok
+  4 | t        | 01-01-1970 |       1 | ok
 (4 rows)
 
 UPDATE test_ref_indirection
-    SET (col_bool, col_date)  = (select true, '1970-06-06'::date)
+    SET (col_bool, col_date)  = (select true, '1970-01-01'::date)
       , (col_int, col_text) = (select 1, 'still ok')
 RETURNING *;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 06-06-1970 |       1 | still ok
-  2 | t        | 06-06-1970 |       1 | still ok
-  3 | t        | 06-06-1970 |       1 | still ok
-  4 | t        | 06-06-1970 |       1 | still ok
+  1 | t        | 01-01-1970 |       1 | still ok
+  2 | t        | 01-01-1970 |       1 | still ok
+  3 | t        | 01-01-1970 |       1 | still ok
+  4 | t        | 01-01-1970 |       1 | still ok
 (4 rows)
 
 UPDATE test_dist_indirection
-    SET (col_bool, col_date)  = (select true, '1970-06-06'::date)
+    SET (col_bool, col_date)  = (select true, '1970-01-01'::date)
       , (col_int, col_text) = (select 1, 'still ok')
 RETURNING *;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 06-06-1970 |       1 | still ok
-  2 | t        | 06-06-1970 |       1 | still ok
-  3 | t        | 06-06-1970 |       1 | still ok
-  4 | t        | 06-06-1970 |       1 | still ok
+  1 | t        | 01-01-1970 |       1 | still ok
+  2 | t        | 01-01-1970 |       1 | still ok
+  3 | t        | 01-01-1970 |       1 | still ok
+  4 | t        | 01-01-1970 |       1 | still ok
+(4 rows)
+
+UPDATE test_ref_indirection
+    SET (col_bool, col_int)  = (select true, 1)
+      , (col_text) = (select 'ok')
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1970 |       1 | ok
+  2 | t        | 01-01-1970 |       1 | ok
+  3 | t        | 01-01-1970 |       1 | ok
+  4 | t        | 01-01-1970 |       1 | ok
 (4 rows)
 
 SELECT * FROM test_ref_indirection ORDER BY id;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 06-06-1970 |       1 | still ok
-  2 | t        | 06-06-1970 |       1 | still ok
-  3 | t        | 06-06-1970 |       1 | still ok
-  4 | t        | 06-06-1970 |       1 | still ok
+  1 | t        | 01-01-1970 |       1 | ok
+  2 | t        | 01-01-1970 |       1 | ok
+  3 | t        | 01-01-1970 |       1 | ok
+  4 | t        | 01-01-1970 |       1 | ok
 (4 rows)
 
 SELECT * FROM test_dist_indirection ORDER BY id;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 06-06-1970 |       1 | still ok
-  2 | t        | 06-06-1970 |       1 | still ok
-  3 | t        | 06-06-1970 |       1 | still ok
-  4 | t        | 06-06-1970 |       1 | still ok
+  1 | t        | 01-01-1970 |       1 | still ok
+  2 | t        | 01-01-1970 |       1 | still ok
+  3 | t        | 01-01-1970 |       1 | still ok
+  4 | t        | 01-01-1970 |       1 | still ok
 (4 rows)
 
 -- but those should not:
--- TODO wrong ERROR
 UPDATE test_ref_indirection
     SET (col_date, col_text, col_int, col_bool)
-        = (SELECT '1970-12-31'::date, 'not ok', 2, false)
+        = (SELECT '1970-06-06'::date, 'not ok', 2, false)
 RETURNING *;
-ERROR:  column "col_bool" is of type boolean but expression is of type date
-HINT:  You will need to rewrite or cast the expression.
-CONTEXT:  while executing command on localhost:xxxxx
--- TODO wrong ERROR
+ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
+HINT:  Sort the columns on the left side by physical order.
 UPDATE test_dist_indirection
     SET (col_date, col_text, col_int, col_bool)
-        = (SELECT '1970-12-31'::date, 'not ok', 2, false)
+        = (SELECT '1970-06-06'::date, 'not ok', 2, false)
 RETURNING *;
-ERROR:  column "col_bool" is of type boolean but expression is of type date
-HINT:  You will need to rewrite or cast the expression.
-CONTEXT:  while executing command on localhost:xxxxx
--- TODO wrong ERROR
+ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
+HINT:  Sort the columns on the left side by physical order.
+UPDATE test_ref_indirection
+    SET (col_int, col_text)  = (select 2, 'not ok')
+      , (col_bool) = (select false)
+RETURNING *;
+ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
+HINT:  Sort the columns on the left side by physical order.
 UPDATE test_ref_indirection
     SET (col_int, col_date)  = (select 2, '1970-06-06'::date)
       , (col_text, col_bool) = (select 'not ok', false)
 RETURNING *;
-ERROR:  column "col_bool" is of type boolean but expression is of type integer
-HINT:  You will need to rewrite or cast the expression.
-CONTEXT:  while executing command on localhost:xxxxx
--- TODO wrong ERROR
+ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
+HINT:  Sort the columns on the left side by physical order.
 UPDATE test_dist_indirection
     SET (col_int, col_date)  = (select 2, '1970-06-06'::date)
       , (col_text, col_bool) = (select 'not ok', false)
 RETURNING *;
-ERROR:  column "col_bool" is of type boolean but expression is of type integer
-HINT:  You will need to rewrite or cast the expression.
-CONTEXT:  while executing command on localhost:xxxxx
+ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
+HINT:  Sort the columns on the left side by physical order.
 --
 -- more restrictive ones, just in case we miss a wrong value
 --
@@ -209,10 +219,10 @@ UPDATE test_ref_indirection
 RETURNING *;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 06-06-1970 |       1 | ok
-  2 | t        | 06-06-1970 |       1 | ok
-  3 | t        | 06-06-1970 |       1 | ok
-  4 | t        | 06-06-1970 |       1 | ok
+  1 | t        | 01-01-1970 |       1 | ok
+  2 | t        | 01-01-1970 |       1 | ok
+  3 | t        | 01-01-1970 |       1 | ok
+  4 | t        | 01-01-1970 |       1 | ok
 (4 rows)
 
 UPDATE test_dist_indirection
@@ -220,10 +230,10 @@ UPDATE test_dist_indirection
 RETURNING *;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 06-06-1970 |       1 | ok
-  2 | t        | 06-06-1970 |       1 | ok
-  3 | t        | 06-06-1970 |       1 | ok
-  4 | t        | 06-06-1970 |       1 | ok
+  1 | t        | 01-01-1970 |       1 | ok
+  2 | t        | 01-01-1970 |       1 | ok
+  3 | t        | 01-01-1970 |       1 | ok
+  4 | t        | 01-01-1970 |       1 | ok
 (4 rows)
 
 UPDATE test_ref_indirection
@@ -232,7 +242,7 @@ WHERE id = 1
 RETURNING *;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 06-06-1970 |       1 | ok
+  1 | t        | 01-01-1970 |       1 | ok
 (1 row)
 
 UPDATE test_dist_indirection
@@ -241,60 +251,51 @@ WHERE id = 1
 RETURNING *;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 06-06-1970 |       1 | ok
+  1 | t        | 01-01-1970 |       1 | ok
 (1 row)
 
 SELECT * FROM test_ref_indirection ORDER BY id;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 06-06-1970 |       1 | ok
-  2 | t        | 06-06-1970 |       1 | ok
-  3 | t        | 06-06-1970 |       1 | ok
-  4 | t        | 06-06-1970 |       1 | ok
+  1 | t        | 01-01-1970 |       1 | ok
+  2 | t        | 01-01-1970 |       1 | ok
+  3 | t        | 01-01-1970 |       1 | ok
+  4 | t        | 01-01-1970 |       1 | ok
 (4 rows)
 
 SELECT * FROM test_dist_indirection ORDER BY id;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 06-06-1970 |       1 | ok
-  2 | t        | 06-06-1970 |       1 | ok
-  3 | t        | 06-06-1970 |       1 | ok
-  4 | t        | 06-06-1970 |       1 | ok
+  1 | t        | 01-01-1970 |       1 | ok
+  2 | t        | 01-01-1970 |       1 | ok
+  3 | t        | 01-01-1970 |       1 | ok
+  4 | t        | 01-01-1970 |       1 | ok
 (4 rows)
 
 -- those should not
--- TODO wrong ERROR
 UPDATE test_ref_indirection
     SET (col_text, col_bool) = (SELECT 'not ok', false)
 RETURNING *;
-ERROR:  column "col_bool" is of type boolean but expression is of type text
-HINT:  You will need to rewrite or cast the expression.
-CONTEXT:  while executing command on localhost:xxxxx
--- TODO wrong ERROR
+ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
+HINT:  Sort the columns on the left side by physical order.
 UPDATE test_dist_indirection
     SET (col_text, col_bool) = (SELECT 'not ok', false)
 RETURNING *;
-ERROR:  column "col_bool" is of type boolean but expression is of type text
-HINT:  You will need to rewrite or cast the expression.
-CONTEXT:  while executing command on localhost:xxxxx
--- TODO wrong ERROR
+ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
+HINT:  Sort the columns on the left side by physical order.
 UPDATE test_ref_indirection
     SET (col_text, col_bool) = (SELECT 'not ok', false)
 WHERE id = 2
 RETURNING *;
-ERROR:  column "col_bool" is of type boolean but expression is of type text
-HINT:  You will need to rewrite or cast the expression.
-CONTEXT:  while executing command on localhost:xxxxx
--- TODO wrong ERROR
+ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
+HINT:  Sort the columns on the left side by physical order.
 UPDATE test_dist_indirection
     SET (col_text, col_bool) = (SELECT 'not ok', false)
 WHERE id = 2
 RETURNING *;
-ERROR:  column "col_bool" is of type boolean but expression is of type text
-HINT:  You will need to rewrite or cast the expression.
-CONTEXT:  while executing command on localhost:xxxxx
+ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
+HINT:  Sort the columns on the left side by physical order.
 -- several updates in CTE shoult not work
--- TODO wrong ERROR
 with qq3 as (
     update test_ref_indirection
     SET (col_text, col_bool)
@@ -310,10 +311,8 @@ qq4 as (
     returning *
 )
 select * from qq3 union all select * from qq4;
-ERROR:  column "col_bool" is of type boolean but expression is of type text
-HINT:  You will need to rewrite or cast the expression.
-CONTEXT:  while executing command on localhost:xxxxx
--- TODO wrong ERROR
+ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
+HINT:  Sort the columns on the left side by physical order.
 with qq3 as (
     update test_dist_indirection
     SET (col_text, col_bool)
@@ -329,9 +328,8 @@ qq4 as (
     returning *
 )
 select * from qq3 union all select * from qq4;
-ERROR:  column "col_bool" is of type boolean but expression is of type text
-HINT:  You will need to rewrite or cast the expression.
-CONTEXT:  while executing command on localhost:xxxxx
+ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
+HINT:  Sort the columns on the left side by physical order.
 DROP TABLE test_dist_indirection;
 DROP TABLE test_dist_indirection_new;
 DROP TABLE test_ref_indirection;
@@ -351,6 +349,8 @@ SELECT create_reference_table('indirections.update_test');
 UPDATE update_test
 SET (b,a) = (select a,b from update_test where b = 41 and c = 'car')
 WHERE a = 100 AND b = 20;
+ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
+HINT:  Sort the columns on the left side by physical order.
 -- https://github.com/citusdata/citus/pull/5692
 DROP SCHEMA indirections CASCADE;
 NOTICE:  drop cascades to 2 other objects
diff --git a/src/test/regress/sql/indirections.sql b/src/test/regress/sql/indirections.sql
index fd8119ff5..73ff78b71 100644
--- a/src/test/regress/sql/indirections.sql
+++ b/src/test/regress/sql/indirections.sql
@@ -65,43 +65,48 @@ SELECT * FROM test_dist_indirection_new ORDER BY id;
 -- now UPDATEs
 UPDATE test_ref_indirection
     SET (col_bool, col_date, col_int, col_text)
-        = (SELECT true, '1970-12-31'::date, 1, 'ok')
+        = (SELECT true, '1970-01-01'::date, 1, 'ok')
 RETURNING *;
 UPDATE test_dist_indirection
     SET (col_bool, col_date, col_int, col_text)
-        = (SELECT true, '1970-12-31'::date, 1, 'ok')
+        = (SELECT true, '1970-01-01'::date, 1, 'ok')
 RETURNING *;
 
 UPDATE test_ref_indirection
-    SET (col_bool, col_date)  = (select true, '1970-06-06'::date)
+    SET (col_bool, col_date)  = (select true, '1970-01-01'::date)
       , (col_int, col_text) = (select 1, 'still ok')
 RETURNING *;
 UPDATE test_dist_indirection
-    SET (col_bool, col_date)  = (select true, '1970-06-06'::date)
+    SET (col_bool, col_date)  = (select true, '1970-01-01'::date)
       , (col_int, col_text) = (select 1, 'still ok')
 RETURNING *;
 
+UPDATE test_ref_indirection
+    SET (col_bool, col_int)  = (select true, 1)
+      , (col_text) = (select 'ok')
+RETURNING *;
+
 SELECT * FROM test_ref_indirection ORDER BY id;
 SELECT * FROM test_dist_indirection ORDER BY id;
 
 -- but those should not:
--- TODO wrong ERROR
 UPDATE test_ref_indirection
     SET (col_date, col_text, col_int, col_bool)
-        = (SELECT '1970-12-31'::date, 'not ok', 2, false)
+        = (SELECT '1970-06-06'::date, 'not ok', 2, false)
 RETURNING *;
--- TODO wrong ERROR
 UPDATE test_dist_indirection
     SET (col_date, col_text, col_int, col_bool)
-        = (SELECT '1970-12-31'::date, 'not ok', 2, false)
+        = (SELECT '1970-06-06'::date, 'not ok', 2, false)
 RETURNING *;
 
--- TODO wrong ERROR
+UPDATE test_ref_indirection
+    SET (col_int, col_text)  = (select 2, 'not ok')
+      , (col_bool) = (select false)
+RETURNING *;
 UPDATE test_ref_indirection
     SET (col_int, col_date)  = (select 2, '1970-06-06'::date)
       , (col_text, col_bool) = (select 'not ok', false)
 RETURNING *;
--- TODO wrong ERROR
 UPDATE test_dist_indirection
     SET (col_int, col_date)  = (select 2, '1970-06-06'::date)
       , (col_text, col_bool) = (select 'not ok', false)
@@ -131,28 +136,23 @@ SELECT * FROM test_ref_indirection ORDER BY id;
 SELECT * FROM test_dist_indirection ORDER BY id;
 
 -- those should not
--- TODO wrong ERROR
 UPDATE test_ref_indirection
     SET (col_text, col_bool) = (SELECT 'not ok', false)
 RETURNING *;
--- TODO wrong ERROR
 UPDATE test_dist_indirection
     SET (col_text, col_bool) = (SELECT 'not ok', false)
 RETURNING *;
 
--- TODO wrong ERROR
 UPDATE test_ref_indirection
     SET (col_text, col_bool) = (SELECT 'not ok', false)
 WHERE id = 2
 RETURNING *;
--- TODO wrong ERROR
 UPDATE test_dist_indirection
     SET (col_text, col_bool) = (SELECT 'not ok', false)
 WHERE id = 2
 RETURNING *;
 
 -- several updates in CTE shoult not work
--- TODO wrong ERROR
 with qq3 as (
     update test_ref_indirection
     SET (col_text, col_bool)
@@ -168,7 +168,6 @@ qq4 as (
     returning *
 )
 select * from qq3 union all select * from qq4;
--- TODO wrong ERROR
 with qq3 as (
     update test_dist_indirection
     SET (col_text, col_bool)

From dcdaaaa88dfd4bdc9b31cf308a2626fa5e47c7cc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Villemain?= <Cedric.Villemain@Data-Bene.io>
Date: Sun, 2 Mar 2025 08:00:02 +0100
Subject: [PATCH 5/6] WIP reorder target list

All OK for basic UPDATE multi SELECT
---
 .../distributed/deparser/citus_ruleutils.c    |  70 ++++
 .../distributed/deparser/ruleutils_16.c       |  72 ++---
 src/include/distributed/citus_ruleutils.h     |   3 +
 src/test/regress/expected/indirections.out    | 301 ++++++++++++------
 src/test/regress/sql/indirections.sql         |  95 +++---
 5 files changed, 359 insertions(+), 182 deletions(-)

diff --git a/src/backend/distributed/deparser/citus_ruleutils.c b/src/backend/distributed/deparser/citus_ruleutils.c
index d590495a6..a7f19c869 100644
--- a/src/backend/distributed/deparser/citus_ruleutils.c
+++ b/src/backend/distributed/deparser/citus_ruleutils.c
@@ -1715,3 +1715,73 @@ RoleSpecString(RoleSpec *spec, bool withQuoteIdentifier)
 		}
 	}
 }
+
+
+/*
+ * list_sort comparator to sort target list by paramid (in MULTIEXPR)
+ * Intended for indirection management: UPDATE SET () = (SELECT )
+ */
+int
+target_list_cmp(const ListCell *a, const ListCell *b)
+{
+	TargetEntry *tleA = lfirst(a);
+	TargetEntry *tleB = lfirst(b);
+
+	if (IsA(tleA->expr, Param) && IsA(tleB->expr, Param))
+	{
+		int la = ((Param *) tleA->expr)->paramid;
+		int lb = ((Param *) tleB->expr)->paramid;
+		return (la > lb) - (la < lb);
+	}
+	else if ((IsA(tleA->expr, Param) && IsA(tleB->expr, SubLink)) ||
+			 (IsA(tleA->expr, SubLink) && IsA(tleB->expr, Param)) ||
+			 (IsA(tleA->expr, SubLink) && IsA(tleB->expr, SubLink)))
+	{
+		return -1;
+	}
+	else
+	{
+		elog(ERROR, "unexpected nodes");
+	}
+}
+
+
+/*
+ * Recursively search an expression for a Param and return its paramid
+ * Intended for indirection management: UPDATE SET () = (SELECT )
+ * Does not cover all options but those supported by Citus.
+ */
+int
+GetParamId(Node *expr)
+{
+	int paramid = 0;
+
+	if (expr == NULL)
+	{
+		return paramid;
+	}
+
+	/* If it's a Param, return its attnum */
+	if (IsA(expr, Param))
+	{
+		Param *param = (Param *) expr;
+		paramid = param->paramid;
+	}
+	/* If it's a FuncExpr, search in arguments */
+	else if (IsA(expr, FuncExpr))
+	{
+		FuncExpr *func = (FuncExpr *) expr;
+		ListCell *lc;
+
+		foreach(lc, func->args)
+		{
+			paramid = GetParamId((Node *) lfirst(lc));
+			if (paramid != 0)
+			{
+				break; /* Stop at the first valid paramid */
+			}
+		}
+	}
+
+	return paramid;
+}
diff --git a/src/backend/distributed/deparser/ruleutils_16.c b/src/backend/distributed/deparser/ruleutils_16.c
index 027fe3242..91eeb4b76 100644
--- a/src/backend/distributed/deparser/ruleutils_16.c
+++ b/src/backend/distributed/deparser/ruleutils_16.c
@@ -3525,9 +3525,6 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 	SubLink    *cur_ma_sublink;
 	List	   *ma_sublinks;
 
-	AttrNumber	previous_attnum = InvalidAttrNumber;
-	int 		paramid_increment = 0;
-
 	/*
 	 * Prepare to deal with MULTIEXPR assignments: collect the source SubLinks
 	 * into a list.  We expect them to appear, in ID order, in resjunk tlist
@@ -3536,13 +3533,21 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 	ma_sublinks = NIL;
 	if (query->hasSubLinks)		/* else there can't be any */
 	{
+		bool saw_junk = false;
+		bool need_to_sort_target_list = false;
+		int previous_paramid = 0;
+
 		foreach(l, targetList)
 		{
 			TargetEntry *tle = (TargetEntry *) lfirst(l);
 
+			// elog(WARNING, "TOP node to string: %s", nodeToString(tle->expr));
+			// elog(WARNING, "TOP node type: %d", (int) nodeTag(tle->expr));
+
 			if (tle->resjunk && IsA(tle->expr, SubLink))
 			{
 				SubLink    *sl = (SubLink *) tle->expr;
+				saw_junk = true;
 
 				if (sl->subLinkType == MULTIEXPR_SUBLINK)
 				{
@@ -3550,7 +3555,30 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 					Assert(sl->subLinkId == list_length(ma_sublinks));
 				}
 			}
+			else if (!tle->resjunk)
+			{
+				int paramid = 0;
+				if (saw_junk)
+					elog(ERROR, "out of order target list");
+
+				paramid = GetParamId((Node *) tle->expr);
+				if (paramid < previous_paramid)
+					need_to_sort_target_list = true;
+
+				previous_paramid = paramid;
+			}
 		}
+
+		/*
+		 * reorder the target list on left side of the update:
+		 * SET () = (SELECT )
+		 * reordering the SELECT side only does not work, consider a case like:
+		 * SET (col_1, col3) = (SELECT 1, 3), (col_2) = (SELECT 2)
+		 * Then default order will lead to:
+		 * SET (col_1, col2) = (SELECT 1, 3), (col_3) = (SELECT 2)
+		 */
+		if (need_to_sort_target_list)
+			list_sort(targetList, target_list_cmp);
 	}
 	next_ma_cell = list_head(ma_sublinks);
 	cur_ma_sublink = NULL;
@@ -3650,54 +3678,18 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 		 */
 		if (cur_ma_sublink != NULL)
 		{
-			AttrNumber attnum = InvalidAttrNumber;
-			if (IsA(expr, Param))
-			{
-				Param *param = (Param *) expr;
-				attnum = param->paramid + paramid_increment;
-			}
-			else if (IsA(expr, FuncExpr))
-			{
-				FuncExpr *func = (FuncExpr *) expr;
-				ListCell *lc;
-
-				/* Iterate through the arguments of the FuncExpr */
-				foreach(lc, func->args)
-				{
-					Node *arg = (Node *) lfirst(lc);
-
-					/* Check if the argument is a PARAM node */
-					if (IsA(arg, Param))
-					{
-						Param *param = (Param *) arg;
-						attnum = param->paramid + paramid_increment;
-
-						break;  /* Exit loop once we find the PARAM node */
-					}
-				}
-			}
-
-			if (previous_attnum >= attnum)
-				ereport(ERROR, (errcode(ERRCODE_WRONG_OBJECT_TYPE),
-								errmsg(
-									"cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order"),
-								errhint("Sort the columns on the left side by physical order.")));
-
-			previous_attnum = attnum;
-
 			if (--remaining_ma_columns > 0)
 				continue;		/* not the last column of multiassignment */
-
 			appendStringInfoChar(buf, ')');
 			expr = (Node *) cur_ma_sublink;
 			cur_ma_sublink = NULL;
-			paramid_increment = previous_attnum;
 		}
 
 		appendStringInfoString(buf, " = ");
 
 		get_rule_expr(expr, context, false);
 	}
+	elog(DEBUG4, "rewriten query: %s", buf->data);
 }
 
 /* ----------
diff --git a/src/include/distributed/citus_ruleutils.h b/src/include/distributed/citus_ruleutils.h
index 3a9c36482..9089ca3a7 100644
--- a/src/include/distributed/citus_ruleutils.h
+++ b/src/include/distributed/citus_ruleutils.h
@@ -60,5 +60,8 @@ extern char * generate_operator_name(Oid operid, Oid arg1, Oid arg2);
 extern List * getOwnedSequences_internal(Oid relid, AttrNumber attnum, char deptype);
 extern void AppendOptionListToString(StringInfo stringData, List *options);
 
+extern int GetParamId(Node *expr);
+extern int target_list_cmp(const ListCell *a, const ListCell *b);
+
 
 #endif /* CITUS_RULEUTILS_H */
diff --git a/src/test/regress/expected/indirections.out b/src/test/regress/expected/indirections.out
index 20bd73d2b..115f5a7b4 100644
--- a/src/test/regress/expected/indirections.out
+++ b/src/test/regress/expected/indirections.out
@@ -126,210 +126,303 @@ RETURNING *;
 (4 rows)
 
 UPDATE test_ref_indirection
-    SET (col_bool, col_date)  = (select true, '1970-01-01'::date)
-      , (col_int, col_text) = (select 1, 'still ok')
+    SET (col_bool, col_date)  = (select false, '1971-01-01'::date)
+      , (col_int, col_text) = (select 2, '2 ok')
 RETURNING *;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 01-01-1970 |       1 | still ok
-  2 | t        | 01-01-1970 |       1 | still ok
-  3 | t        | 01-01-1970 |       1 | still ok
-  4 | t        | 01-01-1970 |       1 | still ok
+  1 | f        | 01-01-1971 |       2 | 2 ok
+  2 | f        | 01-01-1971 |       2 | 2 ok
+  3 | f        | 01-01-1971 |       2 | 2 ok
+  4 | f        | 01-01-1971 |       2 | 2 ok
 (4 rows)
 
 UPDATE test_dist_indirection
-    SET (col_bool, col_date)  = (select true, '1970-01-01'::date)
-      , (col_int, col_text) = (select 1, 'still ok')
+    SET (col_bool, col_date)  = (select false, '1971-01-01'::date)
+      , (col_int, col_text) = (select 2, '2 ok')
 RETURNING *;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 01-01-1970 |       1 | still ok
-  2 | t        | 01-01-1970 |       1 | still ok
-  3 | t        | 01-01-1970 |       1 | still ok
-  4 | t        | 01-01-1970 |       1 | still ok
+  1 | f        | 01-01-1971 |       2 | 2 ok
+  2 | f        | 01-01-1971 |       2 | 2 ok
+  3 | f        | 01-01-1971 |       2 | 2 ok
+  4 | f        | 01-01-1971 |       2 | 2 ok
 (4 rows)
 
 UPDATE test_ref_indirection
-    SET (col_bool, col_int)  = (select true, 1)
-      , (col_text) = (select 'ok')
+    SET (col_bool, col_int)  = (select true, 3)
+      , (col_text) = (select '3 ok')
 RETURNING *;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 01-01-1970 |       1 | ok
-  2 | t        | 01-01-1970 |       1 | ok
-  3 | t        | 01-01-1970 |       1 | ok
-  4 | t        | 01-01-1970 |       1 | ok
+  1 | t        | 01-01-1971 |       3 | 3 ok
+  2 | t        | 01-01-1971 |       3 | 3 ok
+  3 | t        | 01-01-1971 |       3 | 3 ok
+  4 | t        | 01-01-1971 |       3 | 3 ok
 (4 rows)
 
-SELECT * FROM test_ref_indirection ORDER BY id;
+UPDATE test_dist_indirection
+    SET (col_bool, col_int)  = (select true, 3)
+      , (col_text) = (select '3 ok')
+RETURNING *;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 01-01-1970 |       1 | ok
-  2 | t        | 01-01-1970 |       1 | ok
-  3 | t        | 01-01-1970 |       1 | ok
-  4 | t        | 01-01-1970 |       1 | ok
+  1 | t        | 01-01-1971 |       3 | 3 ok
+  2 | t        | 01-01-1971 |       3 | 3 ok
+  3 | t        | 01-01-1971 |       3 | 3 ok
+  4 | t        | 01-01-1971 |       3 | 3 ok
 (4 rows)
 
-SELECT * FROM test_dist_indirection ORDER BY id;
- id | col_bool |  col_date  | col_int | col_text
----------------------------------------------------------------------
-  1 | t        | 01-01-1970 |       1 | still ok
-  2 | t        | 01-01-1970 |       1 | still ok
-  3 | t        | 01-01-1970 |       1 | still ok
-  4 | t        | 01-01-1970 |       1 | still ok
-(4 rows)
-
--- but those should not:
+-- but those should work since 13.X
 UPDATE test_ref_indirection
     SET (col_date, col_text, col_int, col_bool)
-        = (SELECT '1970-06-06'::date, 'not ok', 2, false)
+        = (SELECT '1972-01-01'::date, '4 ok', 4, false)
 RETURNING *;
-ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
-HINT:  Sort the columns on the left side by physical order.
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | f        | 01-01-1972 |       4 | 4 ok
+  2 | f        | 01-01-1972 |       4 | 4 ok
+  3 | f        | 01-01-1972 |       4 | 4 ok
+  4 | f        | 01-01-1972 |       4 | 4 ok
+(4 rows)
+
 UPDATE test_dist_indirection
     SET (col_date, col_text, col_int, col_bool)
-        = (SELECT '1970-06-06'::date, 'not ok', 2, false)
+        = (SELECT '1972-01-01'::date, '4 ok', 4, false)
 RETURNING *;
-ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
-HINT:  Sort the columns on the left side by physical order.
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | f        | 01-01-1972 |       4 | 4 ok
+  2 | f        | 01-01-1972 |       4 | 4 ok
+  3 | f        | 01-01-1972 |       4 | 4 ok
+  4 | f        | 01-01-1972 |       4 | 4 ok
+(4 rows)
+
 UPDATE test_ref_indirection
-    SET (col_int, col_text)  = (select 2, 'not ok')
+    SET (col_int, col_text)  = (select 5, '5 ok')
+      , (col_bool) = (select true)
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1972 |       5 | 5 ok
+  2 | t        | 01-01-1972 |       5 | 5 ok
+  3 | t        | 01-01-1972 |       5 | 5 ok
+  4 | t        | 01-01-1972 |       5 | 5 ok
+(4 rows)
+
+UPDATE test_dist_indirection
+    SET (col_int, col_text)  = (select 5, '5 ok')
+      , (col_bool) = (select true)
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1972 |       5 | 5 ok
+  2 | t        | 01-01-1972 |       5 | 5 ok
+  3 | t        | 01-01-1972 |       5 | 5 ok
+  4 | t        | 01-01-1972 |       5 | 5 ok
+(4 rows)
+
+UPDATE test_ref_indirection
+    SET (col_int, col_date)  = (select 6, '1973-01-01'::date)
+      , (col_text, col_bool) = (select '6 ok', false)
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | f        | 01-01-1973 |       6 | 6 ok
+  2 | f        | 01-01-1973 |       6 | 6 ok
+  3 | f        | 01-01-1973 |       6 | 6 ok
+  4 | f        | 01-01-1973 |       6 | 6 ok
+(4 rows)
+
+UPDATE test_dist_indirection
+    SET (col_int, col_date)  = (select 6, '1973-01-01'::date)
+      , (col_text, col_bool) = (select '6 ok', false)
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | f        | 01-01-1973 |       6 | 6 ok
+  2 | f        | 01-01-1973 |       6 | 6 ok
+  3 | f        | 01-01-1973 |       6 | 6 ok
+  4 | f        | 01-01-1973 |       6 | 6 ok
+(4 rows)
+
+UPDATE test_ref_indirection
+    SET (col_int, col_date, col_text)  = (select 7, '1974-01-01'::date, '7 ok')
+      , (col_bool) = (select true)
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1974 |       7 | 7 ok
+  2 | t        | 01-01-1974 |       7 | 7 ok
+  3 | t        | 01-01-1974 |       7 | 7 ok
+  4 | t        | 01-01-1974 |       7 | 7 ok
+(4 rows)
+
+UPDATE test_dist_indirection
+    SET (col_int, col_date, col_text)  = (select 7, '1974-01-01'::date, '7 ok')
+      , (col_bool) = (select true)
+RETURNING *;
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | t        | 01-01-1974 |       7 | 7 ok
+  2 | t        | 01-01-1974 |       7 | 7 ok
+  3 | t        | 01-01-1974 |       7 | 7 ok
+  4 | t        | 01-01-1974 |       7 | 7 ok
+(4 rows)
+
+UPDATE test_ref_indirection
+    SET (col_date, col_text)  = (select '1975-01-01'::date, '8 ok')
+      , (col_int) = (select 8)
       , (col_bool) = (select false)
 RETURNING *;
-ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
-HINT:  Sort the columns on the left side by physical order.
-UPDATE test_ref_indirection
-    SET (col_int, col_date)  = (select 2, '1970-06-06'::date)
-      , (col_text, col_bool) = (select 'not ok', false)
-RETURNING *;
-ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
-HINT:  Sort the columns on the left side by physical order.
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | f        | 01-01-1975 |       8 | 8 ok
+  2 | f        | 01-01-1975 |       8 | 8 ok
+  3 | f        | 01-01-1975 |       8 | 8 ok
+  4 | f        | 01-01-1975 |       8 | 8 ok
+(4 rows)
+
 UPDATE test_dist_indirection
-    SET (col_int, col_date)  = (select 2, '1970-06-06'::date)
-      , (col_text, col_bool) = (select 'not ok', false)
+    SET (col_date, col_text)  = (select '1975-01-01'::date, '8 ok')
+      , (col_int) = (select 8)
+      , (col_bool) = (select false)
 RETURNING *;
-ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
-HINT:  Sort the columns on the left side by physical order.
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  1 | f        | 01-01-1975 |       8 | 8 ok
+  2 | f        | 01-01-1975 |       8 | 8 ok
+  3 | f        | 01-01-1975 |       8 | 8 ok
+  4 | f        | 01-01-1975 |       8 | 8 ok
+(4 rows)
+
 --
 -- more restrictive ones, just in case we miss a wrong value
 --
 -- those should work
 UPDATE test_ref_indirection
-    SET (col_bool, col_text) = (SELECT true, 'ok')
+    SET (col_bool, col_text) = (SELECT true, '9 ok')
 RETURNING *;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 01-01-1970 |       1 | ok
-  2 | t        | 01-01-1970 |       1 | ok
-  3 | t        | 01-01-1970 |       1 | ok
-  4 | t        | 01-01-1970 |       1 | ok
+  1 | t        | 01-01-1975 |       8 | 9 ok
+  2 | t        | 01-01-1975 |       8 | 9 ok
+  3 | t        | 01-01-1975 |       8 | 9 ok
+  4 | t        | 01-01-1975 |       8 | 9 ok
 (4 rows)
 
 UPDATE test_dist_indirection
-    SET (col_bool, col_text) = (SELECT true, 'ok')
+    SET (col_bool, col_text) = (SELECT true, '9 ok')
 RETURNING *;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 01-01-1970 |       1 | ok
-  2 | t        | 01-01-1970 |       1 | ok
-  3 | t        | 01-01-1970 |       1 | ok
-  4 | t        | 01-01-1970 |       1 | ok
+  1 | t        | 01-01-1975 |       8 | 9 ok
+  2 | t        | 01-01-1975 |       8 | 9 ok
+  3 | t        | 01-01-1975 |       8 | 9 ok
+  4 | t        | 01-01-1975 |       8 | 9 ok
 (4 rows)
 
 UPDATE test_ref_indirection
-    SET (col_bool, col_text) = (SELECT true, 'ok')
+    SET (col_bool, col_text) = (SELECT false, '10 ok')
 WHERE id = 1
 RETURNING *;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 01-01-1970 |       1 | ok
+  1 | f        | 01-01-1975 |       8 | 10 ok
 (1 row)
 
 UPDATE test_dist_indirection
-    SET (col_bool, col_text) = (SELECT true, 'ok')
+    SET (col_bool, col_text) = (SELECT false, '10 ok')
 WHERE id = 1
 RETURNING *;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 01-01-1970 |       1 | ok
+  1 | f        | 01-01-1975 |       8 | 10 ok
 (1 row)
 
-SELECT * FROM test_ref_indirection ORDER BY id;
+UPDATE test_ref_indirection
+    SET (col_text, col_bool) = (SELECT '11 ok', true)
+RETURNING *;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 01-01-1970 |       1 | ok
-  2 | t        | 01-01-1970 |       1 | ok
-  3 | t        | 01-01-1970 |       1 | ok
-  4 | t        | 01-01-1970 |       1 | ok
+  1 | t        | 01-01-1975 |       8 | 11 ok
+  2 | t        | 01-01-1975 |       8 | 11 ok
+  3 | t        | 01-01-1975 |       8 | 11 ok
+  4 | t        | 01-01-1975 |       8 | 11 ok
 (4 rows)
 
-SELECT * FROM test_dist_indirection ORDER BY id;
+UPDATE test_dist_indirection
+    SET (col_text, col_bool) = (SELECT '11 ok', true)
+RETURNING *;
  id | col_bool |  col_date  | col_int | col_text
 ---------------------------------------------------------------------
-  1 | t        | 01-01-1970 |       1 | ok
-  2 | t        | 01-01-1970 |       1 | ok
-  3 | t        | 01-01-1970 |       1 | ok
-  4 | t        | 01-01-1970 |       1 | ok
+  1 | t        | 01-01-1975 |       8 | 11 ok
+  2 | t        | 01-01-1975 |       8 | 11 ok
+  3 | t        | 01-01-1975 |       8 | 11 ok
+  4 | t        | 01-01-1975 |       8 | 11 ok
 (4 rows)
 
--- those should not
 UPDATE test_ref_indirection
-    SET (col_text, col_bool) = (SELECT 'not ok', false)
-RETURNING *;
-ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
-HINT:  Sort the columns on the left side by physical order.
-UPDATE test_dist_indirection
-    SET (col_text, col_bool) = (SELECT 'not ok', false)
-RETURNING *;
-ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
-HINT:  Sort the columns on the left side by physical order.
-UPDATE test_ref_indirection
-    SET (col_text, col_bool) = (SELECT 'not ok', false)
+    SET (col_text, col_bool) = (SELECT '12 ok', false)
 WHERE id = 2
 RETURNING *;
-ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
-HINT:  Sort the columns on the left side by physical order.
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  2 | f        | 01-01-1975 |       8 | 12 ok
+(1 row)
+
 UPDATE test_dist_indirection
-    SET (col_text, col_bool) = (SELECT 'not ok', false)
+    SET (col_text, col_bool) = (SELECT '12 ok', false)
 WHERE id = 2
 RETURNING *;
-ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
-HINT:  Sort the columns on the left side by physical order.
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  2 | f        | 01-01-1975 |       8 | 12 ok
+(1 row)
+
 -- several updates in CTE shoult not work
 with qq3 as (
     update test_ref_indirection
     SET (col_text, col_bool)
-      = (SELECT 'full', true)
+      = (SELECT '13', true)
     where id = 3
     returning *
 ),
 qq4 as (
     update test_ref_indirection
     SET (col_text, col_bool)
-      = (SELECT 'fully', true)
+      = (SELECT '14', false)
     where id = 4
     returning *
 )
 select * from qq3 union all select * from qq4;
-ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
-HINT:  Sort the columns on the left side by physical order.
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  3 | t        | 01-01-1975 |       8 | 13
+  4 | f        | 01-01-1975 |       8 | 14
+(2 rows)
+
 with qq3 as (
     update test_dist_indirection
     SET (col_text, col_bool)
-      = (SELECT 'full', true)
+      = (SELECT '13', true)
     where id = 3
     returning *
 ),
 qq4 as (
     update test_dist_indirection
     SET (col_text, col_bool)
-      = (SELECT 'fully', true)
+      = (SELECT '14', false)
     where id = 4
     returning *
 )
 select * from qq3 union all select * from qq4;
-ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
-HINT:  Sort the columns on the left side by physical order.
+ id | col_bool |  col_date  | col_int | col_text
+---------------------------------------------------------------------
+  3 | t        | 01-01-1975 |       8 | 13
+  4 | f        | 01-01-1975 |       8 | 14
+(2 rows)
+
 DROP TABLE test_dist_indirection;
 DROP TABLE test_dist_indirection_new;
 DROP TABLE test_ref_indirection;
@@ -349,10 +442,6 @@ SELECT create_reference_table('indirections.update_test');
 UPDATE update_test
 SET (b,a) = (select a,b from update_test where b = 41 and c = 'car')
 WHERE a = 100 AND b = 20;
-ERROR:  cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order
-HINT:  Sort the columns on the left side by physical order.
 -- https://github.com/citusdata/citus/pull/5692
+set client_min_messages to ERROR;
 DROP SCHEMA indirections CASCADE;
-NOTICE:  drop cascades to 2 other objects
-DETAIL:  drop cascades to table update_test
-drop cascades to table update_test_750006
diff --git a/src/test/regress/sql/indirections.sql b/src/test/regress/sql/indirections.sql
index 73ff78b71..598124208 100644
--- a/src/test/regress/sql/indirections.sql
+++ b/src/test/regress/sql/indirections.sql
@@ -73,43 +73,69 @@ UPDATE test_dist_indirection
 RETURNING *;
 
 UPDATE test_ref_indirection
-    SET (col_bool, col_date)  = (select true, '1970-01-01'::date)
-      , (col_int, col_text) = (select 1, 'still ok')
+    SET (col_bool, col_date)  = (select false, '1971-01-01'::date)
+      , (col_int, col_text) = (select 2, '2 ok')
 RETURNING *;
 UPDATE test_dist_indirection
-    SET (col_bool, col_date)  = (select true, '1970-01-01'::date)
-      , (col_int, col_text) = (select 1, 'still ok')
+    SET (col_bool, col_date)  = (select false, '1971-01-01'::date)
+      , (col_int, col_text) = (select 2, '2 ok')
 RETURNING *;
 
 UPDATE test_ref_indirection
-    SET (col_bool, col_int)  = (select true, 1)
-      , (col_text) = (select 'ok')
+    SET (col_bool, col_int)  = (select true, 3)
+      , (col_text) = (select '3 ok')
+RETURNING *;
+UPDATE test_dist_indirection
+    SET (col_bool, col_int)  = (select true, 3)
+      , (col_text) = (select '3 ok')
 RETURNING *;
 
-SELECT * FROM test_ref_indirection ORDER BY id;
-SELECT * FROM test_dist_indirection ORDER BY id;
-
--- but those should not:
+-- but those should work since 13.X
 UPDATE test_ref_indirection
     SET (col_date, col_text, col_int, col_bool)
-        = (SELECT '1970-06-06'::date, 'not ok', 2, false)
+        = (SELECT '1972-01-01'::date, '4 ok', 4, false)
 RETURNING *;
 UPDATE test_dist_indirection
     SET (col_date, col_text, col_int, col_bool)
-        = (SELECT '1970-06-06'::date, 'not ok', 2, false)
+        = (SELECT '1972-01-01'::date, '4 ok', 4, false)
 RETURNING *;
 
 UPDATE test_ref_indirection
-    SET (col_int, col_text)  = (select 2, 'not ok')
+    SET (col_int, col_text)  = (select 5, '5 ok')
+      , (col_bool) = (select true)
+RETURNING *;
+UPDATE test_dist_indirection
+    SET (col_int, col_text)  = (select 5, '5 ok')
+      , (col_bool) = (select true)
+RETURNING *;
+
+UPDATE test_ref_indirection
+    SET (col_int, col_date)  = (select 6, '1973-01-01'::date)
+      , (col_text, col_bool) = (select '6 ok', false)
+RETURNING *;
+UPDATE test_dist_indirection
+    SET (col_int, col_date)  = (select 6, '1973-01-01'::date)
+      , (col_text, col_bool) = (select '6 ok', false)
+RETURNING *;
+
+UPDATE test_ref_indirection
+    SET (col_int, col_date, col_text)  = (select 7, '1974-01-01'::date, '7 ok')
+      , (col_bool) = (select true)
+RETURNING *;
+UPDATE test_dist_indirection
+    SET (col_int, col_date, col_text)  = (select 7, '1974-01-01'::date, '7 ok')
+      , (col_bool) = (select true)
+RETURNING *;
+
+UPDATE test_ref_indirection
+    SET (col_date, col_text)  = (select '1975-01-01'::date, '8 ok')
+      , (col_int) = (select 8)
       , (col_bool) = (select false)
 RETURNING *;
-UPDATE test_ref_indirection
-    SET (col_int, col_date)  = (select 2, '1970-06-06'::date)
-      , (col_text, col_bool) = (select 'not ok', false)
-RETURNING *;
 UPDATE test_dist_indirection
-    SET (col_int, col_date)  = (select 2, '1970-06-06'::date)
-      , (col_text, col_bool) = (select 'not ok', false)
+    SET (col_date, col_text)  = (select '1975-01-01'::date, '8 ok')
+      , (col_int) = (select 8)
+      , (col_bool) = (select false)
 RETURNING *;
 
 --
@@ -117,38 +143,34 @@ RETURNING *;
 --
 -- those should work
 UPDATE test_ref_indirection
-    SET (col_bool, col_text) = (SELECT true, 'ok')
+    SET (col_bool, col_text) = (SELECT true, '9 ok')
 RETURNING *;
 UPDATE test_dist_indirection
-    SET (col_bool, col_text) = (SELECT true, 'ok')
+    SET (col_bool, col_text) = (SELECT true, '9 ok')
 RETURNING *;
 
 UPDATE test_ref_indirection
-    SET (col_bool, col_text) = (SELECT true, 'ok')
+    SET (col_bool, col_text) = (SELECT false, '10 ok')
 WHERE id = 1
 RETURNING *;
 UPDATE test_dist_indirection
-    SET (col_bool, col_text) = (SELECT true, 'ok')
+    SET (col_bool, col_text) = (SELECT false, '10 ok')
 WHERE id = 1
 RETURNING *;
 
-SELECT * FROM test_ref_indirection ORDER BY id;
-SELECT * FROM test_dist_indirection ORDER BY id;
-
--- those should not
 UPDATE test_ref_indirection
-    SET (col_text, col_bool) = (SELECT 'not ok', false)
+    SET (col_text, col_bool) = (SELECT '11 ok', true)
 RETURNING *;
 UPDATE test_dist_indirection
-    SET (col_text, col_bool) = (SELECT 'not ok', false)
+    SET (col_text, col_bool) = (SELECT '11 ok', true)
 RETURNING *;
 
 UPDATE test_ref_indirection
-    SET (col_text, col_bool) = (SELECT 'not ok', false)
+    SET (col_text, col_bool) = (SELECT '12 ok', false)
 WHERE id = 2
 RETURNING *;
 UPDATE test_dist_indirection
-    SET (col_text, col_bool) = (SELECT 'not ok', false)
+    SET (col_text, col_bool) = (SELECT '12 ok', false)
 WHERE id = 2
 RETURNING *;
 
@@ -156,14 +178,14 @@ RETURNING *;
 with qq3 as (
     update test_ref_indirection
     SET (col_text, col_bool)
-      = (SELECT 'full', true)
+      = (SELECT '13', true)
     where id = 3
     returning *
 ),
 qq4 as (
     update test_ref_indirection
     SET (col_text, col_bool)
-      = (SELECT 'fully', true)
+      = (SELECT '14', false)
     where id = 4
     returning *
 )
@@ -171,14 +193,14 @@ select * from qq3 union all select * from qq4;
 with qq3 as (
     update test_dist_indirection
     SET (col_text, col_bool)
-      = (SELECT 'full', true)
+      = (SELECT '13', true)
     where id = 3
     returning *
 ),
 qq4 as (
     update test_dist_indirection
     SET (col_text, col_bool)
-      = (SELECT 'fully', true)
+      = (SELECT '14', false)
     where id = 4
     returning *
 )
@@ -201,5 +223,6 @@ SET (b,a) = (select a,b from update_test where b = 41 and c = 'car')
 WHERE a = 100 AND b = 20;
 
 -- https://github.com/citusdata/citus/pull/5692
-
+set client_min_messages to ERROR;
 DROP SCHEMA indirections CASCADE;
+

From 0eb2b4977b225216d281ae00f7e3e92459b8ab9d Mon Sep 17 00:00:00 2001
From: Ibrahim Halatci <ihalatci@gmail.com>
Date: Mon, 23 Jun 2025 13:35:54 +0000
Subject: [PATCH 6/6] Refactor logic in citus_ruleutils, Run ci tools, fix up
 regress tests

---
 .../distributed/deparser/citus_ruleutils.c    | 109 +++++++++++++-----
 .../distributed/deparser/ruleutils_15.c       |  42 +------
 .../distributed/deparser/ruleutils_16.c       |  32 +----
 .../distributed/deparser/ruleutils_17.c       |  42 +------
 .../planner/insert_select_planner.c           |   5 -
 src/include/distributed/citus_ruleutils.h     |   4 +-
 src/test/regress/citus_tests/run_test.py      |   1 +
 .../regress/expected/distributed_types.out    |   1 -
 .../regress/expected/multi_modifications.out  |  24 ++--
 .../expected/multi_mx_modifications.out       |  21 ----
 ...directions.out => multi_update_select.out} |  88 +++++++++++---
 src/test/regress/multi_1_schedule             |   1 +
 src/test/regress/sql/multi_modifications.sql  |  26 +++--
 .../regress/sql/multi_mx_modifications.sql    |  15 ---
 ...directions.sql => multi_update_select.sql} |  51 +++++---
 15 files changed, 229 insertions(+), 233 deletions(-)
 rename src/test/regress/expected/{indirections.out => multi_update_select.out} (84%)
 rename src/test/regress/sql/{indirections.sql => multi_update_select.sql} (80%)

diff --git a/src/backend/distributed/deparser/citus_ruleutils.c b/src/backend/distributed/deparser/citus_ruleutils.c
index a7f19c869..d60ebf7c4 100644
--- a/src/backend/distributed/deparser/citus_ruleutils.c
+++ b/src/backend/distributed/deparser/citus_ruleutils.c
@@ -1717,41 +1717,12 @@ RoleSpecString(RoleSpec *spec, bool withQuoteIdentifier)
 }
 
 
-/*
- * list_sort comparator to sort target list by paramid (in MULTIEXPR)
- * Intended for indirection management: UPDATE SET () = (SELECT )
- */
-int
-target_list_cmp(const ListCell *a, const ListCell *b)
-{
-	TargetEntry *tleA = lfirst(a);
-	TargetEntry *tleB = lfirst(b);
-
-	if (IsA(tleA->expr, Param) && IsA(tleB->expr, Param))
-	{
-		int la = ((Param *) tleA->expr)->paramid;
-		int lb = ((Param *) tleB->expr)->paramid;
-		return (la > lb) - (la < lb);
-	}
-	else if ((IsA(tleA->expr, Param) && IsA(tleB->expr, SubLink)) ||
-			 (IsA(tleA->expr, SubLink) && IsA(tleB->expr, Param)) ||
-			 (IsA(tleA->expr, SubLink) && IsA(tleB->expr, SubLink)))
-	{
-		return -1;
-	}
-	else
-	{
-		elog(ERROR, "unexpected nodes");
-	}
-}
-
-
 /*
  * Recursively search an expression for a Param and return its paramid
  * Intended for indirection management: UPDATE SET () = (SELECT )
  * Does not cover all options but those supported by Citus.
  */
-int
+static int
 GetParamId(Node *expr)
 {
 	int paramid = 0;
@@ -1785,3 +1756,81 @@ GetParamId(Node *expr)
 
 	return paramid;
 }
+
+
+/*
+ * list_sort comparator to sort target list by paramid (in MULTIEXPR)
+ * Intended for indirection management: UPDATE SET () = (SELECT )
+ */
+static int
+target_list_cmp(const ListCell *a, const ListCell *b)
+{
+	TargetEntry *tleA = lfirst(a);
+	TargetEntry *tleB = lfirst(b);
+
+	/*
+	 * Deal with resjunk entries; sublinks are marked resjunk and
+	 * are placed at the end of the target list so this logic
+	 * ensures they stay grouped at the end of the target list:
+	 */
+	if (tleA->resjunk || tleB->resjunk)
+	{
+		return tleA->resjunk - tleB->resjunk;
+	}
+
+	int la = GetParamId((Node *) tleA->expr);
+	int lb = GetParamId((Node *) tleB->expr);
+
+	/*
+	 * Should be looking at legitimate param ids
+	 */
+	Assert(la > 0);
+	Assert(lb > 0);
+
+	/*
+	 * Return -1, 0 or 1 depending on if la is less than,
+	 * equal to or greater than lb
+	 */
+	return (la > lb) - (la < lb);
+}
+
+
+/*
+ * Used by get_update_query_targetlist_def() (in ruleutils) to reorder the target
+ * list on the left side of the update:
+ * SET () = (SELECT )
+ * Reordering the SELECT side only does not work, consider a case like:
+ * SET (col_1, col3) = (SELECT 1, 3), (col_2) = (SELECT 2)
+ * Without ensure_update_targetlist_in_param_order(), this will lead to an incorrect
+ * deparsed query:
+ * SET (col_1, col2) = (SELECT 1, 3), (col_3) = (SELECT 2)
+ */
+void
+ensure_update_targetlist_in_param_order(List *targetList)
+{
+	bool need_to_sort_target_list = false;
+	int previous_paramid = 0;
+	ListCell *l;
+
+	foreach(l, targetList)
+	{
+		TargetEntry *tle = (TargetEntry *) lfirst(l);
+
+		if (!tle->resjunk)
+		{
+			int paramid = GetParamId((Node *) tle->expr);
+			if (paramid < previous_paramid)
+			{
+				need_to_sort_target_list = true;
+				break;
+			}
+
+			previous_paramid = paramid;
+		}
+	}
+
+	if (need_to_sort_target_list)
+	{
+		list_sort(targetList, target_list_cmp);
+	}
+}
diff --git a/src/backend/distributed/deparser/ruleutils_15.c b/src/backend/distributed/deparser/ruleutils_15.c
index 88db9662f..8a8bbcff7 100644
--- a/src/backend/distributed/deparser/ruleutils_15.c
+++ b/src/backend/distributed/deparser/ruleutils_15.c
@@ -3509,9 +3509,6 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 	SubLink    *cur_ma_sublink;
 	List	   *ma_sublinks;
 
-	AttrNumber	previous_attnum = InvalidAttrNumber;
-	int 		paramid_increment = 0;
-
 	/*
 	 * Prepare to deal with MULTIEXPR assignments: collect the source SubLinks
 	 * into a list.  We expect them to appear, in ID order, in resjunk tlist
@@ -3535,6 +3532,8 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 				}
 			}
 		}
+
+		ensure_update_targetlist_in_param_order(targetList);
 	}
 	next_ma_cell = list_head(ma_sublinks);
 	cur_ma_sublink = NULL;
@@ -3634,48 +3633,11 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 		 */
 		if (cur_ma_sublink != NULL)
 		{
-			AttrNumber attnum = InvalidAttrNumber;
-			if (IsA(expr, Param))
-			{
-				Param *param = (Param *) expr;
-				attnum = param->paramid + paramid_increment;
-			}
-			else if (IsA(expr, FuncExpr))
-			{
-				FuncExpr *func = (FuncExpr *) expr;
-				ListCell *lc;
-
-				/* Iterate through the arguments of the FuncExpr */
-				foreach(lc, func->args)
-				{
-					Node *arg = (Node *) lfirst(lc);
-
-					/* Check if the argument is a PARAM node */
-					if (IsA(arg, Param))
-					{
-						Param *param = (Param *) arg;
-						attnum = param->paramid + paramid_increment;
-
-						break;  /* Exit loop once we find the PARAM node */
-					}
-				}
-			}
-
-			if (previous_attnum >= attnum)
-				ereport(ERROR, (errcode(ERRCODE_WRONG_OBJECT_TYPE),
-								errmsg(
-									"cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order"),
-								errhint("Sort the columns on the left side by physical order.")));
-
-			previous_attnum = attnum;
-
 			if (--remaining_ma_columns > 0)
 				continue;		/* not the last column of multiassignment */
-
 			appendStringInfoChar(buf, ')');
 			expr = (Node *) cur_ma_sublink;
 			cur_ma_sublink = NULL;
-			paramid_increment = previous_attnum;
 		}
 
 		appendStringInfoString(buf, " = ");
diff --git a/src/backend/distributed/deparser/ruleutils_16.c b/src/backend/distributed/deparser/ruleutils_16.c
index 91eeb4b76..a13b99ca7 100644
--- a/src/backend/distributed/deparser/ruleutils_16.c
+++ b/src/backend/distributed/deparser/ruleutils_16.c
@@ -3533,21 +3533,13 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 	ma_sublinks = NIL;
 	if (query->hasSubLinks)		/* else there can't be any */
 	{
-		bool saw_junk = false;
-		bool need_to_sort_target_list = false;
-		int previous_paramid = 0;
-
 		foreach(l, targetList)
 		{
 			TargetEntry *tle = (TargetEntry *) lfirst(l);
 
-			// elog(WARNING, "TOP node to string: %s", nodeToString(tle->expr));
-			// elog(WARNING, "TOP node type: %d", (int) nodeTag(tle->expr));
-
 			if (tle->resjunk && IsA(tle->expr, SubLink))
 			{
 				SubLink    *sl = (SubLink *) tle->expr;
-				saw_junk = true;
 
 				if (sl->subLinkType == MULTIEXPR_SUBLINK)
 				{
@@ -3555,30 +3547,9 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 					Assert(sl->subLinkId == list_length(ma_sublinks));
 				}
 			}
-			else if (!tle->resjunk)
-			{
-				int paramid = 0;
-				if (saw_junk)
-					elog(ERROR, "out of order target list");
-
-				paramid = GetParamId((Node *) tle->expr);
-				if (paramid < previous_paramid)
-					need_to_sort_target_list = true;
-
-				previous_paramid = paramid;
-			}
 		}
 
-		/*
-		 * reorder the target list on left side of the update:
-		 * SET () = (SELECT )
-		 * reordering the SELECT side only does not work, consider a case like:
-		 * SET (col_1, col3) = (SELECT 1, 3), (col_2) = (SELECT 2)
-		 * Then default order will lead to:
-		 * SET (col_1, col2) = (SELECT 1, 3), (col_3) = (SELECT 2)
-		 */
-		if (need_to_sort_target_list)
-			list_sort(targetList, target_list_cmp);
+	        ensure_update_targetlist_in_param_order(targetList);
 	}
 	next_ma_cell = list_head(ma_sublinks);
 	cur_ma_sublink = NULL;
@@ -3689,7 +3660,6 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 
 		get_rule_expr(expr, context, false);
 	}
-	elog(DEBUG4, "rewriten query: %s", buf->data);
 }
 
 /* ----------
diff --git a/src/backend/distributed/deparser/ruleutils_17.c b/src/backend/distributed/deparser/ruleutils_17.c
index 64ac077cb..5df8cf354 100644
--- a/src/backend/distributed/deparser/ruleutils_17.c
+++ b/src/backend/distributed/deparser/ruleutils_17.c
@@ -3542,9 +3542,6 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 	SubLink    *cur_ma_sublink;
 	List	   *ma_sublinks;
 
-	AttrNumber	previous_attnum = InvalidAttrNumber;
-	int 		paramid_increment = 0;
-
 	/*
 	 * Prepare to deal with MULTIEXPR assignments: collect the source SubLinks
 	 * into a list.  We expect them to appear, in ID order, in resjunk tlist
@@ -3568,6 +3565,8 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 				}
 			}
 		}
+
+		ensure_update_targetlist_in_param_order(targetList);
 	}
 	next_ma_cell = list_head(ma_sublinks);
 	cur_ma_sublink = NULL;
@@ -3667,48 +3666,11 @@ get_update_query_targetlist_def(Query *query, List *targetList,
 		 */
 		if (cur_ma_sublink != NULL)
 		{
-			AttrNumber attnum = InvalidAttrNumber;
-			if (IsA(expr, Param))
-			{
-				Param *param = (Param *) expr;
-				attnum = param->paramid + paramid_increment;
-			}
-			else if (IsA(expr, FuncExpr))
-			{
-				FuncExpr *func = (FuncExpr *) expr;
-				ListCell *lc;
-
-				/* Iterate through the arguments of the FuncExpr */
-				foreach(lc, func->args)
-				{
-					Node *arg = (Node *) lfirst(lc);
-
-					/* Check if the argument is a PARAM node */
-					if (IsA(arg, Param))
-					{
-						Param *param = (Param *) arg;
-						attnum = param->paramid + paramid_increment;
-
-						break;  /* Exit loop once we find the PARAM node */
-					}
-				}
-			}
-
-			if (previous_attnum >= attnum)
-				ereport(ERROR, (errcode(ERRCODE_WRONG_OBJECT_TYPE),
-								errmsg(
-									"cannot plan distributed UPDATE SET (..) = (SELECT ...) query when not sorted by physical order"),
-								errhint("Sort the columns on the left side by physical order.")));
-
-			previous_attnum = attnum;
-
 			if (--remaining_ma_columns > 0)
 				continue;		/* not the last column of multiassignment */
-
 			appendStringInfoChar(buf, ')');
 			expr = (Node *) cur_ma_sublink;
 			cur_ma_sublink = NULL;
-			paramid_increment = previous_attnum;
 		}
 
 		appendStringInfoString(buf, " = ");
diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c
index f8f0b8b8c..db5c3d4ff 100644
--- a/src/backend/distributed/planner/insert_select_planner.c
+++ b/src/backend/distributed/planner/insert_select_planner.c
@@ -1082,11 +1082,6 @@ ReorderInsertSelectTargetLists(Query *originalQuery, RangeTblEntry *insertRte,
 		TargetEntry *newSubqueryTargetEntry = NULL;
 		AttrNumber originalAttrNo = get_attnum(insertRelationId,
 											   oldInsertTargetEntry->resname);
-		Node *expr;
-
-		/* we need to explore the underlying expression */
-		expr = (Node *) oldInsertTargetEntry->expr;
-		expr = strip_implicit_coercions(expr);
 
 		/* we need to explore the underlying expression */
 		Node *expr = strip_implicit_coercions((Node *) oldInsertTargetEntry->expr);
diff --git a/src/include/distributed/citus_ruleutils.h b/src/include/distributed/citus_ruleutils.h
index 9089ca3a7..8a0f76c8a 100644
--- a/src/include/distributed/citus_ruleutils.h
+++ b/src/include/distributed/citus_ruleutils.h
@@ -60,8 +60,6 @@ extern char * generate_operator_name(Oid operid, Oid arg1, Oid arg2);
 extern List * getOwnedSequences_internal(Oid relid, AttrNumber attnum, char deptype);
 extern void AppendOptionListToString(StringInfo stringData, List *options);
 
-extern int GetParamId(Node *expr);
-extern int target_list_cmp(const ListCell *a, const ListCell *b);
-
+extern void ensure_update_targetlist_in_param_order(List *targetList);
 
 #endif /* CITUS_RULEUTILS_H */
diff --git a/src/test/regress/citus_tests/run_test.py b/src/test/regress/citus_tests/run_test.py
index 9a648c0ab..80d045d94 100755
--- a/src/test/regress/citus_tests/run_test.py
+++ b/src/test/regress/citus_tests/run_test.py
@@ -147,6 +147,7 @@ DEPS = {
     "multi_mx_modifying_xacts": TestDeps(None, ["multi_mx_create_table"]),
     "multi_mx_router_planner": TestDeps(None, ["multi_mx_create_table"]),
     "multi_mx_copy_data": TestDeps(None, ["multi_mx_create_table"]),
+    "multi_mx_modifications": TestDeps(None, ["multi_mx_create_table"]),
     "multi_mx_schema_support": TestDeps(None, ["multi_mx_copy_data"]),
     "multi_simple_queries": TestDeps("base_schedule"),
     "create_single_shard_table": TestDeps("minimal_schedule"),
diff --git a/src/test/regress/expected/distributed_types.out b/src/test/regress/expected/distributed_types.out
index 68b952aae..703614d61 100644
--- a/src/test/regress/expected/distributed_types.out
+++ b/src/test/regress/expected/distributed_types.out
@@ -444,7 +444,6 @@ HINT:  Use the column name to insert or update the composite type as a single va
 UPDATE domain_indirection_test SET domain_array[0].if2 = (SELECT 5);
 ERROR:  inserting or modifying composite type fields is not supported
 HINT:  Use the column name to insert or update the composite type as a single value
-HINT:  Use the column name to insert or update the composite type as a single value
 -- below are supported as we don't do any field indirection
 INSERT INTO field_indirection_test_2 (ct2_col, int_col, ct1_col)
 VALUES ('(1, "text1", 2)', 3, '(4, 5)'), ('(6, "text2", 7)', 8, '(9, 10)');
diff --git a/src/test/regress/expected/multi_modifications.out b/src/test/regress/expected/multi_modifications.out
index 1a2c01a4d..cebef0526 100644
--- a/src/test/regress/expected/multi_modifications.out
+++ b/src/test/regress/expected/multi_modifications.out
@@ -935,10 +935,10 @@ SELECT * FROM summary_table ORDER BY id;
   2 |       400 | 450.0000000000000000 |       |
 (2 rows)
 
---- TODO this one is a silent corruption:
--- UPDATE summary_table SET (average_value, min_value) =
--- 	(SELECT avg(value), min(value) FROM raw_table WHERE id = 2)
--- WHERE id = 2;
+-- try different order of update targets
+UPDATE summary_table SET (average_value, min_value) =
+ 	(SELECT avg(value), min(value) FROM raw_table WHERE id = 2)
+WHERE id = 2;
 SELECT * FROM summary_table ORDER BY id;
  id | min_value |    average_value     | count | uniques
 ---------------------------------------------------------------------
@@ -1122,10 +1122,9 @@ SELECT * FROM reference_summary_table ORDER BY id;
   2 |       400 | 450.0000000000000000 |       |
 (2 rows)
 
---- TODO this one is a silent corruption:
--- UPDATE reference_summary_table SET (average_value, min_value) =
--- 	(SELECT avg(value), min(value) FROM reference_raw_table WHERE id = 2)
--- WHERE id = 2;
+UPDATE reference_summary_table SET (average_value, min_value) =
+	(SELECT avg(value), min(value) FROM reference_raw_table WHERE id = 2)
+WHERE id = 2;
 SELECT * FROM reference_summary_table ORDER BY id;
  id | min_value |    average_value     | count | uniques
 ---------------------------------------------------------------------
@@ -1359,5 +1358,14 @@ DROP TABLE raw_table;
 DROP TABLE summary_table;
 DROP TABLE reference_raw_table;
 DROP TABLE reference_summary_table;
+DROP TABLE limit_orders;
+DROP TABLE multiple_hash;
+DROP TABLE range_partitioned;
+DROP TABLE append_partitioned;
+DROP TABLE bidders;
+DROP FUNCTION stable_append;
+DROP FUNCTION immutable_append;
+DROP FUNCTION temp_strict_func;
+DROP TYPE order_side;
 DROP SCHEMA multi_modifications CASCADE;
 NOTICE:  drop cascades to table multi_modifications.local
diff --git a/src/test/regress/expected/multi_mx_modifications.out b/src/test/regress/expected/multi_mx_modifications.out
index b401a67e3..9e053d3f2 100644
--- a/src/test/regress/expected/multi_mx_modifications.out
+++ b/src/test/regress/expected/multi_mx_modifications.out
@@ -91,8 +91,6 @@ SET client_min_messages TO DEFAULT;
 -- commands with non-constant partition values are unsupported
 INSERT INTO limit_orders_mx VALUES (random() * 100, 'ORCL', 152, '2011-08-25 11:50:45',
 								 'sell', 0.58);
-INSERT INTO limit_orders_mx SELECT random() * 100, 'ORCL', 152, '2011-08-25 11:50:45',
-								 'sell', 0.58;
 -- values for other columns are totally fine
 INSERT INTO limit_orders_mx VALUES (2036, 'GOOG', 5634, now(), 'buy', random());
 -- commands with mutable functions in their quals
@@ -105,22 +103,6 @@ DELETE FROM limit_orders_mx WHERE id = 246 AND placed_at = current_timestamp::ti
 INSERT INTO limit_orders_mx VALUES (2037, 'GOOG', 5634, now(), 'buy', random()),
                                    (2038, 'GOOG', 5634, now(), 'buy', random()),
                                    (2039, 'GOOG', 5634, now(), 'buy', random());
-SELECT count(distinct placed_at) FROM limit_orders_mx WHERE id BETWEEN 2037 AND 2039;
- count
----------------------------------------------------------------------
-     1
-(1 row)
-
-INSERT INTO limit_orders_mx SELECT * FROM (VALUES
-                                   (2040, 'GOOG', 5634, now(), 'buy'::citus_mx_test_schema.order_side_mx, random()),
-                                   (2041, 'GOOG', 5634, now(), 'buy'::citus_mx_test_schema.order_side_mx, random()),
-                                   (2042, 'GOOG', 5634, now(), 'buy'::citus_mx_test_schema.order_side_mx, random()));
-SELECT count(distinct placed_at) FROM limit_orders_mx WHERE id BETWEEN 2040 AND 2042;
- count
----------------------------------------------------------------------
-     1
-(1 row)
-
 -- connect back to the other node
 \c - - - :worker_1_port
 -- commands containing a CTE are supported
@@ -234,12 +216,9 @@ INSERT INTO limit_orders_mx VALUES (275, 'ADR', 140, '2007-07-02 16:32:15', 'sel
 ERROR:  duplicate key value violates unique constraint "limit_orders_mx_pkey_1220093"
 -- multi shard update is supported
 UPDATE limit_orders_mx SET limit_price = 0.00;
-UPDATE limit_orders_mx SET limit_price = (SELECT 0.00);
 -- attempting to change the partition key is unsupported
 UPDATE limit_orders_mx SET id = 0 WHERE id = 246;
 ERROR:  modifying the partition value of rows is not allowed
-UPDATE limit_orders_mx SET id = (SELECT 0) WHERE id = 246;
-ERROR:  modifying the partition value of rows is not allowed
 -- UPDATEs with a FROM clause are unsupported
 UPDATE limit_orders_mx SET limit_price = 0.00 FROM bidders
 					WHERE limit_orders_mx.id = 246 AND
diff --git a/src/test/regress/expected/indirections.out b/src/test/regress/expected/multi_update_select.out
similarity index 84%
rename from src/test/regress/expected/indirections.out
rename to src/test/regress/expected/multi_update_select.out
index 115f5a7b4..bd0267701 100644
--- a/src/test/regress/expected/indirections.out
+++ b/src/test/regress/expected/multi_update_select.out
@@ -1,8 +1,6 @@
-SET citus.shard_count TO 2;
-SET citus.next_shard_id TO 750000;
-SET citus.next_placement_id TO 750000;
-CREATE SCHEMA indirections;
-SET search_path TO indirections;
+CREATE SCHEMA multi_update_select;
+SET search_path TO multi_update_select;
+SET citus.next_shard_id TO 751000;
 -- specific tests related to get_update_query_targetlist_def
 -- we test only queries with sublinks, like:
 -- ( ... SET (...) = (SELECT ...))
@@ -11,7 +9,7 @@ CREATE TABLE test_ref_indirection (
     id bigint primary key
   , col_bool bool , col_date date , col_int integer , col_text text
   );
-SELECT create_reference_table('indirections.test_ref_indirection');
+SELECT create_reference_table('test_ref_indirection');
  create_reference_table
 ---------------------------------------------------------------------
 
@@ -21,7 +19,7 @@ CREATE TABLE test_ref_indirection_new (
     id bigint primary key
   , col_bool bool , col_date date , col_int integer , col_text text
   );
-SELECT create_reference_table('indirections.test_ref_indirection_new');
+SELECT create_reference_table('test_ref_indirection_new');
  create_reference_table
 ---------------------------------------------------------------------
 
@@ -32,7 +30,7 @@ CREATE TABLE test_dist_indirection (
     id bigint primary key
   , col_bool bool , col_date date , col_int integer , col_text text
   );
-SELECT create_distributed_table('indirections.test_dist_indirection', 'id');
+SELECT create_distributed_table('test_dist_indirection', 'id');
  create_distributed_table
 ---------------------------------------------------------------------
 
@@ -42,13 +40,12 @@ CREATE TABLE test_dist_indirection_new (
     id bigint primary key
   , col_bool bool , col_date date , col_int integer , col_text text
   );
-SELECT create_distributed_table('indirections.test_dist_indirection_new', 'id');
+SELECT create_distributed_table('test_dist_indirection_new', 'id');
  create_distributed_table
 ---------------------------------------------------------------------
 
 (1 row)
 
--- Local tables required ?
 -- those should work:
 INSERT INTO test_ref_indirection (id, col_bool, col_date, col_int, col_text)
     SELECT 1, true, '1970-01-01'::date, 1, 'one';
@@ -433,15 +430,76 @@ CREATE TABLE update_test (
     b   INT,
     c   TEXT
 );
-SELECT create_reference_table('indirections.update_test');
+SELECT create_reference_table('update_test');
  create_reference_table
 ---------------------------------------------------------------------
 
 (1 row)
 
+INSERT INTO update_test VALUES (11, 41, 'car');
+INSERT INTO update_test VALUES (100, 20, 'bike');
+INSERT INTO update_test VALUES (100, 20, 'tractor');
+SELECT * FROM update_test;
+  a  | b  |    c
+---------------------------------------------------------------------
+  11 | 41 | car
+ 100 | 20 | bike
+ 100 | 20 | tractor
+(3 rows)
+
 UPDATE update_test
 SET (b,a) = (select a,b from update_test where b = 41 and c = 'car')
-WHERE a = 100 AND b = 20;
--- https://github.com/citusdata/citus/pull/5692
-set client_min_messages to ERROR;
-DROP SCHEMA indirections CASCADE;
+WHERE a = 100 AND b = 20
+RETURNING *;
+ a  | b  |    c
+---------------------------------------------------------------------
+ 41 | 11 | bike
+ 41 | 11 | tractor
+(2 rows)
+
+-- Test that multiple out of order columns and multiple sublinks are handled correctly.
+CREATE TABLE upd2_test (a1 int, b1 int, c1 int, d1 int, e1 int, f1 int, g1 int);
+SELECT create_reference_table('upd2_test');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+INSERT INTO upd2_test SELECT 1, 1, 1, 1, 1, 1, 1 FROM generate_series(1,5) c(i);
+UPDATE upd2_test set (b1, a1) = (SELECT 200, 100), (g1, f1, e1) = (SELECT 700, 600, 500), (d1, c1) = (SELECT 400, 300);
+SELECT * FROM upd2_test;
+ a1  | b1  | c1  | d1  | e1  | f1  | g1
+---------------------------------------------------------------------
+ 100 | 200 | 300 | 400 | 500 | 600 | 700
+ 100 | 200 | 300 | 400 | 500 | 600 | 700
+ 100 | 200 | 300 | 400 | 500 | 600 | 700
+ 100 | 200 | 300 | 400 | 500 | 600 | 700
+ 100 | 200 | 300 | 400 | 500 | 600 | 700
+(5 rows)
+
+UPDATE upd2_test set (g1, a1) = (SELECT 77, 11), (f1, b1) = (SELECT 66, 22), (e1, c1) = (SELECT 55, 33), (d1) = (SELECT 44);
+SELECT * FROM upd2_test;
+ a1 | b1 | c1 | d1 | e1 | f1 | g1
+---------------------------------------------------------------------
+ 11 | 22 | 33 | 44 | 55 | 66 | 77
+ 11 | 22 | 33 | 44 | 55 | 66 | 77
+ 11 | 22 | 33 | 44 | 55 | 66 | 77
+ 11 | 22 | 33 | 44 | 55 | 66 | 77
+ 11 | 22 | 33 | 44 | 55 | 66 | 77
+(5 rows)
+
+UPDATE upd2_test set (g1, a1) = (SELECT 7, 1), (f1) = (SELECT 6), (c1, e1) = (SELECT 3, 5), (b1) = (SELECT 2), (d1) = (SELECT 4);
+SELECT * FROM upd2_test;
+ a1 | b1 | c1 | d1 | e1 | f1 | g1
+---------------------------------------------------------------------
+  1 |  2 |  3 |  4 |  5 |  6 |  7
+  1 |  2 |  3 |  4 |  5 |  6 |  7
+  1 |  2 |  3 |  4 |  5 |  6 |  7
+  1 |  2 |  3 |  4 |  5 |  6 |  7
+  1 |  2 |  3 |  4 |  5 |  6 |  7
+(5 rows)
+
+-- suppress cascade messages
+SET client_min_messages to ERROR;
+DROP SCHEMA multi_update_select CASCADE;
+RESET client_min_messages;
diff --git a/src/test/regress/multi_1_schedule b/src/test/regress/multi_1_schedule
index 39a9e4070..c31ee0b07 100644
--- a/src/test/regress/multi_1_schedule
+++ b/src/test/regress/multi_1_schedule
@@ -204,6 +204,7 @@ test: multi_outer_join
 # ---
 test: multi_complex_count_distinct multi_select_distinct
 test: multi_modifications
+test: multi_update_select
 test: multi_distribution_metadata
 test: multi_prune_shard_list
 test: multi_upsert multi_simple_queries multi_data_types
diff --git a/src/test/regress/sql/multi_modifications.sql b/src/test/regress/sql/multi_modifications.sql
index eeab5536f..958791e44 100644
--- a/src/test/regress/sql/multi_modifications.sql
+++ b/src/test/regress/sql/multi_modifications.sql
@@ -588,10 +588,10 @@ WHERE id = 2;
 
 SELECT * FROM summary_table ORDER BY id;
 
---- TODO this one is a silent corruption:
--- UPDATE summary_table SET (average_value, min_value) =
--- 	(SELECT avg(value), min(value) FROM raw_table WHERE id = 2)
--- WHERE id = 2;
+-- try different order of update targets
+UPDATE summary_table SET (average_value, min_value) =
+ 	(SELECT avg(value), min(value) FROM raw_table WHERE id = 2)
+WHERE id = 2;
 
 SELECT * FROM summary_table ORDER BY id;
 
@@ -723,10 +723,9 @@ WHERE id = 2;
 
 SELECT * FROM reference_summary_table ORDER BY id;
 
---- TODO this one is a silent corruption:
--- UPDATE reference_summary_table SET (average_value, min_value) =
--- 	(SELECT avg(value), min(value) FROM reference_raw_table WHERE id = 2)
--- WHERE id = 2;
+UPDATE reference_summary_table SET (average_value, min_value) =
+	(SELECT avg(value), min(value) FROM reference_raw_table WHERE id = 2)
+WHERE id = 2;
 
 SELECT * FROM reference_summary_table ORDER BY id;
 
@@ -920,4 +919,15 @@ DROP TABLE raw_table;
 DROP TABLE summary_table;
 DROP TABLE reference_raw_table;
 DROP TABLE reference_summary_table;
+DROP TABLE limit_orders;
+DROP TABLE multiple_hash;
+DROP TABLE range_partitioned;
+DROP TABLE append_partitioned;
+DROP TABLE bidders;
+
+DROP FUNCTION stable_append;
+DROP FUNCTION immutable_append;
+DROP FUNCTION temp_strict_func;
+DROP TYPE order_side;
+
 DROP SCHEMA multi_modifications CASCADE;
diff --git a/src/test/regress/sql/multi_mx_modifications.sql b/src/test/regress/sql/multi_mx_modifications.sql
index 185fef451..852bf3a42 100644
--- a/src/test/regress/sql/multi_mx_modifications.sql
+++ b/src/test/regress/sql/multi_mx_modifications.sql
@@ -67,8 +67,6 @@ SET client_min_messages TO DEFAULT;
 -- commands with non-constant partition values are unsupported
 INSERT INTO limit_orders_mx VALUES (random() * 100, 'ORCL', 152, '2011-08-25 11:50:45',
 								 'sell', 0.58);
-INSERT INTO limit_orders_mx SELECT random() * 100, 'ORCL', 152, '2011-08-25 11:50:45',
-								 'sell', 0.58;
 
 -- values for other columns are totally fine
 INSERT INTO limit_orders_mx VALUES (2036, 'GOOG', 5634, now(), 'buy', random());
@@ -85,17 +83,6 @@ INSERT INTO limit_orders_mx VALUES (2037, 'GOOG', 5634, now(), 'buy', random()),
                                    (2038, 'GOOG', 5634, now(), 'buy', random()),
                                    (2039, 'GOOG', 5634, now(), 'buy', random());
 
-SELECT count(distinct placed_at) FROM limit_orders_mx WHERE id BETWEEN 2037 AND 2039;
-
-INSERT INTO limit_orders_mx SELECT * FROM (VALUES
-                                   (2040, 'GOOG', 5634, now(), 'buy'::citus_mx_test_schema.order_side_mx, random()),
-                                   (2041, 'GOOG', 5634, now(), 'buy'::citus_mx_test_schema.order_side_mx, random()),
-                                   (2042, 'GOOG', 5634, now(), 'buy'::citus_mx_test_schema.order_side_mx, random()));
-
-SELECT count(distinct placed_at) FROM limit_orders_mx WHERE id BETWEEN 2040 AND 2042;
-
-
-
 -- connect back to the other node
 \c - - - :worker_1_port
 
@@ -166,11 +153,9 @@ INSERT INTO limit_orders_mx VALUES (275, 'ADR', 140, '2007-07-02 16:32:15', 'sel
 
 -- multi shard update is supported
 UPDATE limit_orders_mx SET limit_price = 0.00;
-UPDATE limit_orders_mx SET limit_price = (SELECT 0.00);
 
 -- attempting to change the partition key is unsupported
 UPDATE limit_orders_mx SET id = 0 WHERE id = 246;
-UPDATE limit_orders_mx SET id = (SELECT 0) WHERE id = 246;
 
 -- UPDATEs with a FROM clause are unsupported
 UPDATE limit_orders_mx SET limit_price = 0.00 FROM bidders
diff --git a/src/test/regress/sql/indirections.sql b/src/test/regress/sql/multi_update_select.sql
similarity index 80%
rename from src/test/regress/sql/indirections.sql
rename to src/test/regress/sql/multi_update_select.sql
index 598124208..4891eb330 100644
--- a/src/test/regress/sql/indirections.sql
+++ b/src/test/regress/sql/multi_update_select.sql
@@ -1,9 +1,7 @@
-SET citus.shard_count TO 2;
-SET citus.next_shard_id TO 750000;
-SET citus.next_placement_id TO 750000;
+CREATE SCHEMA multi_update_select;
+SET search_path TO multi_update_select;
 
-CREATE SCHEMA indirections;
-SET search_path TO indirections;
+SET citus.next_shard_id TO 751000;
 
 -- specific tests related to get_update_query_targetlist_def
 -- we test only queries with sublinks, like:
@@ -14,28 +12,26 @@ CREATE TABLE test_ref_indirection (
     id bigint primary key
   , col_bool bool , col_date date , col_int integer , col_text text
   );
-SELECT create_reference_table('indirections.test_ref_indirection');
+SELECT create_reference_table('test_ref_indirection');
 
 CREATE TABLE test_ref_indirection_new (
     id bigint primary key
   , col_bool bool , col_date date , col_int integer , col_text text
   );
-SELECT create_reference_table('indirections.test_ref_indirection_new');
+SELECT create_reference_table('test_ref_indirection_new');
 
 -- Distributed tables
 CREATE TABLE test_dist_indirection (
     id bigint primary key
   , col_bool bool , col_date date , col_int integer , col_text text
   );
-SELECT create_distributed_table('indirections.test_dist_indirection', 'id');
+SELECT create_distributed_table('test_dist_indirection', 'id');
 
 CREATE TABLE test_dist_indirection_new (
     id bigint primary key
   , col_bool bool , col_date date , col_int integer , col_text text
   );
-SELECT create_distributed_table('indirections.test_dist_indirection_new', 'id');
-
--- Local tables required ?
+SELECT create_distributed_table('test_dist_indirection_new', 'id');
 
 -- those should work:
 INSERT INTO test_ref_indirection (id, col_bool, col_date, col_int, col_text)
@@ -217,12 +213,35 @@ CREATE TABLE update_test (
     b   INT,
     c   TEXT
 );
-SELECT create_reference_table('indirections.update_test');
+
+SELECT create_reference_table('update_test');
+INSERT INTO update_test VALUES (11, 41, 'car');
+INSERT INTO update_test VALUES (100, 20, 'bike');
+INSERT INTO update_test VALUES (100, 20, 'tractor');
+SELECT * FROM update_test;
+
 UPDATE update_test
 SET (b,a) = (select a,b from update_test where b = 41 and c = 'car')
-WHERE a = 100 AND b = 20;
+WHERE a = 100 AND b = 20
+RETURNING *;
 
--- https://github.com/citusdata/citus/pull/5692
-set client_min_messages to ERROR;
-DROP SCHEMA indirections CASCADE;
+-- Test that multiple out of order columns and multiple sublinks are handled correctly.
+CREATE TABLE upd2_test (a1 int, b1 int, c1 int, d1 int, e1 int, f1 int, g1 int);
+SELECT create_reference_table('upd2_test');
+
+INSERT INTO upd2_test SELECT 1, 1, 1, 1, 1, 1, 1 FROM generate_series(1,5) c(i);
+
+UPDATE upd2_test set (b1, a1) = (SELECT 200, 100), (g1, f1, e1) = (SELECT 700, 600, 500), (d1, c1) = (SELECT 400, 300);
+SELECT * FROM upd2_test;
+
+UPDATE upd2_test set (g1, a1) = (SELECT 77, 11), (f1, b1) = (SELECT 66, 22), (e1, c1) = (SELECT 55, 33), (d1) = (SELECT 44);
+SELECT * FROM upd2_test;
+
+UPDATE upd2_test set (g1, a1) = (SELECT 7, 1), (f1) = (SELECT 6), (c1, e1) = (SELECT 3, 5), (b1) = (SELECT 2), (d1) = (SELECT 4);
+SELECT * FROM upd2_test;
+
+-- suppress cascade messages
+SET client_min_messages to ERROR;
+DROP SCHEMA multi_update_select CASCADE;
+RESET client_min_messages;