Re: Optimize planner memory consumption for huge arrays

Tomas Vondra <tomas.vondra@enterprisedb.com> writes:
> On 2/19/24 16:45, Tom Lane wrote:
>> Tomas Vondra <tomas.vondra@enterprisedb.com> writes:
>>> For example, I don't think we expect selectivity functions to allocate
>>> long-lived objects, right? So maybe we could run them in a dedicated
>>> memory context, and reset it aggressively (after each call).

>> That could eliminate a whole lot of potential leaks.  I'm not sure 
>> though how much it moves the needle in terms of overall planner
>> memory consumption.

> I'm not sure about that either, maybe not much - for example it would
> not help with the two other memory usage patches (which are related to
> SpecialJoinInfo and RestrictInfo, outside selectivity functions).

> It was an ad hoc thought, inspired by the issue at hand. Maybe it would
> be possible to find similar "boundaries" in other parts of the planner.

Here's a quick and probably-incomplete implementation of that idea.
I've not tried to study its effects on memory consumption, just made
sure it passes check-world.

The main hazard here is that something invoked inside clause
selectivity might try to cache a data structure for later use.
However, there are already places that do that kind of thing,
and they already explicitly switch into the planner_cxt, because
otherwise they fail under GEQO.  (If we do find places that need
fixing for this, they were probably busted under GEQO already.)
Perhaps it's worth updating the comments at those places, but
I didn't bother in this first cut.

            regards, tom lane

diff --git a/src/backend/optimizer/path/clausesel.c b/src/backend/optimizer/path/clausesel.c
index c949dc1866..669acd7315 100644
--- a/src/backend/optimizer/path/clausesel.c
+++ b/src/backend/optimizer/path/clausesel.c
@@ -24,6 +24,7 @@
 #include "statistics/statistics.h"
 #include "utils/fmgroids.h"
 #include "utils/lsyscache.h"
+#include "utils/memutils.h"
 #include "utils/selfuncs.h"

 /*
@@ -693,6 +694,7 @@ clause_selectivity_ext(PlannerInfo *root,
     Selectivity s1 = 0.5;        /* default for any unhandled clause type */
     RestrictInfo *rinfo = NULL;
     bool        cacheable = false;
+    MemoryContext saved_cxt;

     if (clause == NULL)            /* can this still happen? */
         return s1;
@@ -756,6 +758,21 @@ clause_selectivity_ext(PlannerInfo *root,
             clause = (Node *) rinfo->clause;
     }

