/src/postgres/src/backend/executor/nodeGatherMerge.c

Source
/*-------------------------------------------------------------------------
 *
 * nodeGatherMerge.c
 *    Scan a plan in multiple workers, and do order-preserving merge.
 *
 * Portions Copyright (c) 1996-2025, PostgreSQL Global Development Group
 * Portions Copyright (c) 1994, Regents of the University of California
 *
 * IDENTIFICATION
 *    src/backend/executor/nodeGatherMerge.c
 *
 *-------------------------------------------------------------------------
 */

#include "postgres.h"

#include "access/htup_details.h"
#include "executor/executor.h"
#include "executor/execParallel.h"
#include "executor/nodeGatherMerge.h"
#include "executor/tqueue.h"
#include "lib/binaryheap.h"
#include "miscadmin.h"
#include "optimizer/optimizer.h"

/*
 * When we read tuples from workers, it's a good idea to read several at once
 * for efficiency when possible: this minimizes context-switching overhead.
 * But reading too many at a time wastes memory without improving performance.
 * We'll read up to MAX_TUPLE_STORE tuples (in addition to the first one).
 */
#define MAX_TUPLE_STORE 10

/*
 * Pending-tuple array for each worker.  This holds additional tuples that
 * we were able to fetch from the worker, but can't process yet.  In addition,
 * this struct holds the "done" flag indicating the worker is known to have
 * no more tuples.  (We do not use this struct for the leader; we don't keep
 * any pending tuples for the leader, and the need_to_scan_locally flag serves
 * as its "done" indicator.)
 */
typedef struct GMReaderTupleBuffer
{
  MinimalTuple *tuple;    /* array of length MAX_TUPLE_STORE */
  int     nTuples;    /* number of tuples currently stored */
  int     readCounter;  /* index of next tuple to extract */
  bool    done;     /* true if reader is known exhausted */
} GMReaderTupleBuffer;

static TupleTableSlot *ExecGatherMerge(PlanState *pstate);
static int32 heap_compare_slots(Datum a, Datum b, void *arg);
static TupleTableSlot *gather_merge_getnext(GatherMergeState *gm_state);
static MinimalTuple gm_readnext_tuple(GatherMergeState *gm_state, int nreader,
                    bool nowait, bool *done);
static void ExecShutdownGatherMergeWorkers(GatherMergeState *node);
static void gather_merge_setup(GatherMergeState *gm_state);
static void gather_merge_init(GatherMergeState *gm_state);
static void gather_merge_clear_tuples(GatherMergeState *gm_state);
static bool gather_merge_readnext(GatherMergeState *gm_state, int reader,
                  bool nowait);
static void load_tuple_array(GatherMergeState *gm_state, int reader);

/* ----------------------------------------------------------------
 *    ExecInitGather
 * ----------------------------------------------------------------
 */
