Allowing parallel-safe initplans - Mailing list pgsql-hackers

Pursuant to the discussion at [1], here's a patch that removes our
old restriction that a plan node having initPlans can't be marked
parallel-safe (dating to commit ab77a5a45).  That was really a special
case of the fact that we couldn't transmit subplans to parallel
workers at all.  We fixed that in commit 5e6d8d2bb and follow-ons,
but this case never got addressed.

Along the way, this also takes care of some sloppiness about updating
path costs to match when we move initplans from one place to another
during createplan.c and setrefs.c.  Since all the planning decisions are
already made by that point, this is just cosmetic; but it seems good
to keep EXPLAIN output consistent with where the initplans are.

The diff in query_planner() might be worth remarking on.  I found
that one because after fixing things to allow parallel-safe initplans,
one partition_prune test case changed plans (as shown in the patch)
--- but only when debug_parallel_query was active.  The reason
proved to be that we only bothered to mark Result nodes as potentially
parallel-safe when debug_parallel_query is on.  This neglects the
fact that parallel-safety may be of interest for a sub-query even
though the Result itself doesn't parallelize.

There's only one existing test case that visibly changes plan with
these changes.  The new plan is clearly saner-looking than before,
and testing with some data loaded into the table confirms that it
is faster.  I'm not sure if it's worth devising more test cases.

I'll park this in the July commitfest.

            regards, tom lane

[1] https://www.postgresql.org/message-id/flat/ZDVt6MaNWkRDO1LQ%40telsasoft.com

diff --git a/src/backend/optimizer/plan/createplan.c b/src/backend/optimizer/plan/createplan.c
index 910ffbf1e1..2ec8a537ae 100644
--- a/src/backend/optimizer/plan/createplan.c
+++ b/src/backend/optimizer/plan/createplan.c
@@ -6488,6 +6488,8 @@ materialize_finished_plan(Plan *subplan)
 {
     Plan       *matplan;
     Path        matpath;        /* dummy for result of cost_material */
+    Cost        initplan_cost;
+    bool        unsafe_initplans;

     matplan = (Plan *) make_material(subplan);

@@ -6495,20 +6497,25 @@ materialize_finished_plan(Plan *subplan)
      * XXX horrid kluge: if there are any initPlans attached to the subplan,
      * move them up to the Material node, which is now effectively the top
      * plan node in its query level.  This prevents failure in
-     * SS_finalize_plan(), which see for comments.  We don't bother adjusting
-     * the subplan's cost estimate for this.
+     * SS_finalize_plan(), which see for comments.
      */
     matplan->initPlan = subplan->initPlan;
     subplan->initPlan = NIL;

+    /* Move the initplans' cost delta, as well */
+    SS_compute_initplan_cost(matplan->initPlan,
+                             &initplan_cost, &unsafe_initplans);
+    subplan->startup_cost -= initplan_cost;
+    subplan->total_cost -= initplan_cost;
+
     /* Set cost data */
     cost_material(&matpath,
                   subplan->startup_cost,
                   subplan->total_cost,
                   subplan->plan_rows,
                   subplan->plan_width);
-    matplan->startup_cost = matpath.startup_cost;
-    matplan->total_cost = matpath.total_cost;
+    matplan->startup_cost = matpath.startup_cost + initplan_cost;
+    matplan->total_cost = matpath.total_cost + initplan_cost;
     matplan->plan_rows = subplan->plan_rows;
     matplan->plan_width = subplan->plan_width;
     matplan->parallel_aware = false;
diff --git a/src/backend/optimizer/plan/planmain.c b/src/backend/optimizer/plan/planmain.c
index 7afd434c60..fcc0eacd25 100644
--- a/src/backend/optimizer/plan/planmain.c
+++ b/src/backend/optimizer/plan/planmain.c
@@ -112,14 +112,17 @@ query_planner(PlannerInfo *root,
                  * quals are parallel-restricted.  (We need not check
                  * final_rel->reltarget because it's empty at this point.
                  * Anything parallel-restricted in the query tlist will be
-                 * dealt with later.)  This is normally pretty silly, because
-                 * a Result-only plan would never be interesting to
-                 * parallelize.  However, if debug_parallel_query is on, then
-                 * we want to execute the Result in a parallel worker if
-                 * possible, so we must do this.
+                 * dealt with later.)  We should always do this in a subquery,
+                 * since it might be useful to use the subquery in parallel
+                 * paths in the parent level.  At top level this is normally
+                 * not worth the cycles, because a Result-only plan would
+                 * never be interesting to parallelize.  However, if
+                 * debug_parallel_query is on, then we want to execute the
+                 * Result in a parallel worker if possible, so we must check.
                  */
                 if (root->glob->parallelModeOK &&
-                    debug_parallel_query != DEBUG_PARALLEL_OFF)
+                    (root->query_level > 1 ||
+                     debug_parallel_query != DEBUG_PARALLEL_OFF))
                     final_rel->consider_parallel =
                         is_parallel_safe(root, parse->jointree->quals);

diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index 502ccbcea2..e196f06c07 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -430,17 +430,19 @@ standard_planner(Query *parse, const char *query_string, int cursorOptions,
     /*
      * Optionally add a Gather node for testing purposes, provided this is
      * actually a safe thing to do.
+     *
+     * We could potentially do this even when top_plan has parallel-safe
+     * initPlans, but we'd have to move the initPlans to the Gather node
+     * because of SS_finalize_plan's limitations.  That causes cosmetic
+     * breakage of regression tests, because initPlans that would normally
+     * appear on the top_plan move to the Gather, changing EXPLAIN results.
+     * That doesn't seem worth working hard on, so skip it for now.
      */
-    if (debug_parallel_query != DEBUG_PARALLEL_OFF && top_plan->parallel_safe)
+    if (debug_parallel_query != DEBUG_PARALLEL_OFF &&
+        top_plan->parallel_safe && top_plan->initPlan == NIL)
     {
         Gather       *gather = makeNode(Gather);

-        /*
-         * Top plan must not have any initPlans, else it shouldn't have been
-         * marked parallel-safe.
-         */
-        Assert(top_plan->initPlan == NIL);
-
         gather->plan.targetlist = top_plan->targetlist;
         gather->plan.qual = NIL;
         gather->plan.lefttree = top_plan;
diff --git a/src/backend/optimizer/plan/setrefs.c b/src/backend/optimizer/plan/setrefs.c
index 1812db7f2f..04e4fd1d2f 100644
--- a/src/backend/optimizer/plan/setrefs.c
+++ b/src/backend/optimizer/plan/setrefs.c
@@ -23,6 +23,7 @@
 #include "optimizer/pathnode.h"
 #include "optimizer/planmain.h"
 #include "optimizer/planner.h"
+#include "optimizer/subselect.h"
 #include "optimizer/tlist.h"
 #include "parser/parse_relation.h"
 #include "tcop/utility.h"
@@ -1544,19 +1545,30 @@ clean_up_removed_plan_level(Plan *parent, Plan *child)
 {
     /*
      * We have to be sure we don't lose any initplans, so move any that were
-     * attached to the parent plan to the child.  If we do move any, the child
-     * is no longer parallel-safe.
+     * attached to the parent plan to the child.  If any are parallel-unsafe,
+     * the child is no longer parallel-safe.  As a cosmetic matter, also add
+     * the initplans' run costs to the child's costs.
      */
     if (parent->initPlan)
-        child->parallel_safe = false;
+    {
+        Cost        initplan_cost;
+        bool        unsafe_initplans;

-    /*
-     * Attach plans this way so that parent's initplans are processed before
-     * any pre-existing initplans of the child.  Probably doesn't matter, but
-     * let's preserve the ordering just in case.
-     */
-    child->initPlan = list_concat(parent->initPlan,
-                                  child->initPlan);
+        SS_compute_initplan_cost(parent->initPlan,
+                                 &initplan_cost, &unsafe_initplans);
+        child->startup_cost += initplan_cost;
+        child->total_cost += initplan_cost;
+        if (unsafe_initplans)
+            child->parallel_safe = false;
+
+        /*
+         * Attach plans this way so that parent's initplans are processed
+         * before any pre-existing initplans of the child.  Probably doesn't
+         * matter, but let's preserve the ordering just in case.
+         */
+        child->initPlan = list_concat(parent->initPlan,
+                                      child->initPlan);
+    }

     /*
      * We also have to transfer the parent's column labeling info into the
diff --git a/src/backend/optimizer/plan/subselect.c b/src/backend/optimizer/plan/subselect.c
index 5f12b2ef9b..c9285ad625 100644
--- a/src/backend/optimizer/plan/subselect.c
+++ b/src/backend/optimizer/plan/subselect.c
@@ -1016,8 +1016,7 @@ SS_process_ctes(PlannerInfo *root)

         /*
          * CTE scans are not considered for parallelism (cf
-         * set_rel_consider_parallel), and even if they were, initPlans aren't
-         * parallel-safe.
+         * set_rel_consider_parallel).
          */
         splan->parallel_safe = false;
         splan->setParam = NIL;
@@ -2120,8 +2119,8 @@ SS_identify_outer_params(PlannerInfo *root)
  * If any initPlans have been created in the current query level, they will
  * get attached to the Plan tree created from whichever Path we select from
  * the given rel.  Increment all that rel's Paths' costs to account for them,
- * and make sure the paths get marked as parallel-unsafe, since we can't
- * currently transmit initPlans to parallel workers.
+ * and if any of the initPlans are parallel-unsafe, mark all the rel's Paths
+ * parallel-unsafe as well.
  *
  * This is separate from SS_attach_initplans because we might conditionally
  * create more initPlans during create_plan(), depending on which Path we
@@ -2132,6 +2131,7 @@ void
 SS_charge_for_initplans(PlannerInfo *root, RelOptInfo *final_rel)
 {
     Cost        initplan_cost;
+    bool        unsafe_initplans;
     ListCell   *lc;

     /* Nothing to do if no initPlans */
@@ -2140,17 +2140,10 @@ SS_charge_for_initplans(PlannerInfo *root, RelOptInfo *final_rel)

     /*
      * Compute the cost increment just once, since it will be the same for all
-     * Paths.  We assume each initPlan gets run once during top plan startup.
-     * This is a conservative overestimate, since in fact an initPlan might be
-     * executed later than plan startup, or even not at all.
+     * Paths.  Also check for parallel-unsafe initPlans.
      */
-    initplan_cost = 0;
-    foreach(lc, root->init_plans)
-    {
-        SubPlan    *initsubplan = (SubPlan *) lfirst(lc);
-
-        initplan_cost += initsubplan->startup_cost + initsubplan->per_call_cost;
-    }
+    SS_compute_initplan_cost(root->init_plans,
+                             &initplan_cost, &unsafe_initplans);

     /*
      * Now adjust the costs and parallel_safe flags.
@@ -2161,19 +2154,71 @@ SS_charge_for_initplans(PlannerInfo *root, RelOptInfo *final_rel)

         path->startup_cost += initplan_cost;
         path->total_cost += initplan_cost;
-        path->parallel_safe = false;
+        if (unsafe_initplans)
+            path->parallel_safe = false;
     }

     /*
-     * Forget about any partial paths and clear consider_parallel, too;
-     * they're not usable if we attached an initPlan.
+     * Adjust partial paths' costs too, or forget them entirely if we must
+     * consider the rel parallel-unsafe.
      */
-    final_rel->partial_pathlist = NIL;
-    final_rel->consider_parallel = false;
+    if (unsafe_initplans)
+    {
+        final_rel->partial_pathlist = NIL;
+        final_rel->consider_parallel = false;
+    }
+    else
+    {
+        foreach(lc, final_rel->partial_pathlist)
+        {
+            Path       *path = (Path *) lfirst(lc);
+
+            path->startup_cost += initplan_cost;
+            path->total_cost += initplan_cost;
+        }
+    }

     /* We needn't do set_cheapest() here, caller will do it */
 }

+/*
+ * SS_compute_initplan_cost - count up the cost delta for some initplans
+ *
+ * The total cost returned in *initplan_cost_p should be added to both the
+ * startup and total costs of the plan node the initplans get attached to.
+ * We also report whether any of the initplans are not parallel-safe.
+ *
+ * The primary user of this is SS_charge_for_initplans, but it's also
+ * used in adjusting costs when we move initplans to another plan node.
+ */
+void
+SS_compute_initplan_cost(List *init_plans,
+                         Cost *initplan_cost_p,
+                         bool *unsafe_initplans_p)
+{
+    Cost        initplan_cost;
+    bool        unsafe_initplans;
+    ListCell   *lc;
+
+    /*
+     * We assume each initPlan gets run once during top plan startup.  This is
+     * a conservative overestimate, since in fact an initPlan might be
+     * executed later than plan startup, or even not at all.
+     */
+    initplan_cost = 0;
+    unsafe_initplans = false;
+    foreach(lc, init_plans)
+    {
+        SubPlan    *initsubplan = lfirst_node(SubPlan, lc);
+
+        initplan_cost += initsubplan->startup_cost + initsubplan->per_call_cost;
+        if (!initsubplan->parallel_safe)
+            unsafe_initplans = true;
+    }
+    *initplan_cost_p = initplan_cost;
+    *unsafe_initplans_p = unsafe_initplans;
+}
+
 /*
  * SS_attach_initplans - attach initplans to topmost plan node
  *
@@ -2967,6 +3012,7 @@ SS_make_initplan_from_plan(PlannerInfo *root,
                                node->plan_id, prm->paramid);
     get_first_col_type(plan, &node->firstColType, &node->firstColTypmod,
                        &node->firstColCollation);
+    node->parallel_safe = plan->parallel_safe;
     node->setParam = list_make1_int(prm->paramid);

     root->init_plans = lappend(root->init_plans, node);
diff --git a/src/backend/optimizer/util/pathnode.c b/src/backend/optimizer/util/pathnode.c
index 5f5596841c..f123fcb41e 100644
--- a/src/backend/optimizer/util/pathnode.c
+++ b/src/backend/optimizer/util/pathnode.c
@@ -3348,8 +3348,7 @@ create_minmaxagg_path(PlannerInfo *root,
     /* For now, assume we are above any joins, so no parameterization */
     pathnode->path.param_info = NULL;
     pathnode->path.parallel_aware = false;
-    /* A MinMaxAggPath implies use of initplans, so cannot be parallel-safe */
-    pathnode->path.parallel_safe = false;
+    pathnode->path.parallel_safe = true;    /* might change below */
     pathnode->path.parallel_workers = 0;
     /* Result is one unordered row */
     pathnode->path.rows = 1;
@@ -3358,13 +3357,15 @@ create_minmaxagg_path(PlannerInfo *root,
     pathnode->mmaggregates = mmaggregates;
     pathnode->quals = quals;

-    /* Calculate cost of all the initplans ... */
+    /* Calculate cost of all the initplans, and check parallel safety */
     initplan_cost = 0;
     foreach(lc, mmaggregates)
     {
         MinMaxAggInfo *mminfo = (MinMaxAggInfo *) lfirst(lc);

         initplan_cost += mminfo->pathcost;
+        if (!mminfo->path->parallel_safe)
+            pathnode->path.parallel_safe = false;
     }

     /* add tlist eval cost for each output row, plus cpu_tuple_cost */
@@ -3385,6 +3386,17 @@ create_minmaxagg_path(PlannerInfo *root,
         pathnode->path.total_cost += qual_cost.startup + qual_cost.per_tuple;
     }

+    /*
+     * If the initplans were all parallel-safe, also check safety of the
+     * target and quals.  (The Result node itself isn't parallelizable, but if
+     * we are in a subquery then it can be useful for the outer query to know
+     * that this one is parallel-safe.)
+     */
+    if (pathnode->path.parallel_safe)
+        pathnode->path.parallel_safe =
+            is_parallel_safe(root, (Node *) target->exprs) &&
+            is_parallel_safe(root, (Node *) quals);
+
     return pathnode;
 }

diff --git a/src/include/optimizer/subselect.h b/src/include/optimizer/subselect.h
index c03ffc56bf..44bc0bda7e 100644
--- a/src/include/optimizer/subselect.h
+++ b/src/include/optimizer/subselect.h
@@ -28,6 +28,9 @@ extern Node *SS_replace_correlation_vars(PlannerInfo *root, Node *expr);
 extern Node *SS_process_sublinks(PlannerInfo *root, Node *expr, bool isQual);
 extern void SS_identify_outer_params(PlannerInfo *root);
 extern void SS_charge_for_initplans(PlannerInfo *root, RelOptInfo *final_rel);
+extern void SS_compute_initplan_cost(List *init_plans,
+                                     Cost *initplan_cost_p,
+                                     bool *unsafe_initplans_p);
 extern void SS_attach_initplans(PlannerInfo *root, Plan *plan);
 extern void SS_finalize_plan(PlannerInfo *root, Plan *plan);
 extern Param *SS_make_initplan_output_param(PlannerInfo *root,
diff --git a/src/test/regress/expected/partition_prune.out b/src/test/regress/expected/partition_prune.out
index d700c00629..f61fbb2809 100644
--- a/src/test/regress/expected/partition_prune.out
+++ b/src/test/regress/expected/partition_prune.out
@@ -3698,30 +3698,25 @@ select explain_parallel_append(
 select * from listp where a = (select 2);');
                               explain_parallel_append
 -----------------------------------------------------------------------------------
- Append (actual rows=N loops=N)
-   ->  Gather (actual rows=N loops=N)
-         Workers Planned: 2
-         Params Evaluated: $0
-         Workers Launched: N
-         InitPlan 1 (returns $0)
-           ->  Result (actual rows=N loops=N)
-         ->  Parallel Append (actual rows=N loops=N)
-               ->  Seq Scan on listp_12_1 listp_1 (actual rows=N loops=N)
-                     Filter: (a = $0)
-               ->  Parallel Seq Scan on listp_12_2 listp_2 (never executed)
-                     Filter: (a = $0)
-   ->  Gather (actual rows=N loops=N)
-         Workers Planned: 2
-         Params Evaluated: $1
-         Workers Launched: N
-         InitPlan 2 (returns $1)
-           ->  Result (actual rows=N loops=N)
+ Gather (actual rows=N loops=N)
+   Workers Planned: 2
+   Workers Launched: N
+   ->  Parallel Append (actual rows=N loops=N)
          ->  Parallel Append (actual rows=N loops=N)
-               ->  Seq Scan on listp_12_1 listp_4 (never executed)
+               InitPlan 2 (returns $1)
+                 ->  Result (actual rows=N loops=N)
+               ->  Seq Scan on listp_12_1 listp_1 (never executed)
                      Filter: (a = $1)
-               ->  Parallel Seq Scan on listp_12_2 listp_5 (actual rows=N loops=N)
+               ->  Parallel Seq Scan on listp_12_2 listp_2 (actual rows=N loops=N)
                      Filter: (a = $1)
-(23 rows)
+         ->  Parallel Append (actual rows=N loops=N)
+               InitPlan 1 (returns $0)
+                 ->  Result (actual rows=N loops=N)
+               ->  Seq Scan on listp_12_1 listp_4 (actual rows=N loops=N)
+                     Filter: (a = $0)
+               ->  Parallel Seq Scan on listp_12_2 listp_5 (never executed)
+                     Filter: (a = $0)
+(18 rows)

 drop table listp;
 reset parallel_tuple_cost;