Evaluate nextval in the target list on the coordinator

src/backend/distributed/executor/insert_select_executor.c

@@ -1059,6 +1059,17 @@ IsRedistributablePlan(Plan *selectPlan)
 		return false;
 	}
 
+	if (distSelectPlan->masterQuery != NULL)
+	{
+		Query *masterQuery = (Query *) distSelectPlan->masterQuery;
+
+		if (contain_nextval_expression_walker((Node *) masterQuery->targetList, NULL))
+		{
+			/* nextval needs to be evaluated on the coordinator */
+			return false;
+		}
+	}
+
 	return true;
 }
 

src/backend/distributed/planner/insert_select_planner.c

@@ -1164,18 +1164,6 @@ CoordinatorInsertSelectSupported(Query *insertSelectQuery)
 							 "not supported", NULL, NULL);
 	}
 
-	RangeTblEntry *subqueryRte = ExtractSelectRangeTableEntry(insertSelectQuery);
-	Query *subquery = (Query *) subqueryRte->subquery;
-
-	if (NeedsDistributedPlanning(subquery) &&
-		contain_nextval_expression_walker((Node *) insertSelectQuery->targetList, NULL))
-	{
-		return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED,
-							 "INSERT ... SELECT cannot generate sequence values when "
-							 "selecting from a distributed table",
-							 NULL, NULL);
-	}
-
 	return NULL;
 }
 

src/backend/distributed/planner/multi_logical_optimizer.c

@@ -3589,17 +3589,22 @@ static bool
 CanPushDownExpression(Node *expression,
 					  const ExtendedOpNodeProperties *extendedOpNodeProperties)
 {
+	if (contain_nextval_expression_walker(expression, NULL))
+	{
+		/* nextval can only be evaluated on the coordinator */
+		return false;
+	}
+
 	bool hasAggregate = contain_aggs_of_level(expression, 0);
 	bool hasWindowFunction = contain_window_function(expression);
-	bool hasPushableWindowFunction =
-		hasWindowFunction && extendedOpNodeProperties->onlyPushableWindowFunctions;
-
 	if (!hasAggregate && !hasWindowFunction)
 	{
 		return true;
 	}
 
 	/* aggregates inside pushed down window functions can be pushed down */
+	bool hasPushableWindowFunction =
+		hasWindowFunction && extendedOpNodeProperties->onlyPushableWindowFunctions;
 	if (hasPushableWindowFunction)
 	{
 		return true;

src/backend/distributed/planner/multi_router_planner.c

@@ -3140,6 +3140,17 @@ MultiRouterPlannableQuery(Query *query)
 							 NULL, NULL);
 	}
 
