Allow subqueries in LIMIT/OFFSET

src/backend/distributed/planner/distributed_planner.c

@@ -258,8 +258,27 @@ distributed_planner(Query *parse, int cursorOptions, ParamListInfo boundParams)
 			planContext.plan = standard_planner(planContext.query,
 												planContext.cursorOptions,
 												planContext.boundParams);
+
 			if (needsDistributedPlanning)
 			{
+				/*
+				 * standard_planner rewrites simple queries like 'select 10' to PARAM_EXEC nodes,
+				 * which we're unable to handle. Meanwhile we only optimize rewrites to Const.
+				 * So deoptimize non-Const LIMIT/OFFSET, standard_planner will handle it again later.
+				 */
+				if (planContext.query->limitCount &&
+					!IsA(planContext.query->limitCount, Const))
+				{
+					planContext.query->limitCount = planContext.originalQuery->limitCount;
+				}
+
+				if (planContext.query->limitOffset &&
+					!IsA(planContext.query->limitOffset, Const))
+				{
+					planContext.query->limitOffset =
+						planContext.originalQuery->limitOffset;
+				}
+
 				result = PlanDistributedStmt(&planContext, rteIdCounter);
 			}
 			else if ((result = TryToDelegateFunctionCall(&planContext)) == NULL)
@@ -806,7 +825,7 @@ InlineCtesAndCreateDistributedPlannedStmt(uint64 planId,
 	/* after inlining, we shouldn't have any inlinable CTEs */
 	Assert(!QueryTreeContainsInlinableCTE(copyOfOriginalQuery));
 
-	#if PG_VERSION_NUM < PG_VERSION_12
+#if PG_VERSION_NUM < PG_VERSION_12
 	Query *query = planContext->query;
 
 	/*

src/backend/distributed/planner/multi_logical_planner.c

@@ -1013,18 +1013,6 @@ DeferErrorIfQueryNotSupported(Query *queryTree)
 		errorHint = filterHint;
 	}
 
-	if (FindNodeCheck((Node *) queryTree->limitCount, IsNodeSubquery))
-	{
-		preconditionsSatisfied = false;
-		errorMessage = "subquery in LIMIT is not supported in multi-shard queries";
-	}
-
-	if (FindNodeCheck((Node *) queryTree->limitOffset, IsNodeSubquery))
-	{
-		preconditionsSatisfied = false;
-		errorMessage = "subquery in OFFSET is not supported in multi-shard queries";
-	}
-
 	/* finally check and error out if not satisfied */
 	if (!preconditionsSatisfied)
 	{

src/test/regress/expected/multi_limit_clause.out

@@ -553,8 +553,40 @@ SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT 10 OFFSET my_limi
 
 DROP FUNCTION my_limit();
 -- subqueries should error out
+SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT (SELECT min(l_linenumber) FROM lineitem_hash) OFFSET (SELECT (count(*)/2)::int FROM lineitem_hash);
+ l_orderkey
+---------------------------------------------------------------------
+       8997
+(1 row)
+
 SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT (SELECT 10);
-ERROR:  subquery in LIMIT is not supported in multi-shard queries
+ l_orderkey
+---------------------------------------------------------------------
+          1
+          1
+          1
+          1
+          1
+          1
+          2
+          3
+          3
+          3
+(10 rows)
+
 SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT 10 OFFSET (SELECT 10);
-ERROR:  subquery in OFFSET is not supported in multi-shard queries
+ l_orderkey
+---------------------------------------------------------------------
+          3
+          3
+          3
+          4
+          5
+          5
+          5
+          6
+          7
+          7
+(10 rows)
+
 DROP TABLE lineitem_hash;

src/test/regress/sql/multi_limit_clause.sql

@@ -237,6 +237,8 @@ SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT 10 OFFSET my_limi
 DROP FUNCTION my_limit();
 
 -- subqueries should error out
+SELECT min(l_linenumber) FROM lineitem;
+SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT (SELECT min(l_linenumber) FROM lineitem) OFFSET (SELECT (count(*)/2)::int FROM lineitem_hash);
 SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT (SELECT 10);
 SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT 10 OFFSET (SELECT 10);