Add tests verifying that updates return correct tuple counts.

This unfortunately requires adding a new table, triggering renumbering
of a number of shard ids.

src/test/regress/expected/multi_modifications.out

@@ -12,6 +12,10 @@ CREATE TABLE limit_orders (
 	kind order_side NOT NULL,
 	limit_price decimal NOT NULL DEFAULT 0.00 CHECK (limit_price >= 0.00)
 );
+CREATE TABLE multiple_hash (
+	category text NOT NULL,
+	data text NOT NULL
+);
 CREATE TABLE insufficient_shards ( LIKE limit_orders );
 CREATE TABLE range_partitioned ( LIKE limit_orders );
 CREATE TABLE append_partitioned ( LIKE limit_orders );
@@ -21,6 +25,12 @@ SELECT master_create_distributed_table('limit_orders', 'id', 'hash');
  
 (1 row)
 
+SELECT master_create_distributed_table('multiple_hash', 'category', 'hash');
+ master_create_distributed_table 
+---------------------------------
+ 
+(1 row)
+
 SELECT master_create_distributed_table('insufficient_shards', 'id', 'hash');
  master_create_distributed_table 
 ---------------------------------
@@ -45,6 +55,12 @@ SELECT master_create_worker_shards('limit_orders', 2, 2);
  
 (1 row)
 
+SELECT master_create_worker_shards('multiple_hash', 2, 2);
+ master_create_worker_shards 
+-----------------------------
+ 
+(1 row)
+
 -- make a single shard that covers no partition values
 SELECT master_create_worker_shards('insufficient_shards', 1, 1);
  master_create_worker_shards 
@@ -95,7 +111,7 @@ INSERT INTO append_partitioned VALUES (414123, 'AAPL', 9580, '2004-10-19 10:23:5
 SET client_min_messages TO 'DEBUG2';
 SET citus.task_executor_type TO 'real-time';
 SELECT * FROM range_partitioned WHERE id = 32743;
-DEBUG:  predicate pruning for shardId 750004
+DEBUG:  predicate pruning for shardId 750006
 DEBUG:  Plan is router executable
   id   | symbol | bidder_id |        placed_at         | kind | limit_price 
 -------+--------+-----------+--------------------------+------+-------------
@@ -103,7 +119,7 @@ DEBUG:  Plan is router executable
 (1 row)
 
 SELECT * FROM append_partitioned WHERE id = 414123;
-DEBUG:  predicate pruning for shardId 750006
+DEBUG:  predicate pruning for shardId 750008
 DEBUG:  Plan is router executable
    id   | symbol | bidder_id |        placed_at         | kind | limit_price 
 --------+--------+-----------+--------------------------+------+-------------
@@ -380,3 +396,46 @@ ERROR:  functions used in modification queries on distributed tables must be mar
 -- cursors are not supported
 UPDATE limit_orders SET symbol = 'GM' WHERE CURRENT OF cursor_name;
 ERROR:  distributed modifications must target exactly one shard
+-- ensure returned row counters are correct
+\set QUIET off
+INSERT INTO multiple_hash VALUES ('1', '1');
+INSERT 0 1
+INSERT INTO multiple_hash VALUES ('1', '2');
+INSERT 0 1
+INSERT INTO multiple_hash VALUES ('1', '3');
+INSERT 0 1
+INSERT INTO multiple_hash VALUES ('2', '1');
+INSERT 0 1
+INSERT INTO multiple_hash VALUES ('2', '2');
+INSERT 0 1
+INSERT INTO multiple_hash VALUES ('2', '3');
+INSERT 0 1
+-- check that update return the right number of rows
+-- one row
+UPDATE multiple_hash SET data = data ||'-1' WHERE category = '1' AND data = '1';
+UPDATE 1
+-- three rows
+UPDATE multiple_hash SET data = data ||'-2' WHERE category = '1';
+UPDATE 3
+-- check
+SELECT * FROM multiple_hash WHERE category = '1' ORDER BY category, data;
+ category | data  
+----------+-------
+ 1        | 1-1-2
+ 1        | 2-2
+ 1        | 3-2
+(3 rows)
+
+-- check that deletes return the right number of rows
+-- one row
+DELETE FROM multiple_hash WHERE category = '2' AND data = '1';
+DELETE 1
+-- two rows
+DELETE FROM multiple_hash WHERE category = '2';
+DELETE 2
+-- check
+SELECT * FROM multiple_hash WHERE category = '2' ORDER BY category, data;
+ category | data 
+----------+------
+(0 rows)
+

src/test/regress/sql/multi_modifications.sql

@@ -18,16 +18,24 @@ CREATE TABLE limit_orders (
 	limit_price decimal NOT NULL DEFAULT 0.00 CHECK (limit_price >= 0.00)
 );
 
+
+CREATE TABLE multiple_hash (
+	category text NOT NULL,
+	data text NOT NULL
+);
+
 CREATE TABLE insufficient_shards ( LIKE limit_orders );
 CREATE TABLE range_partitioned ( LIKE limit_orders );
 CREATE TABLE append_partitioned ( LIKE limit_orders );
 
 SELECT master_create_distributed_table('limit_orders', 'id', 'hash');
+SELECT master_create_distributed_table('multiple_hash', 'category', 'hash');
 SELECT master_create_distributed_table('insufficient_shards', 'id', 'hash');
 SELECT master_create_distributed_table('range_partitioned', 'id', 'range');
 SELECT master_create_distributed_table('append_partitioned', 'id', 'append');
 
 SELECT master_create_worker_shards('limit_orders', 2, 2);
+SELECT master_create_worker_shards('multiple_hash', 2, 2);
 
 -- make a single shard that covers no partition values
 SELECT master_create_worker_shards('insufficient_shards', 1, 1);
@@ -292,3 +300,29 @@ UPDATE limit_orders SET placed_at = now() WHERE id = 246;
 
 -- cursors are not supported
 UPDATE limit_orders SET symbol = 'GM' WHERE CURRENT OF cursor_name;
+
+
+-- ensure returned row counters are correct
+\set QUIET off
+INSERT INTO multiple_hash VALUES ('1', '1');
+INSERT INTO multiple_hash VALUES ('1', '2');
+INSERT INTO multiple_hash VALUES ('1', '3');
+INSERT INTO multiple_hash VALUES ('2', '1');
+INSERT INTO multiple_hash VALUES ('2', '2');
+INSERT INTO multiple_hash VALUES ('2', '3');
+
+-- check that update return the right number of rows
+-- one row
+UPDATE multiple_hash SET data = data ||'-1' WHERE category = '1' AND data = '1';
+-- three rows
+UPDATE multiple_hash SET data = data ||'-2' WHERE category = '1';
+-- check
+SELECT * FROM multiple_hash WHERE category = '1' ORDER BY category, data;
+
+-- check that deletes return the right number of rows
+-- one row
+DELETE FROM multiple_hash WHERE category = '2' AND data = '1';
+-- two rows
+DELETE FROM multiple_hash WHERE category = '2';
+-- check
+SELECT * FROM multiple_hash WHERE category = '2' ORDER BY category, data;