GatherMergeState *
ExecInitGatherMerge(GatherMerge *node, EState *estate, int eflags)
{
  GatherMergeState *gm_state;
  Plan     *outerNode;
  TupleDesc tupDesc;

  /* Gather merge node doesn't have innerPlan node. */
  Assert(innerPlan(node) == NULL);

  /*
   * create state structure
   */
  gm_state = makeNode(GatherMergeState);
  gm_state->ps.plan = (Plan *) node;
  gm_state->ps.state = estate;
  gm_state->ps.ExecProcNode = ExecGatherMerge;

  gm_state->initialized = false;
  gm_state->gm_initialized = false;
  gm_state->tuples_needed = -1;

  /*
   * Miscellaneous initialization
   *
   * create expression context for node
   */
  ExecAssignExprContext(estate, &gm_state->ps);

  /*
   * GatherMerge doesn't support checking a qual (it's always more efficient
   * to do it in the child node).
   */
  Assert(!node->plan.qual);

  /*
   * now initialize outer plan
   */
  outerNode = outerPlan(node);
  outerPlanState(gm_state) = ExecInitNode(outerNode, estate, eflags);

  /*
   * Leader may access ExecProcNode result directly (if
   * need_to_scan_locally), or from workers via tuple queue.  So we can't
   * trivially rely on the slot type being fixed for expressions evaluated
   * within this node.
   */
  gm_state->ps.outeropsset = true;
  gm_state->ps.outeropsfixed = false;

  /*
   * Store the tuple descriptor into gather merge state, so we can use it
   * while initializing the gather merge slots.
   */
  tupDesc = ExecGetResultType(outerPlanState(gm_state));
  gm_state->tupDesc = tupDesc;

  /*
   * Initialize result type and projection.
   */
  ExecInitResultTypeTL(&gm_state->ps);
  ExecConditionalAssignProjectionInfo(&gm_state->ps, tupDesc, OUTER_VAR);

  /*
   * Without projections result slot type is not trivially known, see
   * comment above.
   */
  if (gm_state->ps.ps_ProjInfo == NULL)
  {
    gm_state->ps.resultopsset = true;
    gm_state->ps.resultopsfixed = false;
  }

  /*
   * initialize sort-key information
   */
  if (node->numCols)
  {
    int     i;

    gm_state->gm_nkeys = node->numCols;
    gm_state->gm_sortkeys =
      palloc0(sizeof(SortSupportData) * node->numCols);

    for (i = 0; i < node->numCols; i++)
    {
      SortSupport sortKey = gm_state->gm_sortkeys + i;

      sortKey->ssup_cxt = CurrentMemoryContext;
      sortKey->ssup_collation = node->collations[i];
      sortKey->ssup_nulls_first = node->nullsFirst[i];
      sortKey->ssup_attno = node->sortColIdx[i];

      /*
       * We don't perform abbreviated key conversion here, for the same
       * reasons that it isn't used in MergeAppend
       */
      sortKey->abbreviate = false;

      PrepareSortSupportFromOrderingOp(node->sortOperators[i], sortKey);
    }
  }

  /* Now allocate the workspace for gather merge */
  gather_merge_setup(gm_state);

  return gm_state;
}

/* ----------------------------------------------------------------
 *    ExecGatherMerge(node)
 *
 *    Scans the relation via multiple workers and returns
 *    the next qualifying tuple.
 * ----------------------------------------------------------------
 */
static TupleTableSlot *
ExecGatherMerge(PlanState *pstate)
{
  GatherMergeState *node = castNode(GatherMergeState, pstate);
  TupleTableSlot *slot;
  ExprContext *econtext;

  CHECK_FOR_INTERRUPTS();

  /*
   * As with Gather, we don't launch workers until this node is actually
   * executed.
   */
  if (!node->initialized)
  {
    EState     *estate = node->ps.state;
    GatherMerge *gm = castNode(GatherMerge, node->ps.plan);

    /*
     * Sometimes we might have to run without parallelism; but if parallel
     * mode is active then we can try to fire up some workers.
     */
    if (gm->num_workers > 0 && estate->es_use_parallel_mode)
    {
      ParallelContext *pcxt;

      /* Initialize, or re-initialize, shared state needed by workers. */
      if (!node->pei)
        node->pei = ExecInitParallelPlan(outerPlanState(node),
                         estate,
                         gm->initParam,
                         gm->num_workers,
                         node->tuples_needed);
      else
        ExecParallelReinitialize(outerPlanState(node),
                     node->pei,
                     gm->initParam);

      /* Try to launch workers. */
      pcxt = node->pei->pcxt;
      LaunchParallelWorkers(pcxt);
      /* We save # workers launched for the benefit of EXPLAIN */
      node->nworkers_launched = pcxt->nworkers_launched;

      /*
       * Count number of workers originally wanted and actually
       * launched.
       */
      estate->es_parallel_workers_to_launch += pcxt->nworkers_to_launch;
      estate->es_parallel_workers_launched += pcxt->nworkers_launched;

      /* Set up tuple queue readers to read the results. */
      if (pcxt->nworkers_launched > 0)
      {
        ExecParallelCreateReaders(node->pei);
        /* Make a working array showing the active readers */
        node->nreaders = pcxt->nworkers_launched;
        node->reader = (TupleQueueReader **)
          palloc(node->nreaders * sizeof(TupleQueueReader *));
        memcpy(node->reader, node->pei->reader,
             node->nreaders * sizeof(TupleQueueReader *));
      }
      else
      {
        /* No workers?  Then never mind. */
        node->nreaders = 0;
        node->reader = NULL;
      }
    }

    /* allow leader to participate if enabled or no choice */
    if (parallel_leader_participation || node->nreaders == 0)
      node->need_to_scan_locally = true;
    node->initialized = true;
  }

  /*
   * Reset per-tuple memory context to free any expression evaluation
   * storage allocated in the previous tuple cycle.
   */
  econtext = node->ps.ps_ExprContext;
  ResetExprContext(econtext);

  /*
   * Get next tuple, either from one of our workers, or by running the plan
   * ourselves.
   */
  slot = gather_merge_getnext(node);
  if (TupIsNull(slot))
    return NULL;

  /* If no projection is required, we're done. */
  if (node->ps.ps_ProjInfo == NULL)
    return slot;

  /*
   * Form the result tuple using ExecProject(), and return it.
   */
  econtext->ecxt_outertuple = slot;
  return ExecProject(node->ps.ps_ProjInfo);
}