+    /*
+     * Run all the remaining work in the short-lived planner_tmp_cxt, which
+     * we'll reset at the bottom.  This allows selectivity-related code to not
+     * be too concerned about leaking memory.
+     */
+    saved_cxt = MemoryContextSwitchTo(root->planner_tmp_cxt);
+
+    /*
+     * This function can be called recursively, in which case we'd better not
+     * reset planner_tmp_cxt until we exit the topmost level.  Use of
+     * planner_tmp_cxt_depth also makes it safe for other places to use and
+     * reset planner_tmp_cxt in the same fashion.
+     */
+    root->planner_tmp_cxt_depth++;
+
     if (IsA(clause, Var))
     {
         Var           *var = (Var *) clause;
@@ -967,6 +984,12 @@ clause_selectivity_ext(PlannerInfo *root,
             rinfo->outer_selec = s1;
     }

+    /* Exit and clean up the planner_tmp_cxt */
+    MemoryContextSwitchTo(saved_cxt);
+    Assert(root->planner_tmp_cxt_depth > 0);
+    if (--root->planner_tmp_cxt_depth == 0)
+        MemoryContextReset(root->planner_tmp_cxt);
+
 #ifdef SELECTIVITY_DEBUG
     elog(DEBUG4, "clause_selectivity: s1 %f", s1);
 #endif                            /* SELECTIVITY_DEBUG */
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index be4e182869..44b9555dbf 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -643,6 +643,17 @@ subquery_planner(PlannerGlobal *glob, Query *parse,
     root->plan_params = NIL;
     root->outer_params = NULL;
     root->planner_cxt = CurrentMemoryContext;
+    /* a subquery can share the main query's planner_tmp_cxt */
+    if (parent_root)
+    {
+        root->planner_tmp_cxt = parent_root->planner_tmp_cxt;
+        Assert(parent_root->planner_tmp_cxt_depth == 0);
+    }
+    else
+        root->planner_tmp_cxt = AllocSetContextCreate(root->planner_cxt,
+                                                      "Planner temp context",
+                                                      ALLOCSET_DEFAULT_SIZES);
+    root->planner_tmp_cxt_depth = 0;
     root->init_plans = NIL;
     root->cte_plan_ids = NIL;
     root->multiexpr_params = NIL;
@@ -6514,6 +6525,10 @@ plan_cluster_use_sort(Oid tableOid, Oid indexOid)
     root->glob = glob;
     root->query_level = 1;
     root->planner_cxt = CurrentMemoryContext;
+    root->planner_tmp_cxt = AllocSetContextCreate(root->planner_cxt,
+                                                  "Planner temp context",
+                                                  ALLOCSET_DEFAULT_SIZES);
+    root->planner_tmp_cxt_depth = 0;
     root->wt_param_id = -1;
     root->join_domains = list_make1(makeNode(JoinDomain));

@@ -6552,7 +6567,10 @@ plan_cluster_use_sort(Oid tableOid, Oid indexOid)
      * trust the index contents but use seqscan-and-sort.
      */
     if (lc == NULL)                /* not in the list? */
+    {
+        MemoryContextDelete(root->planner_tmp_cxt);
         return true;            /* use sort */
+    }

     /*
      * Rather than doing all the pushups that would be needed to use
@@ -6584,6 +6602,9 @@ plan_cluster_use_sort(Oid tableOid, Oid indexOid)
                                       ForwardScanDirection, false,
                                       NULL, 1.0, false);

+    /* We assume this won't free *indexScanPath */
+    MemoryContextDelete(root->planner_tmp_cxt);
+
     return (seqScanAndSortPath.total_cost < indexScanPath->path.total_cost);
 }

@@ -6636,6 +6657,10 @@ plan_create_index_workers(Oid tableOid, Oid indexOid)
     root->glob = glob;
     root->query_level = 1;
     root->planner_cxt = CurrentMemoryContext;
+    root->planner_tmp_cxt = AllocSetContextCreate(root->planner_cxt,
+                                                  "Planner temp context",
+                                                  ALLOCSET_DEFAULT_SIZES);
+    root->planner_tmp_cxt_depth = 0;
     root->wt_param_id = -1;
     root->join_domains = list_make1(makeNode(JoinDomain));

@@ -6725,6 +6750,8 @@ plan_create_index_workers(Oid tableOid, Oid indexOid)
         parallel_workers--;

 done:
+    MemoryContextDelete(root->planner_tmp_cxt);
+
     index_close(index, NoLock);
     table_close(heap, NoLock);

diff --git a/src/backend/optimizer/prep/prepjointree.c b/src/backend/optimizer/prep/prepjointree.c
index aa83dd3636..0ea8361190 100644
--- a/src/backend/optimizer/prep/prepjointree.c
+++ b/src/backend/optimizer/prep/prepjointree.c
@@ -987,7 +987,10 @@ pull_up_simple_subquery(PlannerInfo *root, Node *jtnode, RangeTblEntry *rte,
     subroot->parent_root = root->parent_root;
     subroot->plan_params = NIL;
     subroot->outer_params = NULL;
-    subroot->planner_cxt = CurrentMemoryContext;
+    subroot->planner_cxt = root->planner_cxt;
+    subroot->planner_tmp_cxt = root->planner_tmp_cxt;
+    Assert(root->planner_tmp_cxt_depth == 0);
+    subroot->planner_tmp_cxt_depth = 0;
     subroot->init_plans = NIL;
     subroot->cte_plan_ids = NIL;
     subroot->multiexpr_params = NIL;
diff --git a/src/include/nodes/pathnodes.h b/src/include/nodes/pathnodes.h
index 534692bee1..d278a9773a 100644
--- a/src/include/nodes/pathnodes.h
+++ b/src/include/nodes/pathnodes.h
@@ -471,6 +471,11 @@ struct PlannerInfo
     /* context holding PlannerInfo */
     MemoryContext planner_cxt pg_node_attr(read_write_ignore);

+    /* short-lived context for purposes such as calling selectivity functions */
+    MemoryContext planner_tmp_cxt pg_node_attr(read_write_ignore);
+    /* nesting depth of uses of planner_tmp_cxt; reset it only at level 0 */
+    int            planner_tmp_cxt_depth;
+
     /* # of pages in all non-dummy tables of query */
     Cardinality total_table_pages;