nbtree: fix read page recheck typo.
[pgsql.git] / src / backend / executor / nodeGatherMerge.c
blobbc99c0b4485a4766fd9f4e4bb0634914877f0176
1 /*-------------------------------------------------------------------------
3 * nodeGatherMerge.c
4 * Scan a plan in multiple workers, and do order-preserving merge.
6 * Portions Copyright (c) 1996-2024, PostgreSQL Global Development Group
7 * Portions Copyright (c) 1994, Regents of the University of California
9 * IDENTIFICATION
10 * src/backend/executor/nodeGatherMerge.c
12 *-------------------------------------------------------------------------
15 #include "postgres.h"
17 #include "executor/executor.h"
18 #include "executor/execParallel.h"
19 #include "executor/nodeGatherMerge.h"
20 #include "executor/tqueue.h"
21 #include "lib/binaryheap.h"
22 #include "miscadmin.h"
23 #include "optimizer/optimizer.h"
26 * When we read tuples from workers, it's a good idea to read several at once
27 * for efficiency when possible: this minimizes context-switching overhead.
28 * But reading too many at a time wastes memory without improving performance.
29 * We'll read up to MAX_TUPLE_STORE tuples (in addition to the first one).
31 #define MAX_TUPLE_STORE 10
34 * Pending-tuple array for each worker. This holds additional tuples that
35 * we were able to fetch from the worker, but can't process yet. In addition,
36 * this struct holds the "done" flag indicating the worker is known to have
37 * no more tuples. (We do not use this struct for the leader; we don't keep
38 * any pending tuples for the leader, and the need_to_scan_locally flag serves
39 * as its "done" indicator.)
41 typedef struct GMReaderTupleBuffer
43 MinimalTuple *tuple; /* array of length MAX_TUPLE_STORE */
44 int nTuples; /* number of tuples currently stored */
45 int readCounter; /* index of next tuple to extract */
46 bool done; /* true if reader is known exhausted */
47 } GMReaderTupleBuffer;
49 static TupleTableSlot *ExecGatherMerge(PlanState *pstate);
50 static int32 heap_compare_slots(Datum a, Datum b, void *arg);
51 static TupleTableSlot *gather_merge_getnext(GatherMergeState *gm_state);
52 static MinimalTuple gm_readnext_tuple(GatherMergeState *gm_state, int nreader,
53 bool nowait, bool *done);
54 static void ExecShutdownGatherMergeWorkers(GatherMergeState *node);
55 static void gather_merge_setup(GatherMergeState *gm_state);
56 static void gather_merge_init(GatherMergeState *gm_state);
57 static void gather_merge_clear_tuples(GatherMergeState *gm_state);
58 static bool gather_merge_readnext(GatherMergeState *gm_state, int reader,
59 bool nowait);
60 static void load_tuple_array(GatherMergeState *gm_state, int reader);
62 /* ----------------------------------------------------------------
63 * ExecInitGather
64 * ----------------------------------------------------------------
66 GatherMergeState *
67 ExecInitGatherMerge(GatherMerge *node, EState *estate, int eflags)
69 GatherMergeState *gm_state;
70 Plan *outerNode;
71 TupleDesc tupDesc;
73 /* Gather merge node doesn't have innerPlan node. */
74 Assert(innerPlan(node) == NULL);
77 * create state structure
79 gm_state = makeNode(GatherMergeState);
80 gm_state->ps.plan = (Plan *) node;
81 gm_state->ps.state = estate;
82 gm_state->ps.ExecProcNode = ExecGatherMerge;
84 gm_state->initialized = false;
85 gm_state->gm_initialized = false;
86 gm_state->tuples_needed = -1;
89 * Miscellaneous initialization
91 * create expression context for node
93 ExecAssignExprContext(estate, &gm_state->ps);
96 * GatherMerge doesn't support checking a qual (it's always more efficient
97 * to do it in the child node).
99 Assert(!node->plan.qual);
102 * now initialize outer plan
104 outerNode = outerPlan(node);
105 outerPlanState(gm_state) = ExecInitNode(outerNode, estate, eflags);
108 * Leader may access ExecProcNode result directly (if
109 * need_to_scan_locally), or from workers via tuple queue. So we can't
110 * trivially rely on the slot type being fixed for expressions evaluated
111 * within this node.
113 gm_state->ps.outeropsset = true;
114 gm_state->ps.outeropsfixed = false;
117 * Store the tuple descriptor into gather merge state, so we can use it
118 * while initializing the gather merge slots.
120 tupDesc = ExecGetResultType(outerPlanState(gm_state));
121 gm_state->tupDesc = tupDesc;
124 * Initialize result type and projection.
126 ExecInitResultTypeTL(&gm_state->ps);
127 ExecConditionalAssignProjectionInfo(&gm_state->ps, tupDesc, OUTER_VAR);
130 * Without projections result slot type is not trivially known, see
131 * comment above.
133 if (gm_state->ps.ps_ProjInfo == NULL)
135 gm_state->ps.resultopsset = true;
136 gm_state->ps.resultopsfixed = false;
140 * initialize sort-key information
142 if (node->numCols)
144 int i;
146 gm_state->gm_nkeys = node->numCols;
147 gm_state->gm_sortkeys =
148 palloc0(sizeof(SortSupportData) * node->numCols);
150 for (i = 0; i < node->numCols; i++)
152 SortSupport sortKey = gm_state->gm_sortkeys + i;
154 sortKey->ssup_cxt = CurrentMemoryContext;
155 sortKey->ssup_collation = node->collations[i];
156 sortKey->ssup_nulls_first = node->nullsFirst[i];
157 sortKey->ssup_attno = node->sortColIdx[i];
160 * We don't perform abbreviated key conversion here, for the same
161 * reasons that it isn't used in MergeAppend
163 sortKey->abbreviate = false;
165 PrepareSortSupportFromOrderingOp(node->sortOperators[i], sortKey);
169 /* Now allocate the workspace for gather merge */
170 gather_merge_setup(gm_state);
172 return gm_state;
175 /* ----------------------------------------------------------------
176 * ExecGatherMerge(node)
178 * Scans the relation via multiple workers and returns
179 * the next qualifying tuple.
180 * ----------------------------------------------------------------
182 static TupleTableSlot *
183 ExecGatherMerge(PlanState *pstate)
185 GatherMergeState *node = castNode(GatherMergeState, pstate);
186 TupleTableSlot *slot;
187 ExprContext *econtext;
189 CHECK_FOR_INTERRUPTS();
192 * As with Gather, we don't launch workers until this node is actually
193 * executed.
195 if (!node->initialized)
197 EState *estate = node->ps.state;
198 GatherMerge *gm = castNode(GatherMerge, node->ps.plan);
201 * Sometimes we might have to run without parallelism; but if parallel
202 * mode is active then we can try to fire up some workers.
204 if (gm->num_workers > 0 && estate->es_use_parallel_mode)
206 ParallelContext *pcxt;
208 /* Initialize, or re-initialize, shared state needed by workers. */
209 if (!node->pei)
210 node->pei = ExecInitParallelPlan(outerPlanState(node),
211 estate,
212 gm->initParam,
213 gm->num_workers,
214 node->tuples_needed);
215 else
216 ExecParallelReinitialize(outerPlanState(node),
217 node->pei,
218 gm->initParam);
220 /* Try to launch workers. */
221 pcxt = node->pei->pcxt;
222 LaunchParallelWorkers(pcxt);
223 /* We save # workers launched for the benefit of EXPLAIN */
224 node->nworkers_launched = pcxt->nworkers_launched;
227 * Count number of workers originally wanted and actually
228 * launched.
230 estate->es_parallel_workers_to_launch += pcxt->nworkers_to_launch;
231 estate->es_parallel_workers_launched += pcxt->nworkers_launched;
233 /* Set up tuple queue readers to read the results. */
234 if (pcxt->nworkers_launched > 0)
236 ExecParallelCreateReaders(node->pei);
237 /* Make a working array showing the active readers */
238 node->nreaders = pcxt->nworkers_launched;
239 node->reader = (TupleQueueReader **)
240 palloc(node->nreaders * sizeof(TupleQueueReader *));
241 memcpy(node->reader, node->pei->reader,
242 node->nreaders * sizeof(TupleQueueReader *));
244 else
246 /* No workers? Then never mind. */
247 node->nreaders = 0;
248 node->reader = NULL;
252 /* allow leader to participate if enabled or no choice */
253 if (parallel_leader_participation || node->nreaders == 0)
254 node->need_to_scan_locally = true;
255 node->initialized = true;
259 * Reset per-tuple memory context to free any expression evaluation
260 * storage allocated in the previous tuple cycle.
262 econtext = node->ps.ps_ExprContext;
263 ResetExprContext(econtext);
266 * Get next tuple, either from one of our workers, or by running the plan
267 * ourselves.
269 slot = gather_merge_getnext(node);
270 if (TupIsNull(slot))
271 return NULL;
273 /* If no projection is required, we're done. */
274 if (node->ps.ps_ProjInfo == NULL)
275 return slot;
278 * Form the result tuple using ExecProject(), and return it.
280 econtext->ecxt_outertuple = slot;
281 return ExecProject(node->ps.ps_ProjInfo);
284 /* ----------------------------------------------------------------
285 * ExecEndGatherMerge
287 * frees any storage allocated through C routines.
288 * ----------------------------------------------------------------
290 void
291 ExecEndGatherMerge(GatherMergeState *node)
293 ExecEndNode(outerPlanState(node)); /* let children clean up first */
294 ExecShutdownGatherMerge(node);
297 /* ----------------------------------------------------------------
298 * ExecShutdownGatherMerge
300 * Destroy the setup for parallel workers including parallel context.
301 * ----------------------------------------------------------------
303 void
304 ExecShutdownGatherMerge(GatherMergeState *node)
306 ExecShutdownGatherMergeWorkers(node);
308 /* Now destroy the parallel context. */
309 if (node->pei != NULL)
311 ExecParallelCleanup(node->pei);
312 node->pei = NULL;
316 /* ----------------------------------------------------------------
317 * ExecShutdownGatherMergeWorkers
319 * Stop all the parallel workers.
320 * ----------------------------------------------------------------
322 static void
323 ExecShutdownGatherMergeWorkers(GatherMergeState *node)
325 if (node->pei != NULL)
326 ExecParallelFinish(node->pei);
328 /* Flush local copy of reader array */
329 if (node->reader)
330 pfree(node->reader);
331 node->reader = NULL;
334 /* ----------------------------------------------------------------
335 * ExecReScanGatherMerge
337 * Prepare to re-scan the result of a GatherMerge.
338 * ----------------------------------------------------------------
340 void
341 ExecReScanGatherMerge(GatherMergeState *node)
343 GatherMerge *gm = (GatherMerge *) node->ps.plan;
344 PlanState *outerPlan = outerPlanState(node);
346 /* Make sure any existing workers are gracefully shut down */
347 ExecShutdownGatherMergeWorkers(node);
349 /* Free any unused tuples, so we don't leak memory across rescans */
350 gather_merge_clear_tuples(node);
352 /* Mark node so that shared state will be rebuilt at next call */
353 node->initialized = false;
354 node->gm_initialized = false;
357 * Set child node's chgParam to tell it that the next scan might deliver a
358 * different set of rows within the leader process. (The overall rowset
359 * shouldn't change, but the leader process's subset might; hence nodes
360 * between here and the parallel table scan node mustn't optimize on the
361 * assumption of an unchanging rowset.)
363 if (gm->rescan_param >= 0)
364 outerPlan->chgParam = bms_add_member(outerPlan->chgParam,
365 gm->rescan_param);
368 * If chgParam of subnode is not null then plan will be re-scanned by
369 * first ExecProcNode. Note: because this does nothing if we have a
370 * rescan_param, it's currently guaranteed that parallel-aware child nodes
371 * will not see a ReScan call until after they get a ReInitializeDSM call.
372 * That ordering might not be something to rely on, though. A good rule
373 * of thumb is that ReInitializeDSM should reset only shared state, ReScan
374 * should reset only local state, and anything that depends on both of
375 * those steps being finished must wait until the first ExecProcNode call.
377 if (outerPlan->chgParam == NULL)
378 ExecReScan(outerPlan);
382 * Set up the data structures that we'll need for Gather Merge.
384 * We allocate these once on the basis of gm->num_workers, which is an
385 * upper bound for the number of workers we'll actually have. During
386 * a rescan, we reset the structures to empty. This approach simplifies
387 * not leaking memory across rescans.
389 * In the gm_slots[] array, index 0 is for the leader, and indexes 1 to n
390 * are for workers. The values placed into gm_heap correspond to indexes
391 * in gm_slots[]. The gm_tuple_buffers[] array, however, is indexed from
392 * 0 to n-1; it has no entry for the leader.
394 static void
395 gather_merge_setup(GatherMergeState *gm_state)
397 GatherMerge *gm = castNode(GatherMerge, gm_state->ps.plan);
398 int nreaders = gm->num_workers;
399 int i;
402 * Allocate gm_slots for the number of workers + one more slot for leader.
403 * Slot 0 is always for the leader. Leader always calls ExecProcNode() to
404 * read the tuple, and then stores it directly into its gm_slots entry.
405 * For other slots, code below will call ExecInitExtraTupleSlot() to
406 * create a slot for the worker's results. Note that during any single
407 * scan, we might have fewer than num_workers available workers, in which
408 * case the extra array entries go unused.
410 gm_state->gm_slots = (TupleTableSlot **)
411 palloc0((nreaders + 1) * sizeof(TupleTableSlot *));
413 /* Allocate the tuple slot and tuple array for each worker */
414 gm_state->gm_tuple_buffers = (GMReaderTupleBuffer *)
415 palloc0(nreaders * sizeof(GMReaderTupleBuffer));
417 for (i = 0; i < nreaders; i++)
419 /* Allocate the tuple array with length MAX_TUPLE_STORE */
420 gm_state->gm_tuple_buffers[i].tuple =
421 (MinimalTuple *) palloc0(sizeof(MinimalTuple) * MAX_TUPLE_STORE);
423 /* Initialize tuple slot for worker */
424 gm_state->gm_slots[i + 1] =
425 ExecInitExtraTupleSlot(gm_state->ps.state, gm_state->tupDesc,
426 &TTSOpsMinimalTuple);
429 /* Allocate the resources for the merge */
430 gm_state->gm_heap = binaryheap_allocate(nreaders + 1,
431 heap_compare_slots,
432 gm_state);
436 * Initialize the Gather Merge.
438 * Reset data structures to ensure they're empty. Then pull at least one
439 * tuple from leader + each worker (or set its "done" indicator), and set up
440 * the heap.
442 static void
443 gather_merge_init(GatherMergeState *gm_state)
445 int nreaders = gm_state->nreaders;
446 bool nowait = true;
447 int i;
449 /* Assert that gather_merge_setup made enough space */
450 Assert(nreaders <= castNode(GatherMerge, gm_state->ps.plan)->num_workers);
452 /* Reset leader's tuple slot to empty */
453 gm_state->gm_slots[0] = NULL;
455 /* Reset the tuple slot and tuple array for each worker */
456 for (i = 0; i < nreaders; i++)
458 /* Reset tuple array to empty */
459 gm_state->gm_tuple_buffers[i].nTuples = 0;
460 gm_state->gm_tuple_buffers[i].readCounter = 0;
461 /* Reset done flag to not-done */
462 gm_state->gm_tuple_buffers[i].done = false;
463 /* Ensure output slot is empty */
464 ExecClearTuple(gm_state->gm_slots[i + 1]);
467 /* Reset binary heap to empty */
468 binaryheap_reset(gm_state->gm_heap);
471 * First, try to read a tuple from each worker (including leader) in
472 * nowait mode. After this, if not all workers were able to produce a
473 * tuple (or a "done" indication), then re-read from remaining workers,
474 * this time using wait mode. Add all live readers (those producing at
475 * least one tuple) to the heap.
477 reread:
478 for (i = 0; i <= nreaders; i++)
480 CHECK_FOR_INTERRUPTS();
482 /* skip this source if already known done */
483 if ((i == 0) ? gm_state->need_to_scan_locally :
484 !gm_state->gm_tuple_buffers[i - 1].done)
486 if (TupIsNull(gm_state->gm_slots[i]))
488 /* Don't have a tuple yet, try to get one */
489 if (gather_merge_readnext(gm_state, i, nowait))
490 binaryheap_add_unordered(gm_state->gm_heap,
491 Int32GetDatum(i));
493 else
496 * We already got at least one tuple from this worker, but
497 * might as well see if it has any more ready by now.
499 load_tuple_array(gm_state, i);
504 /* need not recheck leader, since nowait doesn't matter for it */
505 for (i = 1; i <= nreaders; i++)
507 if (!gm_state->gm_tuple_buffers[i - 1].done &&
508 TupIsNull(gm_state->gm_slots[i]))
510 nowait = false;
511 goto reread;
515 /* Now heapify the heap. */
516 binaryheap_build(gm_state->gm_heap);
518 gm_state->gm_initialized = true;
522 * Clear out the tuple table slot, and any unused pending tuples,
523 * for each gather merge input.
525 static void
526 gather_merge_clear_tuples(GatherMergeState *gm_state)
528 int i;
530 for (i = 0; i < gm_state->nreaders; i++)
532 GMReaderTupleBuffer *tuple_buffer = &gm_state->gm_tuple_buffers[i];
534 while (tuple_buffer->readCounter < tuple_buffer->nTuples)
535 pfree(tuple_buffer->tuple[tuple_buffer->readCounter++]);
537 ExecClearTuple(gm_state->gm_slots[i + 1]);
542 * Read the next tuple for gather merge.
544 * Fetch the sorted tuple out of the heap.
546 static TupleTableSlot *
547 gather_merge_getnext(GatherMergeState *gm_state)
549 int i;
551 if (!gm_state->gm_initialized)
554 * First time through: pull the first tuple from each participant, and
555 * set up the heap.
557 gather_merge_init(gm_state);
559 else
562 * Otherwise, pull the next tuple from whichever participant we
563 * returned from last time, and reinsert that participant's index into
564 * the heap, because it might now compare differently against the
565 * other elements of the heap.
567 i = DatumGetInt32(binaryheap_first(gm_state->gm_heap));
569 if (gather_merge_readnext(gm_state, i, false))
570 binaryheap_replace_first(gm_state->gm_heap, Int32GetDatum(i));
571 else
573 /* reader exhausted, remove it from heap */
574 (void) binaryheap_remove_first(gm_state->gm_heap);
578 if (binaryheap_empty(gm_state->gm_heap))
580 /* All the queues are exhausted, and so is the heap */
581 gather_merge_clear_tuples(gm_state);
582 return NULL;
584 else
586 /* Return next tuple from whichever participant has the leading one */
587 i = DatumGetInt32(binaryheap_first(gm_state->gm_heap));
588 return gm_state->gm_slots[i];
593 * Read tuple(s) for given reader in nowait mode, and load into its tuple
594 * array, until we have MAX_TUPLE_STORE of them or would have to block.
596 static void
597 load_tuple_array(GatherMergeState *gm_state, int reader)
599 GMReaderTupleBuffer *tuple_buffer;
600 int i;
602 /* Don't do anything if this is the leader. */
603 if (reader == 0)
604 return;
606 tuple_buffer = &gm_state->gm_tuple_buffers[reader - 1];
608 /* If there's nothing in the array, reset the counters to zero. */
609 if (tuple_buffer->nTuples == tuple_buffer->readCounter)
610 tuple_buffer->nTuples = tuple_buffer->readCounter = 0;
612 /* Try to fill additional slots in the array. */
613 for (i = tuple_buffer->nTuples; i < MAX_TUPLE_STORE; i++)
615 MinimalTuple tuple;
617 tuple = gm_readnext_tuple(gm_state,
618 reader,
619 true,
620 &tuple_buffer->done);
621 if (!tuple)
622 break;
623 tuple_buffer->tuple[i] = tuple;
624 tuple_buffer->nTuples++;
629 * Store the next tuple for a given reader into the appropriate slot.
631 * Returns true if successful, false if not (either reader is exhausted,
632 * or we didn't want to wait for a tuple). Sets done flag if reader
633 * is found to be exhausted.
635 static bool
636 gather_merge_readnext(GatherMergeState *gm_state, int reader, bool nowait)
638 GMReaderTupleBuffer *tuple_buffer;
639 MinimalTuple tup;
642 * If we're being asked to generate a tuple from the leader, then we just
643 * call ExecProcNode as normal to produce one.
645 if (reader == 0)
647 if (gm_state->need_to_scan_locally)
649 PlanState *outerPlan = outerPlanState(gm_state);
650 TupleTableSlot *outerTupleSlot;
651 EState *estate = gm_state->ps.state;
653 /* Install our DSA area while executing the plan. */
654 estate->es_query_dsa = gm_state->pei ? gm_state->pei->area : NULL;
655 outerTupleSlot = ExecProcNode(outerPlan);
656 estate->es_query_dsa = NULL;
658 if (!TupIsNull(outerTupleSlot))
660 gm_state->gm_slots[0] = outerTupleSlot;
661 return true;
663 /* need_to_scan_locally serves as "done" flag for leader */
664 gm_state->need_to_scan_locally = false;
666 return false;
669 /* Otherwise, check the state of the relevant tuple buffer. */
670 tuple_buffer = &gm_state->gm_tuple_buffers[reader - 1];
672 if (tuple_buffer->nTuples > tuple_buffer->readCounter)
674 /* Return any tuple previously read that is still buffered. */
675 tup = tuple_buffer->tuple[tuple_buffer->readCounter++];
677 else if (tuple_buffer->done)
679 /* Reader is known to be exhausted. */
680 return false;
682 else
684 /* Read and buffer next tuple. */
685 tup = gm_readnext_tuple(gm_state,
686 reader,
687 nowait,
688 &tuple_buffer->done);
689 if (!tup)
690 return false;
693 * Attempt to read more tuples in nowait mode and store them in the
694 * pending-tuple array for the reader.
696 load_tuple_array(gm_state, reader);
699 Assert(tup);
701 /* Build the TupleTableSlot for the given tuple */
702 ExecStoreMinimalTuple(tup, /* tuple to store */
703 gm_state->gm_slots[reader], /* slot in which to
704 * store the tuple */
705 true); /* pfree tuple when done with it */
707 return true;
711 * Attempt to read a tuple from given worker.
713 static MinimalTuple
714 gm_readnext_tuple(GatherMergeState *gm_state, int nreader, bool nowait,
715 bool *done)
717 TupleQueueReader *reader;
718 MinimalTuple tup;
720 /* Check for async events, particularly messages from workers. */
721 CHECK_FOR_INTERRUPTS();
724 * Attempt to read a tuple.
726 * Note that TupleQueueReaderNext will just return NULL for a worker which
727 * fails to initialize. We'll treat that worker as having produced no
728 * tuples; WaitForParallelWorkersToFinish will error out when we get
729 * there.
731 reader = gm_state->reader[nreader - 1];
732 tup = TupleQueueReaderNext(reader, nowait, done);
735 * Since we'll be buffering these across multiple calls, we need to make a
736 * copy.
738 return tup ? heap_copy_minimal_tuple(tup) : NULL;
742 * We have one slot for each item in the heap array. We use SlotNumber
743 * to store slot indexes. This doesn't actually provide any formal
744 * type-safety, but it makes the code more self-documenting.
746 typedef int32 SlotNumber;
749 * Compare the tuples in the two given slots.
751 static int32
752 heap_compare_slots(Datum a, Datum b, void *arg)
754 GatherMergeState *node = (GatherMergeState *) arg;
755 SlotNumber slot1 = DatumGetInt32(a);
756 SlotNumber slot2 = DatumGetInt32(b);
758 TupleTableSlot *s1 = node->gm_slots[slot1];
759 TupleTableSlot *s2 = node->gm_slots[slot2];
760 int nkey;
762 Assert(!TupIsNull(s1));
763 Assert(!TupIsNull(s2));
765 for (nkey = 0; nkey < node->gm_nkeys; nkey++)
767 SortSupport sortKey = node->gm_sortkeys + nkey;
768 AttrNumber attno = sortKey->ssup_attno;
769 Datum datum1,
770 datum2;
771 bool isNull1,
772 isNull2;
773 int compare;
775 datum1 = slot_getattr(s1, attno, &isNull1);
776 datum2 = slot_getattr(s2, attno, &isNull2);
778 compare = ApplySortComparator(datum1, isNull1,
779 datum2, isNull2,
780 sortKey);
781 if (compare != 0)
783 INVERT_COMPARE_RESULT(compare);
784 return compare;
787 return 0;