Reserve memory for one partition at a time

src/backend/utils/resource_manager/memquota.c

@@ -25,6 +25,7 @@
 #include "miscadmin.h"
 #include "cdb/cdbvars.h"
 #include "optimizer/clauses.h"
+#include "optimizer/walkers.h"
 #include "parser/parsetree.h"
 #include "tcop/pquery.h"
 
@@ -54,6 +55,7 @@ typedef struct PolicyAutoContext
  * Forward declarations.
  */
 static bool PolicyAutoPrelimWalker(Node *node, PolicyAutoContext *context);
+static bool PolicyAutoPrelimBranches(Append *append, PolicyAutoContext *context);
 static bool	PolicyAutoAssignWalker(Node *node, PolicyAutoContext *context);
 static bool IsAggMemoryIntensive(Agg *agg);
 static bool IsMemoryIntensiveOperator(Node *node, PlannedStmt *stmt);
@@ -266,6 +268,23 @@ IsMemoryIntensiveOperator(Node *node, PlannedStmt *stmt)
 	}
 }
 
+/*
+ * RunsBranchesOneAtATime
+ *    Return true if the given node runs its branches one after another.
+ *
+ * Only one branch of such a node runs at a time, so memory is reserved for
+ * its largest branch instead of the sum of all branches.  Otherwise a table
+ * with hundreds of partitions needs tens of megabytes just to start a query.
+ *
+ * This holds for an Append.  Asynchronous branches run together, and a
+ * MergeAppend reads all branches at once, so those do not count.
+ */
+static bool
+RunsBranchesOneAtATime(Node *node)
+{
+	return IsA(node, Append) && ((Append *) node)->nasyncplans == 0;
+}
+
 /*
  * IsRootOperatorInGroup
  *    Return true if the given node is the root operator in an operator group.
@@ -305,10 +324,55 @@ static bool PolicyAutoPrelimWalker(Node *node, PolicyAutoContext *context)
 		{
 			context->numNonMemIntensiveOperators++;
 		}
+
+		if (RunsBranchesOneAtATime(node))
+		{
+			return PolicyAutoPrelimBranches((Append *) node, context);
+		}
 	}
 	return plan_tree_walker(node, PolicyAutoPrelimWalker, context, true);
 }
 
+/*
+ * PolicyAutoPrelimBranches
+ *    Count the operators of an Append, keeping only its largest branch.
+ *    See RunsBranchesOneAtATime().
+ */
+static bool
+PolicyAutoPrelimBranches(Append *append, PolicyAutoContext *context)
+{
+	uint64		baseMemIntense;
+	uint64		baseNonMemIntense;
+	uint64		maxMemIntense = 0;
+	uint64		maxNonMemIntense = 0;
+	ListCell   *lc;
+
+	if (walk_plan_node_fields((Plan *) append, PolicyAutoPrelimWalker, context))
+		return true;
+
+	baseMemIntense = context->numMemIntensiveOperators;
+	baseNonMemIntense = context->numNonMemIntensiveOperators;
+
+	foreach(lc, append->appendplans)
+	{
+		context->numMemIntensiveOperators = baseMemIntense;
+		context->numNonMemIntensiveOperators = baseNonMemIntense;
+
+		if (PolicyAutoPrelimWalker((Node *) lfirst(lc), context))
+			return true;
+
+		maxMemIntense = Max(maxMemIntense,
+							context->numMemIntensiveOperators - baseMemIntense);
+		maxNonMemIntense = Max(maxNonMemIntense,
+							   context->numNonMemIntensiveOperators - baseNonMemIntense);
+	}
+
+	context->numMemIntensiveOperators = baseMemIntense + maxMemIntense;
+	context->numNonMemIntensiveOperators = baseNonMemIntense + maxNonMemIntense;
+
+	return false;
+}
+
 /**
  * This walker assigns specific amount of memory to each operator in a plan.
  * It allocates a fixed size to each non-memory intensive operator and distributes
@@ -679,6 +743,9 @@ ComputeMemLimitForChildGroups(OperatorGroupNode *parentGroupNode)
  * node (except for the leaves of the leave groups). At the same time,
  * we collect some stats information about operators in each group.
  */
+static bool PolicyEagerFreePrelimBranches(Append *append,
+										  PolicyEagerFreeContext *context);
+
 static bool
 PolicyEagerFreePrelimWalker(Node *node, PolicyEagerFreeContext *context)
 {
@@ -715,7 +782,12 @@ PolicyEagerFreePrelimWalker(Node *node, PolicyEagerFreeContext *context)
 		}
 	}
 