+	if (contain_nextval_expression_walker((Node *) query->targetList, NULL))
+	{
+		/*
+		 * We let queries with nextval in the target list fall through to
+		 * the logical planner, which knows how to handle those queries.
+		 */
+		return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED,
+							 "Sequences cannot be used in router queries",
+							 NULL, NULL);
+	}
+
 	ExtractRangeTableRelationWalker((Node *) query, &rangeTableRelationList);
 	foreach(rangeTableRelationCell, rangeTableRelationList)
 	{

src/test/regress/expected/multi_insert_select.out

@@ -2283,7 +2283,6 @@ FROM
   table_with_defaults
 GROUP BY
   store_id;
-ERROR:  INSERT ... SELECT cannot generate sequence values when selecting from a distributed table
 -- do some more error/error message checks
 SET citus.shard_count TO 4;
 SET citus.shard_replication_factor TO 1;
@@ -2701,7 +2700,7 @@ SELECT create_distributed_table('dist_table_with_sequence', 'user_id');
 -- from local query
 INSERT INTO dist_table_with_sequence (value_1)
 SELECT s FROM generate_series(1,5) s;
-SELECT * FROM dist_table_with_sequence ORDER BY user_id;
+SELECT * FROM dist_table_with_sequence ORDER BY user_id, value_1;
  user_id | value_1
 ---------------------------------------------------------------------
        1 |       1
@@ -2713,9 +2712,26 @@ SELECT * FROM dist_table_with_sequence ORDER BY user_id;
 
 -- from a distributed query
 INSERT INTO dist_table_with_sequence (value_1)
-SELECT value_1 FROM dist_table_with_sequence;
+SELECT value_1 FROM dist_table_with_sequence ORDER BY value_1;
-ERROR:  INSERT ... SELECT cannot generate sequence values when selecting from a distributed table
+SELECT * FROM dist_table_with_sequence ORDER BY user_id, value_1;
-SELECT * FROM dist_table_with_sequence ORDER BY user_id;
+ user_id | value_1
+---------------------------------------------------------------------
+       1 |       1
+       2 |       2
+       3 |       3
+       4 |       4
+       5 |       5
+       6 |       1
+       7 |       2
+       8 |       3
+       9 |       4
+      10 |       5
+(10 rows)
+
+TRUNCATE dist_table_with_sequence;
+INSERT INTO dist_table_with_sequence (user_id)
+SELECT user_id FROM raw_events_second ORDER BY user_id;
+SELECT * FROM dist_table_with_sequence ORDER BY user_id, value_1;
  user_id | value_1
 ---------------------------------------------------------------------
        1 |       1
@@ -2725,8 +2741,38 @@ SELECT * FROM dist_table_with_sequence ORDER BY user_id;
        5 |       5
 (5 rows)
 
+WITH top10 AS (
+  SELECT user_id FROM raw_events_second WHERE value_1 IS NOT NULL ORDER BY value_1 LIMIT 10
+)
+INSERT INTO dist_table_with_sequence (value_1)
+SELECT * FROM top10;
+ERROR:  cannot handle complex subqueries when the router executor is disabled
+SELECT * FROM dist_table_with_sequence ORDER BY user_id, value_1;
+ user_id | value_1
+---------------------------------------------------------------------
+       1 |       1
+       2 |       2
+       3 |       3
+       4 |       4
+       5 |       5
+(5 rows)
+
+-- router queries become logical planner queries when there is a nextval call
+INSERT INTO dist_table_with_sequence (user_id)
+SELECT user_id FROM dist_table_with_sequence WHERE user_id = 1;
+SELECT * FROM dist_table_with_sequence ORDER BY user_id, value_1;
+ user_id | value_1
+---------------------------------------------------------------------
+       1 |       1
+       1 |       6
+       2 |       2
+       3 |       3
+       4 |       4
+       5 |       5
+(6 rows)
+
 -- Select from distributed table into reference table
-CREATE TABLE ref_table (user_id int, value_1 int);
+CREATE TABLE ref_table (user_id serial, value_1 int);
 SELECT create_reference_table('ref_table');
  create_reference_table
 ---------------------------------------------------------------------
@@ -2745,6 +2791,49 @@ SELECT * FROM ref_table ORDER BY user_id, value_1;
        5 |
 (5 rows)
 
+INSERT INTO ref_table (value_1)
+SELECT value_1 FROM raw_events_second ORDER BY value_1;
+SELECT * FROM ref_table ORDER BY user_id, value_1;
+ user_id | value_1
+---------------------------------------------------------------------
+       1 |
+       1 |
+       2 |
+       2 |
+       3 |
+       3 |
+       4 |
+       4 |
+       5 |
+       5 |
+(10 rows)
+
+INSERT INTO ref_table SELECT * FROM ref_table;
+SELECT * FROM ref_table ORDER BY user_id, value_1;
+ user_id | value_1
+---------------------------------------------------------------------
+       1 |
+       1 |
+       1 |
+       1 |
+       2 |
+       2 |
+       2 |
+       2 |
+       3 |
+       3 |
+       3 |
+       3 |
+       4 |
+       4 |
+       4 |
+       4 |
+       5 |
+       5 |
+       5 |
+       5 |
+(20 rows)
+
 DROP TABLE ref_table;
 -- Select from reference table into reference table
 CREATE TABLE ref1 (d timestamptz);

src/test/regress/expected/multi_simple_queries.out

@@ -722,3 +722,47 @@ DETAIL:  distribution column value: 1
 (5 rows)
 
 SET client_min_messages to 'NOTICE';
+-- we should be able to use nextval in the target list
+CREATE SEQUENCE query_seq;
+SELECT nextval('query_seq') FROM articles WHERE author_id = 1;
+ nextval
+---------------------------------------------------------------------
+       1
+       2
+       3
+       4
+       5
+(5 rows)
+
+SELECT nextval('query_seq') FROM articles LIMIT 3;
+ nextval
+---------------------------------------------------------------------
+       6
+       7
+       8
+(3 rows)
+
+SELECT nextval('query_seq')*2 FROM articles LIMIT 3;
+ ?column?
+---------------------------------------------------------------------
+       18
+       20
+       22
+(3 rows)
+
+SELECT * FROM (SELECT nextval('query_seq') FROM articles LIMIT 3) vals;
+ nextval
+---------------------------------------------------------------------
+      12
+      13
+      14
+(3 rows)
+
+-- but not elsewhere
+SELECT sum(nextval('query_seq')) FROM articles;
+ERROR:  relation "public.query_seq" does not exist
+CONTEXT:  while executing command on localhost:xxxxx
+SELECT n FROM (SELECT nextval('query_seq') n, random() FROM articles) vals;
+ERROR:  relation "public.query_seq" does not exist
+CONTEXT:  while executing command on localhost:xxxxx
+DROP SEQUENCE query_seq;

src/test/regress/expected/multi_simple_queries_0.out

@@ -666,3 +666,47 @@ DETAIL:  distribution column value: 1
 (5 rows)
 
 SET client_min_messages to 'NOTICE';
+-- we should be able to use nextval in the target list
+CREATE SEQUENCE query_seq;
+SELECT nextval('query_seq') FROM articles WHERE author_id = 1;
+ nextval
+---------------------------------------------------------------------
+       1
+       2
+       3
+       4
+       5
+(5 rows)
+
+SELECT nextval('query_seq') FROM articles LIMIT 3;
+ nextval
+---------------------------------------------------------------------
+       6
+       7
+       8
+(3 rows)
+
+SELECT nextval('query_seq')*2 FROM articles LIMIT 3;
+ ?column?
+---------------------------------------------------------------------
+       18
+       20
+       22
+(3 rows)
+
+SELECT * FROM (SELECT nextval('query_seq') FROM articles LIMIT 3) vals;
+ nextval
+---------------------------------------------------------------------
+      12
+      13
+      14
+(3 rows)
+
+-- but not elsewhere
+SELECT sum(nextval('query_seq')) FROM articles;
+ERROR:  relation "public.query_seq" does not exist
+CONTEXT:  while executing command on localhost:xxxxx
+SELECT n FROM (SELECT nextval('query_seq') n, random() FROM articles) vals;
+ERROR:  relation "public.query_seq" does not exist
+CONTEXT:  while executing command on localhost:xxxxx
+DROP SEQUENCE query_seq;

src/test/regress/sql/multi_insert_select.sql

@@ -1982,16 +1982,37 @@ SELECT create_distributed_table('dist_table_with_sequence', 'user_id');
 INSERT INTO dist_table_with_sequence (value_1)
 SELECT s FROM generate_series(1,5) s;
 
-SELECT * FROM dist_table_with_sequence ORDER BY user_id;
+SELECT * FROM dist_table_with_sequence ORDER BY user_id, value_1;
 
 -- from a distributed query
 INSERT INTO dist_table_with_sequence (value_1)
-SELECT value_1 FROM dist_table_with_sequence;
+SELECT value_1 FROM dist_table_with_sequence ORDER BY value_1;
 
-SELECT * FROM dist_table_with_sequence ORDER BY user_id;
+SELECT * FROM dist_table_with_sequence ORDER BY user_id, value_1;
+
+TRUNCATE dist_table_with_sequence;
+
+INSERT INTO dist_table_with_sequence (user_id)
+SELECT user_id FROM raw_events_second ORDER BY user_id;
+
+SELECT * FROM dist_table_with_sequence ORDER BY user_id, value_1;
+
+WITH top10 AS (
+  SELECT user_id FROM raw_events_second WHERE value_1 IS NOT NULL ORDER BY value_1 LIMIT 10
+)
+INSERT INTO dist_table_with_sequence (value_1)
+SELECT * FROM top10;
+
+SELECT * FROM dist_table_with_sequence ORDER BY user_id, value_1;
+
+-- router queries become logical planner queries when there is a nextval call
+INSERT INTO dist_table_with_sequence (user_id)
+SELECT user_id FROM dist_table_with_sequence WHERE user_id = 1;
+
+SELECT * FROM dist_table_with_sequence ORDER BY user_id, value_1;
 
 -- Select from distributed table into reference table
-CREATE TABLE ref_table (user_id int, value_1 int);
+CREATE TABLE ref_table (user_id serial, value_1 int);
 SELECT create_reference_table('ref_table');
 
 INSERT INTO ref_table
@@ -1999,6 +2020,14 @@ SELECT user_id, value_1 FROM raw_events_second;
 
 SELECT * FROM ref_table ORDER BY user_id, value_1;
 
+INSERT INTO ref_table (value_1)
+SELECT value_1 FROM raw_events_second ORDER BY value_1;
+
+SELECT * FROM ref_table ORDER BY user_id, value_1;
+
+INSERT INTO ref_table SELECT * FROM ref_table;
+SELECT * FROM ref_table ORDER BY user_id, value_1;
+
 DROP TABLE ref_table;
 
 -- Select from reference table into reference table

src/test/regress/sql/multi_simple_queries.sql

@@ -328,3 +328,16 @@ SELECT * FROM articles TABLESAMPLE SYSTEM (100) WHERE author_id = 1 ORDER BY id;
 SELECT * FROM articles TABLESAMPLE BERNOULLI (100) WHERE author_id = 1 ORDER BY id;
 
 SET client_min_messages to 'NOTICE';
+
+-- we should be able to use nextval in the target list
+CREATE SEQUENCE query_seq;
+SELECT nextval('query_seq') FROM articles WHERE author_id = 1;
+SELECT nextval('query_seq') FROM articles LIMIT 3;
+SELECT nextval('query_seq')*2 FROM articles LIMIT 3;
+SELECT * FROM (SELECT nextval('query_seq') FROM articles LIMIT 3) vals;
+
+-- but not elsewhere
+SELECT sum(nextval('query_seq')) FROM articles;
+SELECT n FROM (SELECT nextval('query_seq') n, random() FROM articles) vals;
+
+DROP SEQUENCE query_seq;