/* ----------------------------------------------------------------
 *    ExecEndGatherMerge
 *
 *    frees any storage allocated through C routines.
 * ----------------------------------------------------------------
 */
void
ExecEndGatherMerge(GatherMergeState *node)
{
  ExecEndNode(outerPlanState(node));  /* let children clean up first */
  ExecShutdownGatherMerge(node);
}

/* ----------------------------------------------------------------
 *    ExecShutdownGatherMerge
 *
 *    Destroy the setup for parallel workers including parallel context.
 * ----------------------------------------------------------------
 */
void
ExecShutdownGatherMerge(GatherMergeState *node)
{
  ExecShutdownGatherMergeWorkers(node);

  /* Now destroy the parallel context. */
  if (node->pei != NULL)
  {
    ExecParallelCleanup(node->pei);
    node->pei = NULL;
  }
}

/* ----------------------------------------------------------------
 *    ExecShutdownGatherMergeWorkers
 *
 *    Stop all the parallel workers.
 * ----------------------------------------------------------------
 */
static void
ExecShutdownGatherMergeWorkers(GatherMergeState *node)
{
  if (node->pei != NULL)
    ExecParallelFinish(node->pei);

  /* Flush local copy of reader array */
  if (node->reader)
    pfree(node->reader);
  node->reader = NULL;
}

/* ----------------------------------------------------------------
 *    ExecReScanGatherMerge
 *
 *    Prepare to re-scan the result of a GatherMerge.
 * ----------------------------------------------------------------
 */
void
ExecReScanGatherMerge(GatherMergeState *node)
{
  GatherMerge *gm = (GatherMerge *) node->ps.plan;
  PlanState  *outerPlan = outerPlanState(node);

  /* Make sure any existing workers are gracefully shut down */
  ExecShutdownGatherMergeWorkers(node);

  /* Free any unused tuples, so we don't leak memory across rescans */
  gather_merge_clear_tuples(node);

  /* Mark node so that shared state will be rebuilt at next call */
  node->initialized = false;
  node->gm_initialized = false;

  /*
   * Set child node's chgParam to tell it that the next scan might deliver a
   * different set of rows within the leader process.  (The overall rowset
   * shouldn't change, but the leader process's subset might; hence nodes
   * between here and the parallel table scan node mustn't optimize on the
   * assumption of an unchanging rowset.)
   */
  if (gm->rescan_param >= 0)
    outerPlan->chgParam = bms_add_member(outerPlan->chgParam,
                       gm->rescan_param);

  /*
   * If chgParam of subnode is not null then plan will be re-scanned by
   * first ExecProcNode.  Note: because this does nothing if we have a
   * rescan_param, it's currently guaranteed that parallel-aware child nodes
   * will not see a ReScan call until after they get a ReInitializeDSM call.
   * That ordering might not be something to rely on, though.  A good rule
   * of thumb is that ReInitializeDSM should reset only shared state, ReScan
   * should reset only local state, and anything that depends on both of
   * those steps being finished must wait until the first ExecProcNode call.
   */
  if (outerPlan->chgParam == NULL)
    ExecReScan(outerPlan);
}

/*
 * Set up the data structures that we'll need for Gather Merge.
 *
 * We allocate these once on the basis of gm->num_workers, which is an
 * upper bound for the number of workers we'll actually have.  During
 * a rescan, we reset the structures to empty.  This approach simplifies
 * not leaking memory across rescans.
 *
 * In the gm_slots[] array, index 0 is for the leader, and indexes 1 to n
 * are for workers.  The values placed into gm_heap correspond to indexes
 * in gm_slots[].  The gm_tuple_buffers[] array, however, is indexed from
 * 0 to n-1; it has no entry for the leader.
 */