-	bool result = plan_tree_walker(node, PolicyEagerFreePrelimWalker, context, true);
+	bool result;
+
+	if (is_plan_node(node) && RunsBranchesOneAtATime(node))
+		result = PolicyEagerFreePrelimBranches((Append *) node, context);
+	else
+		result = plan_tree_walker(node, PolicyEagerFreePrelimWalker, context, true);
 	Assert(!result);
 
 	/*
@@ -753,6 +825,55 @@ PolicyEagerFreePrelimWalker(Node *node, PolicyEagerFreeContext *context)
 	return result;
 }
 
+/*
+ * PolicyEagerFreePrelimBranches
+ *    Same as PolicyAutoPrelimBranches(), for the eager free policy.
+ *
+ * Only operators in the Append's own group are counted this way.  A branch
+ * that starts its own group, for example with a sort, is still added in full,
+ * so we may reserve a little too much, but never too little.
+ *
+ * Visit the branches in the same order as plan_tree_walker(), so the groups
+ * get the same numbers when the memory is assigned later.
+ */
+static bool
+PolicyEagerFreePrelimBranches(Append *append, PolicyEagerFreeContext *context)
+{
+	OperatorGroupNode *groupNode = context->groupNode;
+	uint64		baseMemIntense;
+	uint64		baseNonMemIntense;
+	uint64		maxMemIntense = 0;
+	uint64		maxNonMemIntense = 0;
+	ListCell   *lc;
+
+	Assert(groupNode != NULL);
+
+	if (walk_plan_node_fields((Plan *) append, PolicyEagerFreePrelimWalker, context))
+		return true;
+
+	baseMemIntense = groupNode->numMemIntenseOps;
+	baseNonMemIntense = groupNode->numNonMemIntenseOps;
+
+	foreach(lc, append->appendplans)
+	{
+		groupNode->numMemIntenseOps = baseMemIntense;
+		groupNode->numNonMemIntenseOps = baseNonMemIntense;
+
+		if (PolicyEagerFreePrelimWalker((Node *) lfirst(lc), context))
+			return true;
+
+		maxMemIntense = Max(maxMemIntense,
+							groupNode->numMemIntenseOps - baseMemIntense);
+		maxNonMemIntense = Max(maxNonMemIntense,
+							   groupNode->numNonMemIntenseOps - baseNonMemIntense);
+	}
+
+	groupNode->numMemIntenseOps = baseMemIntense + maxMemIntense;
+	groupNode->numNonMemIntenseOps = baseNonMemIntense + maxNonMemIntense;
+
+	return false;
+}
+
 /*
  * PolicyEagerFreeAssignWalker
  *    Walk the plan tree and assign the memory to each plan node.

src/test/regress/expected/memquota_partitions.out

@@ -0,0 +1,28 @@
+-- Memory reserved for a query over a table with many partitions.
+--
+-- The Append runs one partition at a time, so the branches never hold their
+-- memory together.  Reserving memory for all of them at once used to make the
+-- statement fail outright.
+CREATE TABLE memquota_parts (id int, d date) DISTRIBUTED BY (id)
+  PARTITION BY RANGE (d)
+  (START ('2020-01-01'::date) END ('2020-07-19'::date) EVERY ('1 day'::interval));
+-- ORCA scans the partitions with a single Dynamic Seq Scan, so the Append this
+-- test is about is only built by the Postgres planner.
+SET optimizer = off;
+SET statement_mem = '2MB';
+SELECT count(*) FROM memquota_parts;
+ count 
+-------
+     0
+(1 row)
+
+-- the same through a subquery, so the Append is not the top node
+SELECT count(*) FROM (SELECT * FROM memquota_parts WHERE id > 0) x;
+ count 
+-------
+     0
+(1 row)
+
+RESET statement_mem;
+RESET optimizer;
+DROP TABLE memquota_parts;

src/test/regress/greenplum_schedule

@@ -71,6 +71,7 @@ test: guc_gp
 test: toast
 test: misc_jiras
 test: statement_mem_for_windowagg
+test: memquota_partitions
 test: write_gang_idle_in_transaction_session_timeout
 
 # namespace_gp test will show diff if concurrent tests use temporary tables.

src/test/regress/sql/memquota_partitions.sql

@@ -0,0 +1,22 @@
+-- Memory reserved for a query over a table with many partitions.
+--
+-- The Append runs one partition at a time, so the branches never hold their
+-- memory together.  Reserving memory for all of them at once used to make the
+-- statement fail outright.
+CREATE TABLE memquota_parts (id int, d date) DISTRIBUTED BY (id)
+  PARTITION BY RANGE (d)
+  (START ('2020-01-01'::date) END ('2020-07-19'::date) EVERY ('1 day'::interval));
+
+-- ORCA scans the partitions with a single Dynamic Seq Scan, so the Append this
+-- test is about is only built by the Postgres planner.
+SET optimizer = off;
+SET statement_mem = '2MB';
+
+SELECT count(*) FROM memquota_parts;
+
+-- the same through a subquery, so the Append is not the top node
+SELECT count(*) FROM (SELECT * FROM memquota_parts WHERE id > 0) x;
+
+RESET statement_mem;
+RESET optimizer;
+DROP TABLE memquota_parts;