Merge pull request #3690 from citusdata/fix/limit_non_const

Correctly handle non-constant LIMIT/OFFSET clauses

src/backend/distributed/planner/multi_logical_optimizer.c

@@ -4320,19 +4320,30 @@ WorkerLimitCount(Node *limitCount, Node *limitOffset, OrderByLimitReference
 	Node *workerLimitNode = NULL;
 	LimitPushdownable canPushDownLimit = LIMIT_CANNOT_PUSHDOWN;
 
-	/* no limit node to push down */
 	if (limitCount == NULL)
 	{
+		/* no limit node to push down */
+		return NULL;
+	}
+
+	if (!IsA(limitCount, Const))
+	{
+		/*
+		 * We only push down constant LIMIT clauses to make sure we get back
+		 * the minimum number of rows.
+		 */
+		return NULL;
+	}
+
+	if (limitOffset != NULL && !IsA(limitOffset, Const))
+	{
+		/*
+		 * If OFFSET is not a constant then we cannot calculate the LIMIT to
+		 * push down.
+		 */
 		return NULL;
 	}
 
-	/*
-	 * During subquery pushdown planning original query is used. In that case,
-	 * certain expressions such as parameters are not evaluated and converted
-	 * into Consts on the op node.
-	 */
-	Assert(IsA(limitCount, Const));
-	Assert(limitOffset == NULL || IsA(limitOffset, Const));
 
 	/*
 	 * If window functions are computed on coordinator, we cannot push down LIMIT.

src/backend/distributed/planner/multi_logical_planner.c

@@ -1024,6 +1024,17 @@ 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

@@ -520,4 +520,41 @@ SELECT
 (1 row)
 
 SET client_min_messages TO NOTICE;
+-- non constants should not push down
+CREATE OR REPLACE FUNCTION my_limit()
+RETURNS INT AS $$
+BEGIN
+  RETURN 5;
+END; $$ language plpgsql VOLATILE;
+SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT my_limit();
+ l_orderkey
+---------------------------------------------------------------------
+          1
+          1
+          1
+          1
+          1
+(5 rows)
+
+SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT 10 OFFSET my_limit();
+ l_orderkey
+---------------------------------------------------------------------
+          1
+          2
+          3
+          3
+          3
+          3
+          3
+          3
+          4
+          5
+(10 rows)
+
+DROP FUNCTION my_limit();
+-- subqueries should error out
+SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT (SELECT 10);
+ERROR:  subquery in LIMIT is not supported in multi-shard queries
+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
 DROP TABLE lineitem_hash;

src/test/regress/sql/multi_limit_clause.sql

@@ -223,4 +223,21 @@ SELECT
 	LIMIT 5;
 
 SET client_min_messages TO NOTICE;
+
+-- non constants should not push down
+CREATE OR REPLACE FUNCTION my_limit()
+RETURNS INT AS $$
+BEGIN
+  RETURN 5;
+END; $$ language plpgsql VOLATILE;
+
+SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT my_limit();
+SELECT l_orderkey FROM lineitem_hash ORDER BY l_orderkey LIMIT 10 OFFSET my_limit();
+
+DROP FUNCTION my_limit();
+
+-- subqueries should error out
+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);
+
 DROP TABLE lineitem_hash;