static void
gather_merge_setup(GatherMergeState *gm_state)
{
  GatherMerge *gm = castNode(GatherMerge, gm_state->ps.plan);
  int     nreaders = gm->num_workers;
  int     i;

  /*
   * Allocate gm_slots for the number of workers + one more slot for leader.
   * Slot 0 is always for the leader.  Leader always calls ExecProcNode() to
   * read the tuple, and then stores it directly into its gm_slots entry.
   * For other slots, code below will call ExecInitExtraTupleSlot() to
   * create a slot for the worker's results.  Note that during any single
   * scan, we might have fewer than num_workers available workers, in which
   * case the extra array entries go unused.
   */
  gm_state->gm_slots = (TupleTableSlot **)
    palloc0((nreaders + 1) * sizeof(TupleTableSlot *));

  /* Allocate the tuple slot and tuple array for each worker */
  gm_state->gm_tuple_buffers = (GMReaderTupleBuffer *)
    palloc0(nreaders * sizeof(GMReaderTupleBuffer));

  for (i = 0; i < nreaders; i++)
  {
    /* Allocate the tuple array with length MAX_TUPLE_STORE */
    gm_state->gm_tuple_buffers[i].tuple =
      (MinimalTuple *) palloc0(sizeof(MinimalTuple) * MAX_TUPLE_STORE);

    /* Initialize tuple slot for worker */
    gm_state->gm_slots[i + 1] =
      ExecInitExtraTupleSlot(gm_state->ps.state, gm_state->tupDesc,
                   &TTSOpsMinimalTuple);
  }

  /* Allocate the resources for the merge */
  gm_state->gm_heap = binaryheap_allocate(nreaders + 1,
                      heap_compare_slots,
                      gm_state);
}

/*
 * Initialize the Gather Merge.
 *
 * Reset data structures to ensure they're empty.  Then pull at least one
 * tuple from leader + each worker (or set its "done" indicator), and set up
 * the heap.
 */
static void
gather_merge_init(GatherMergeState *gm_state)
{
  int     nreaders = gm_state->nreaders;
  bool    nowait = true;
  int     i;

  /* Assert that gather_merge_setup made enough space */
  Assert(nreaders <= castNode(GatherMerge, gm_state->ps.plan)->num_workers);

  /* Reset leader's tuple slot to empty */
  gm_state->gm_slots[0] = NULL;

  /* Reset the tuple slot and tuple array for each worker */
  for (i = 0; i < nreaders; i++)
  {
    /* Reset tuple array to empty */
    gm_state->gm_tuple_buffers[i].nTuples = 0;
    gm_state->gm_tuple_buffers[i].readCounter = 0;
    /* Reset done flag to not-done */
    gm_state->gm_tuple_buffers[i].done = false;
    /* Ensure output slot is empty */
    ExecClearTuple(gm_state->gm_slots[i + 1]);
  }

  /* Reset binary heap to empty */
  binaryheap_reset(gm_state->gm_heap);

  /*
   * First, try to read a tuple from each worker (including leader) in
   * nowait mode.  After this, if not all workers were able to produce a
   * tuple (or a "done" indication), then re-read from remaining workers,
   * this time using wait mode.  Add all live readers (those producing at
   * least one tuple) to the heap.
   */
reread:
  for (i = 0; i <= nreaders; i++)
  {
    CHECK_FOR_INTERRUPTS();

    /* skip this source if already known done */
    if ((i == 0) ? gm_state->need_to_scan_locally :
      !gm_state->gm_tuple_buffers[i - 1].done)
    {
      if (TupIsNull(gm_state->gm_slots[i]))
      {
        /* Don't have a tuple yet, try to get one */
        if (gather_merge_readnext(gm_state, i, nowait))
          binaryheap_add_unordered(gm_state->gm_heap,
                       Int32GetDatum(i));
      }
      else
      {
        /*
         * We already got at least one tuple from this worker, but
         * might as well see if it has any more ready by now.
         */
        load_tuple_array(gm_state, i);
      }
    }
  }

  /* need not recheck leader, since nowait doesn't matter for it */
  for (i = 1; i <= nreaders; i++)
  {
    if (!gm_state->gm_tuple_buffers[i - 1].done &&
      TupIsNull(gm_state->gm_slots[i]))
    {
      nowait = false;
      goto reread;
    }
  }

  /* Now heapify the heap. */
  binaryheap_build(gm_state->gm_heap);

  gm_state->gm_initialized = true;
}

