Refactor logic in citus_ruleutils, Run ci tools, fix up regress tests

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);
+	}
+}

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, " = ");

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;
-			}
 		}
 
-		/*
+	        ensure_update_targetlist_in_param_order(targetList);
-		 * 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;
@@ -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);
 }
 
 /* ----------

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, " = ");

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);

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 void ensure_update_targetlist_in_param_order(List *targetList);
-extern int target_list_cmp(const ListCell *a, const ListCell *b);
-
 
 #endif /* CITUS_RULEUTILS_H */

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"),

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)');

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:
+-- try different order of update targets
--- UPDATE summary_table SET (average_value, min_value) =
+UPDATE summary_table SET (average_value, min_value) =
--- 	(SELECT avg(value), min(value) FROM raw_table WHERE id = 2)
+ 	(SELECT avg(value), min(value) FROM raw_table WHERE id = 2)
--- 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) =
--- UPDATE reference_summary_table SET (average_value, min_value) =
+	(SELECT avg(value), min(value) FROM reference_raw_table WHERE id = 2)
--- 	(SELECT avg(value), min(value) FROM reference_raw_table WHERE id = 2)
+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

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

src/test/regress/expected/multi_update_select.out

@@ -1,8 +1,6 @@
-SET citus.shard_count TO 2;
+CREATE SCHEMA multi_update_select;
-SET citus.next_shard_id TO 750000;
+SET search_path TO multi_update_select;
-SET citus.next_placement_id TO 750000;
+SET citus.next_shard_id TO 751000;
-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 ...))
@@ -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;
+WHERE a = 100 AND b = 20
--- https://github.com/citusdata/citus/pull/5692
+RETURNING *;
-set client_min_messages to ERROR;
+ a  | b  |    c
-DROP SCHEMA indirections CASCADE;
+---------------------------------------------------------------------
+ 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;

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

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:
+-- try different order of update targets
--- UPDATE summary_table SET (average_value, min_value) =
+UPDATE summary_table SET (average_value, min_value) =
--- 	(SELECT avg(value), min(value) FROM raw_table WHERE id = 2)
+ 	(SELECT avg(value), min(value) FROM raw_table WHERE id = 2)
--- 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) =
--- UPDATE reference_summary_table SET (average_value, min_value) =
+	(SELECT avg(value), min(value) FROM reference_raw_table WHERE id = 2)
--- 	(SELECT avg(value), min(value) FROM reference_raw_table WHERE id = 2)
+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;

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

src/test/regress/sql/multi_update_select.sql

@@ -1,9 +1,7 @@
-SET citus.shard_count TO 2;
+CREATE SCHEMA multi_update_select;
-SET citus.next_shard_id TO 750000;
+SET search_path TO multi_update_select;
-SET citus.next_placement_id TO 750000;
 
-CREATE SCHEMA indirections;
+SET citus.next_shard_id TO 751000;
-SET search_path TO indirections;
 
 -- 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');
+SELECT create_distributed_table('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)
@@ -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
+-- Test that multiple out of order columns and multiple sublinks are handled correctly.
-set client_min_messages to ERROR;
+CREATE TABLE upd2_test (a1 int, b1 int, c1 int, d1 int, e1 int, f1 int, g1 int);
-DROP SCHEMA indirections CASCADE;
+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;