From 8c959fd6f91c24c21425105ab425564de1d459c4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philip=20Dub=C3=A9?= <gunkmute@gmail.com>
Date: Fri, 10 Apr 2020 02:37:36 +0000
Subject: [PATCH] Allow subqueries in LIMIT/OFFSET

---
 .../distributed/planner/distributed_planner.c | 21 ++++++++++-
 .../planner/multi_logical_planner.c           | 12 -------
 .../regress/expected/multi_limit_clause.out   | 36 +++++++++++++++++--
 src/test/regress/sql/multi_limit_clause.sql   |  2 ++
 4 files changed, 56 insertions(+), 15 deletions(-)

diff --git a/src/backend/distributed/planner/distributed_planner.c b/src/backend/distributed/planner/distributed_planner.c
index 719e7d498..6222d3ebf 100644
--- a/src/backend/distributed/planner/distributed_planner.c
+++ b/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;
 
 	/*
diff --git a/src/backend/distributed/planner/multi_logical_planner.c b/src/backend/distributed/planner/multi_logical_planner.c
index 549900a5e..4d30a11d3 100644
--- a/src/backend/distributed/planner/multi_logical_planner.c
+++ b/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)
 	{
diff --git a/src/test/regress/expected/multi_limit_clause.out b/src/test/regress/expected/multi_limit_clause.out
index b3e0367a3..121f16992 100644
--- a/src/test/regress/expected/multi_limit_clause.out
+++ b/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;
diff --git a/src/test/regress/sql/multi_limit_clause.sql b/src/test/regress/sql/multi_limit_clause.sql
index 8d14bbbc8..cc8008acb 100644
--- a/src/test/regress/sql/multi_limit_clause.sql
+++ b/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);