/*
 * Clear out the tuple table slot, and any unused pending tuples,
 * for each gather merge input.
 */
static void
gather_merge_clear_tuples(GatherMergeState *gm_state)
{
  int     i;

  for (i = 0; i < gm_state->nreaders; i++)
  {
    GMReaderTupleBuffer *tuple_buffer = &gm_state->gm_tuple_buffers[i];

    while (tuple_buffer->readCounter < tuple_buffer->nTuples)
      pfree(tuple_buffer->tuple[tuple_buffer->readCounter++]);

    ExecClearTuple(gm_state->gm_slots[i + 1]);
  }
}

/*
 * Read the next tuple for gather merge.
 *
 * Fetch the sorted tuple out of the heap.
 */
static TupleTableSlot *
gather_merge_getnext(GatherMergeState *gm_state)
{
  int     i;

  if (!gm_state->gm_initialized)
  {
    /*
     * First time through: pull the first tuple from each participant, and
     * set up the heap.
     */
    gather_merge_init(gm_state);
  }
  else
  {
    /*
     * Otherwise, pull the next tuple from whichever participant we
     * returned from last time, and reinsert that participant's index into
     * the heap, because it might now compare differently against the
     * other elements of the heap.
     */
    i = DatumGetInt32(binaryheap_first(gm_state->gm_heap));

    if (gather_merge_readnext(gm_state, i, false))
      binaryheap_replace_first(gm_state->gm_heap, Int32GetDatum(i));
    else
    {
      /* reader exhausted, remove it from heap */
      (void) binaryheap_remove_first(gm_state->gm_heap);
    }
  }

  if (binaryheap_empty(gm_state->gm_heap))
  {
    /* All the queues are exhausted, and so is the heap */
    gather_merge_clear_tuples(gm_state);
    return NULL;
  }
  else
  {
    /* Return next tuple from whichever participant has the leading one */
    i = DatumGetInt32(binaryheap_first(gm_state->gm_heap));
    return gm_state->gm_slots[i];
  }
}

/*
 * Read tuple(s) for given reader in nowait mode, and load into its tuple
 * array, until we have MAX_TUPLE_STORE of them or would have to block.
 */
static void
load_tuple_array(GatherMergeState *gm_state, int reader)
{
  GMReaderTupleBuffer *tuple_buffer;
  int     i;

  /* Don't do anything if this is the leader. */
  if (reader == 0)
    return;

  tuple_buffer = &gm_state->gm_tuple_buffers[reader - 1];

  /* If there's nothing in the array, reset the counters to zero. */
  if (tuple_buffer->nTuples == tuple_buffer->readCounter)
    tuple_buffer->nTuples = tuple_buffer->readCounter = 0;

  /* Try to fill additional slots in the array. */
  for (i = tuple_buffer->nTuples; i < MAX_TUPLE_STORE; i++)
  {
    MinimalTuple tuple;

    tuple = gm_readnext_tuple(gm_state,
                  reader,
                  true,
                  &tuple_buffer->done);
    if (!tuple)
      break;
    tuple_buffer->tuple[i] = tuple;
    tuple_buffer->nTuples++;
  }
}

/*
 * Store the next tuple for a given reader into the appropriate slot.
 *
 * Returns true if successful, false if not (either reader is exhausted,
 * or we didn't want to wait for a tuple).  Sets done flag if reader
 * is found to be exhausted.
 */
static bool
gather_merge_readnext(GatherMergeState *gm_state, int reader, bool nowait)
{
  GMReaderTupleBuffer *tuple_buffer;
  MinimalTuple tup;

  /*
   * If we're being asked to generate a tuple from the leader, then we just
   * call ExecProcNode as normal to produce one.
   */
  if (reader == 0)
  {
    if (gm_state->need_to_scan_locally)
    {
      PlanState  *outerPlan = outerPlanState(gm_state);
      TupleTableSlot *outerTupleSlot;
      EState     *estate = gm_state->ps.state;

      /* Install our DSA area while executing the plan. */
      estate->es_query_dsa = gm_state->pei ? gm_state->pei->area : NULL;
      outerTupleSlot = ExecProcNode(outerPlan);
      estate->es_query_dsa = NULL;

      if (!TupIsNull(outerTupleSlot))
      {
        gm_state->gm_slots[0] = outerTupleSlot;
        return true;
      }
      /* need_to_scan_locally serves as "done" flag for leader */
      gm_state->need_to_scan_locally = false;
    }
    return false;
  }

  /* Otherwise, check the state of the relevant tuple buffer. */
  tuple_buffer = &gm_state->gm_tuple_buffers[reader - 1];

  if (tuple_buffer->nTuples > tuple_buffer->readCounter)
  {
    /* Return any tuple previously read that is still buffered. */
    tup = tuple_buffer->tuple[tuple_buffer->readCounter++];
  }
  else if (tuple_buffer->done)
  {
    /* Reader is known to be exhausted. */
    return false;
  }
  else
  {
    /* Read and buffer next tuple. */
    tup = gm_readnext_tuple(gm_state,
                reader,
                nowait,
                &tuple_buffer->done);
    if (!tup)
      return false;

    /*
     * Attempt to read more tuples in nowait mode and store them in the
     * pending-tuple array for the reader.
     */
    load_tuple_array(gm_state, reader);
  }

  Assert(tup);

  /* Build the TupleTableSlot for the given tuple */
  ExecStoreMinimalTuple(tup,  /* tuple to store */
              gm_state->gm_slots[reader], /* slot in which to
                             * store the tuple */
              true);  /* pfree tuple when done with it */

  return true;
}

/*
 * Attempt to read a tuple from given worker.
 */
static MinimalTuple
gm_readnext_tuple(GatherMergeState *gm_state, int nreader, bool nowait,
          bool *done)
{
  TupleQueueReader *reader;
  MinimalTuple tup;

  /* Check for async events, particularly messages from workers. */
  CHECK_FOR_INTERRUPTS();

  /*
   * Attempt to read a tuple.
   *
   * Note that TupleQueueReaderNext will just return NULL for a worker which
   * fails to initialize.  We'll treat that worker as having produced no
   * tuples; WaitForParallelWorkersToFinish will error out when we get
   * there.
   */
  reader = gm_state->reader[nreader - 1];
  tup = TupleQueueReaderNext(reader, nowait, done);

  /*
   * Since we'll be buffering these across multiple calls, we need to make a
   * copy.
   */
  return tup ? heap_copy_minimal_tuple(tup, 0) : NULL;
}

/*
 * We have one slot for each item in the heap array.  We use SlotNumber
 * to store slot indexes.  This doesn't actually provide any formal
 * type-safety, but it makes the code more self-documenting.
 */
typedef int32 SlotNumber;

/*
 * Compare the tuples in the two given slots.
 */
static int32
heap_compare_slots(Datum a, Datum b, void *arg)
{
  GatherMergeState *node = (GatherMergeState *) arg;
  SlotNumber  slot1 = DatumGetInt32(a);
  SlotNumber  slot2 = DatumGetInt32(b);

  TupleTableSlot *s1 = node->gm_slots[slot1];
  TupleTableSlot *s2 = node->gm_slots[slot2];
  int     nkey;

  Assert(!TupIsNull(s1));
  Assert(!TupIsNull(s2));

  for (nkey = 0; nkey < node->gm_nkeys; nkey++)
  {
    SortSupport sortKey = node->gm_sortkeys + nkey;
    AttrNumber  attno = sortKey->ssup_attno;
    Datum   datum1,
          datum2;
    bool    isNull1,
          isNull2;
    int     compare;

    datum1 = slot_getattr(s1, attno, &isNull1);
    datum2 = slot_getattr(s2, attno, &isNull2);

    compare = ApplySortComparator(datum1, isNull1,
                    datum2, isNull2,
                    sortKey);
    if (compare != 0)
    {
      INVERT_COMPARE_RESULT(compare);
      return compare;
    }
  }
  return 0;
}

Coverage Report

Created: 2025-10-09 06:07