From 8f5b6d6ef405c9d1a56fdae2c5f1bbc95ff617f6 Mon Sep 17 00:00:00 2001 From: Manuel Reyes Bravo Date: Sat, 3 Oct 2026 16:17:22 -0300 Subject: [PATCH v1] Add a per-node working-memory limit to Plan nodes Each Plan node gets a "workmem" field. When it is zero, which is what the planner leaves in it, the executor uses work_mem as before; when it is set, the executor enforces that value for the node's sorts, hash tables and tuplestores instead. Hash-based nodes multiply it by hash_mem_multiplier, as they do with work_mem. Nothing in core sets the field, so behavior does not change. It lets a planner hook give individual nodes their own limits, for example to split a per-query memory budget among the nodes of a plan, which is the first step Jeff Davis suggested in the query_work_mem thread. Code outside the executor that runs on behalf of a node follows the node's limit too: the limit is also kept in the node's ExprContext, so set-returning functions in materialize mode (InitMaterializedSRF() and the jsonb record-set functions) use it for their tuplestores, and ordered-set aggregates get the Agg node's limit through the new AggGetWorkMem(). EXPLAIN VERBOSE shows the limit of a node that has one. The test_plan_workmem module sets the field from a planner hook and checks, for each kind of node that uses working memory, that the node follows it both below and above work_mem. ExecChooseHashTableSize() and hash_agg_set_limits() now take the memory limit as an argument, since the planner calls them with the work_mem setting and the executor with the node's limit. The new CreateWorkExprContextExtended() sizes a work context for a given limit; CreateWorkExprContext() keeps its signature and uses work_mem. Discussion: https://postgr.es/m/CAJVSvF6s1LgXF6KB2Cz68sHzk+v+O_vmwEkaon=H8O9VcOr-tQ@mail.gmail.com --- src/backend/commands/explain.c | 4 + src/backend/executor/execSRF.c | 8 +- src/backend/executor/execUtils.c | 39 ++- src/backend/executor/nodeAgg.c | 64 ++-- src/backend/executor/nodeBitmapIndexscan.c | 2 +- src/backend/executor/nodeBitmapOr.c | 2 +- src/backend/executor/nodeCtescan.c | 3 +- src/backend/executor/nodeHash.c | 26 +- src/backend/executor/nodeIncrementalSort.c | 4 +- src/backend/executor/nodeMaterial.c | 3 +- src/backend/executor/nodeMemoize.c | 2 +- src/backend/executor/nodeRecursiveunion.c | 6 +- src/backend/executor/nodeSort.c | 4 +- src/backend/executor/nodeTableFuncscan.c | 3 +- src/backend/executor/nodeWindowAgg.c | 3 +- src/backend/optimizer/path/costsize.c | 4 +- src/backend/utils/adt/jsonfuncs.c | 4 +- src/backend/utils/adt/orderedsetaggs.c | 4 +- src/backend/utils/fmgr/funcapi.c | 4 +- src/include/executor/executor.h | 4 + src/include/executor/hashjoin.h | 2 + src/include/executor/nodeAgg.h | 3 +- src/include/executor/nodeHash.h | 1 + src/include/fmgr.h | 1 + src/include/miscadmin.h | 1 + src/include/nodes/execnodes.h | 3 + src/include/nodes/plannodes.h | 2 + src/test/modules/Makefile | 1 + src/test/modules/meson.build | 1 + src/test/modules/test_plan_workmem/.gitignore | 3 + src/test/modules/test_plan_workmem/Makefile | 21 ++ .../expected/test_plan_workmem.out | 285 ++++++++++++++++++ .../modules/test_plan_workmem/meson.build | 28 ++ .../sql/test_plan_workmem.sql | 186 ++++++++++++ .../test_plan_workmem/test_plan_workmem.c | 117 +++++++ 35 files changed, 802 insertions(+), 46 deletions(-) create mode 100644 src/test/modules/test_plan_workmem/.gitignore create mode 100644 src/test/modules/test_plan_workmem/Makefile create mode 100644 src/test/modules/test_plan_workmem/expected/test_plan_workmem.out create mode 100644 src/test/modules/test_plan_workmem/meson.build create mode 100644 src/test/modules/test_plan_workmem/sql/test_plan_workmem.sql create mode 100644 src/test/modules/test_plan_workmem/test_plan_workmem.c diff --git a/src/backend/commands/explain.c b/src/backend/commands/explain.c index 96f2f0e2e74..af01b8fabb8 100644 --- a/src/backend/commands/explain.c +++ b/src/backend/commands/explain.c @@ -1946,6 +1946,10 @@ ExplainNode(PlanState *planstate, List *ancestors, if (es->verbose) show_plan_tlist(planstate, ancestors, es); + /* node's own working-memory limit, if one was set */ + if (es->verbose && plan->workmem > 0) + ExplainPropertyInteger("Work Mem", "kB", plan->workmem, es); + /* unique join */ switch (nodeTag(plan)) { diff --git a/src/backend/executor/execSRF.c b/src/backend/executor/execSRF.c index f839f3bbde7..56484e87c66 100644 --- a/src/backend/executor/execSRF.c +++ b/src/backend/executor/execSRF.c @@ -262,7 +262,9 @@ ExecMakeTableFunctionResult(SetExprState *setexpr, MemoryContext oldcontext = MemoryContextSwitchTo(econtext->ecxt_per_query_memory); - tupstore = tuplestore_begin_heap(randomAccess, false, work_mem); + tupstore = tuplestore_begin_heap(randomAccess, false, + econtext->ecxt_workmem > 0 ? + econtext->ecxt_workmem : work_mem); rsinfo.setResult = tupstore; if (!returnsTuple) { @@ -398,7 +400,9 @@ no_function_result: MemoryContext oldcontext = MemoryContextSwitchTo(econtext->ecxt_per_query_memory); - tupstore = tuplestore_begin_heap(randomAccess, false, work_mem); + tupstore = tuplestore_begin_heap(randomAccess, false, + econtext->ecxt_workmem > 0 ? + econtext->ecxt_workmem : work_mem); rsinfo.setResult = tupstore; MemoryContextSwitchTo(oldcontext); diff --git a/src/backend/executor/execUtils.c b/src/backend/executor/execUtils.c index c0c4276acbf..6bf84ae339a 100644 --- a/src/backend/executor/execUtils.c +++ b/src/backend/executor/execUtils.c @@ -315,6 +315,30 @@ CreateExprContext(EState *estate) } +/* + * ExecGetWorkMem + * + * Working memory, in kB, for the given plan node: the node's own limit if + * one was set (for example by a planner hook), else work_mem. + */ +int +ExecGetWorkMem(PlanState *ps) +{ + return ps->plan->workmem > 0 ? ps->plan->workmem : work_mem; +} + +/* + * ExecGetHashMemoryLimit + * + * Memory limit, in bytes, for a hash table belonging to the given plan + * node: the node's working memory times hash_mem_multiplier. + */ +size_t +ExecGetHashMemoryLimit(PlanState *ps) +{ + return compute_hash_memory_limit(ExecGetWorkMem(ps), hash_mem_multiplier); +} + /* ---------------- * CreateWorkExprContext * @@ -325,10 +349,22 @@ CreateExprContext(EState *estate) */ ExprContext * CreateWorkExprContext(EState *estate) +{ + return CreateWorkExprContextExtended(estate, work_mem); +} + +/* + * CreateWorkExprContextExtended + * + * Like CreateWorkExprContext(), for a node whose working memory, in kB, is + * given rather than taken from work_mem. + */ +ExprContext * +CreateWorkExprContextExtended(EState *estate, int workmem) { Size maxBlockSize; - maxBlockSize = pg_prevpower2_size_t(work_mem * (Size) 1024 / 16); + maxBlockSize = pg_prevpower2_size_t(workmem * (Size) 1024 / 16); /* But no bigger than ALLOCSET_DEFAULT_MAXSIZE */ maxBlockSize = Min(maxBlockSize, ALLOCSET_DEFAULT_MAXSIZE); @@ -490,6 +526,7 @@ void ExecAssignExprContext(EState *estate, PlanState *planstate) { planstate->ps_ExprContext = CreateExprContext(estate); + planstate->ps_ExprContext->ecxt_workmem = planstate->plan->workmem; } /* ---------------- diff --git a/src/backend/executor/nodeAgg.c b/src/backend/executor/nodeAgg.c index 29037cf3122..0d98470f170 100644 --- a/src/backend/executor/nodeAgg.c +++ b/src/backend/executor/nodeAgg.c @@ -413,6 +413,7 @@ static double hash_choose_num_buckets(double hashentrysize, static int hash_choose_num_partitions(double input_groups, double hashentrysize, int used_bits, + Size hash_mem_limit, int *log2_npartitions); static void initialize_hash_entry(AggState *aggstate, TupleHashTable hashtable, @@ -433,7 +434,8 @@ static HashAggBatch *hashagg_batch_new(LogicalTape *input_tape, int setno, int64 input_tuples, double input_card, int used_bits); static MinimalTuple hashagg_batch_read(HashAggBatch *batch, uint32 *hashp); -static void hashagg_spill_init(HashAggSpill *spill, LogicalTapeSet *tapeset, +static void hashagg_spill_init(HashAggSpill *spill, Size hash_mem_limit, + LogicalTapeSet *tapeset, int used_bits, double input_groups, double hashentrysize); static Size hashagg_spill_tuple(AggState *aggstate, HashAggSpill *spill, @@ -530,7 +532,7 @@ initialize_phase(AggState *aggstate, int newphase) sortnode->sortOperators, sortnode->collations, sortnode->nullsFirst, - work_mem, + ExecGetWorkMem(&aggstate->ss.ps), NULL, TUPLESORT_NONE); } @@ -608,7 +610,8 @@ initialize_aggregate(AggState *aggstate, AggStatePerTrans pertrans, pertrans->sortOperators[0], pertrans->sortCollations[0], pertrans->sortNullsFirst[0], - work_mem, NULL, TUPLESORT_NONE); + ExecGetWorkMem(&aggstate->ss.ps), NULL, + TUPLESORT_NONE); } else pertrans->sortstates[aggstate->current_set] = @@ -618,7 +621,8 @@ initialize_aggregate(AggState *aggstate, AggStatePerTrans pertrans, pertrans->sortOperators, pertrans->sortCollations, pertrans->sortNullsFirst, - work_mem, NULL, TUPLESORT_NONE); + ExecGetWorkMem(&aggstate->ss.ps), NULL, + TUPLESORT_NONE); } /* @@ -1804,12 +1808,11 @@ hashagg_recompile_expressions(AggState *aggstate, bool minslot, bool nullcheck) */ void hash_agg_set_limits(double hashentrysize, double input_groups, int used_bits, - Size *mem_limit, uint64 *ngroups_limit, - int *num_partitions) + Size hash_mem_limit, Size *mem_limit, + uint64 *ngroups_limit, int *num_partitions) { int npartitions; Size partition_mem; - Size hash_mem_limit = get_hash_memory_limit(); /* if not expected to spill, use all of hash_mem */ if (input_groups * hashentrysize <= hash_mem_limit) @@ -1828,7 +1831,7 @@ hash_agg_set_limits(double hashentrysize, double input_groups, int used_bits, */ npartitions = hash_choose_num_partitions(input_groups, hashentrysize, - used_bits, + used_bits, hash_mem_limit, NULL); if (num_partitions != NULL) *num_partitions = npartitions; @@ -1927,7 +1930,8 @@ hash_agg_enter_spill_mode(AggState *aggstate) AggStatePerHash perhash = &aggstate->perhash[setno]; HashAggSpill *spill = &aggstate->hash_spills[setno]; - hashagg_spill_init(spill, aggstate->hash_tapeset, 0, + hashagg_spill_init(spill, ExecGetHashMemoryLimit(&aggstate->ss.ps), + aggstate->hash_tapeset, 0, perhash->aggnode->numGroups, aggstate->hashentrysize); } @@ -2002,7 +2006,8 @@ hash_create_memory(AggState *aggstate) * The hashcontext's per-tuple memory will be used for byref transition * values and returned by AggCheckCallContext(). */ - aggstate->hashcontext = CreateWorkExprContext(aggstate->ss.ps.state); + aggstate->hashcontext = CreateWorkExprContextExtended(aggstate->ss.ps.state, + ExecGetWorkMem(&aggstate->ss.ps)); /* * The meta context will be used for the bucket array of @@ -2032,7 +2037,8 @@ hash_create_memory(AggState *aggstate) * Like CreateWorkExprContext(), use smaller sizings for smaller work_mem, * to avoid large jumps in memory usage. */ - maxBlockSize = pg_prevpower2_size_t(work_mem * (Size) 1024 / 16); + maxBlockSize = pg_prevpower2_size_t(ExecGetWorkMem(&aggstate->ss.ps) * + (Size) 1024 / 16); /* But no bigger than ALLOCSET_DEFAULT_MAXSIZE */ maxBlockSize = Min(maxBlockSize, ALLOCSET_DEFAULT_MAXSIZE); @@ -2082,9 +2088,9 @@ hash_choose_num_buckets(double hashentrysize, double ngroups, Size memory) */ static int hash_choose_num_partitions(double input_groups, double hashentrysize, - int used_bits, int *log2_npartitions) + int used_bits, Size hash_mem_limit, + int *log2_npartitions) { - Size hash_mem_limit = get_hash_memory_limit(); double partition_limit; double mem_wanted; double dpartitions; @@ -2218,7 +2224,8 @@ lookup_hash_entries(AggState *aggstate) TupleTableSlot *slot = aggstate->tmpcontext->ecxt_outertuple; if (spill->partitions == NULL) - hashagg_spill_init(spill, aggstate->hash_tapeset, 0, + hashagg_spill_init(spill, ExecGetHashMemoryLimit(&aggstate->ss.ps), + aggstate->hash_tapeset, 0, perhash->aggnode->numGroups, aggstate->hashentrysize); @@ -2694,7 +2701,9 @@ agg_refill_hash_table(AggState *aggstate) aggstate->hash_batches = list_delete_last(aggstate->hash_batches); hash_agg_set_limits(aggstate->hashentrysize, batch->input_card, - batch->used_bits, &aggstate->hash_mem_limit, + batch->used_bits, + ExecGetHashMemoryLimit(&aggstate->ss.ps), + &aggstate->hash_mem_limit, &aggstate->hash_ngroups_limit, NULL); /* @@ -2782,7 +2791,8 @@ agg_refill_hash_table(AggState *aggstate) * that we don't assign tapes that will never be used. */ spill_initialized = true; - hashagg_spill_init(&spill, tapeset, batch->used_bits, + hashagg_spill_init(&spill, ExecGetHashMemoryLimit(&aggstate->ss.ps), + tapeset, batch->used_bits, batch->input_card, aggstate->hashentrysize); } /* no memory for a new group, spill */ @@ -2981,14 +2991,16 @@ agg_retrieve_hash_table_in_memory(AggState *aggstate) * of partitions to create, and initializes them. */ static void -hashagg_spill_init(HashAggSpill *spill, LogicalTapeSet *tapeset, int used_bits, +hashagg_spill_init(HashAggSpill *spill, Size hash_mem_limit, + LogicalTapeSet *tapeset, int used_bits, double input_groups, double hashentrysize) { int npartitions; int partition_bits; npartitions = hash_choose_num_partitions(input_groups, hashentrysize, - used_bits, &partition_bits); + used_bits, hash_mem_limit, + &partition_bits); #ifdef USE_INJECTION_POINTS if (IS_INJECTION_POINT_ATTACHED("hash-aggregate-single-partition")) @@ -3707,6 +3719,7 @@ ExecInitAgg(Agg *node, EState *estate, int eflags) totalGroups += aggstate->perhash[k].aggnode->numGroups; hash_agg_set_limits(aggstate->hashentrysize, totalGroups, 0, + ExecGetHashMemoryLimit(&aggstate->ss.ps), &aggstate->hash_mem_limit, &aggstate->hash_ngroups_limit, &aggstate->hash_planned_partitions); @@ -4727,6 +4740,21 @@ AggStateIsShared(FunctionCallInfo fcinfo) return true; } +/* + * AggGetWorkMem - working memory (kB) for an aggregate support function + * + * Ordered-set aggregates sort their input in the transition function. They + * should use the working memory of the Agg node that calls them, which is + * work_mem unless the node has its own limit. + */ +int +AggGetWorkMem(FunctionCallInfo fcinfo) +{ + if (fcinfo->context && IsA(fcinfo->context, AggState)) + return ExecGetWorkMem(&((AggState *) fcinfo->context)->ss.ps); + return work_mem; +} + /* * AggRegisterCallback - register a cleanup callback for an aggregate * diff --git a/src/backend/executor/nodeBitmapIndexscan.c b/src/backend/executor/nodeBitmapIndexscan.c index 90b010f9b71..5563b41c623 100644 --- a/src/backend/executor/nodeBitmapIndexscan.c +++ b/src/backend/executor/nodeBitmapIndexscan.c @@ -93,7 +93,7 @@ MultiExecBitmapIndexScan(BitmapIndexScanState *node) else { /* XXX should we use less than work_mem for this? */ - tbm = tbm_create(work_mem * (Size) 1024, + tbm = tbm_create(ExecGetWorkMem(&node->ss.ps) * (Size) 1024, ((BitmapIndexScan *) node->ss.ps.plan)->isshared ? node->ss.ps.state->es_query_dsa : NULL); } diff --git a/src/backend/executor/nodeBitmapOr.c b/src/backend/executor/nodeBitmapOr.c index e84f97eebd4..b842982d979 100644 --- a/src/backend/executor/nodeBitmapOr.c +++ b/src/backend/executor/nodeBitmapOr.c @@ -146,7 +146,7 @@ MultiExecBitmapOr(BitmapOrState *node) if (result == NULL) /* first subplan */ { /* XXX should we use less than work_mem for this? */ - result = tbm_create(work_mem * (Size) 1024, + result = tbm_create(ExecGetWorkMem(&node->ps) * (Size) 1024, ((BitmapOr *) node->ps.plan)->isshared ? node->ps.state->es_query_dsa : NULL); } diff --git a/src/backend/executor/nodeCtescan.c b/src/backend/executor/nodeCtescan.c index d3a8551e801..f6c4e6c9e7b 100644 --- a/src/backend/executor/nodeCtescan.c +++ b/src/backend/executor/nodeCtescan.c @@ -233,7 +233,8 @@ ExecInitCteScan(CteScan *node, EState *estate, int eflags) /* I am the leader */ prmdata->value = PointerGetDatum(scanstate); scanstate->leader = scanstate; - scanstate->cte_table = tuplestore_begin_heap(true, false, work_mem); + scanstate->cte_table = tuplestore_begin_heap(true, false, + ExecGetWorkMem(&scanstate->ss.ps)); tuplestore_set_eflags(scanstate->cte_table, scanstate->eflags); scanstate->readptr = 0; } diff --git a/src/backend/executor/nodeHash.c b/src/backend/executor/nodeHash.c index 8825bb6fa23..5fa1f43fa2c 100644 --- a/src/backend/executor/nodeHash.c +++ b/src/backend/executor/nodeHash.c @@ -501,6 +501,7 @@ ExecHashTableCreate(HashState *state) state->parallel_state != NULL, state->parallel_state != NULL ? state->parallel_state->nparticipants - 1 : 0, + ExecGetHashMemoryLimit(&state->ps), &space_allowed, &nbuckets, &nbatch, &num_skew_mcvs); @@ -540,6 +541,8 @@ ExecHashTableCreate(HashState *state) hashtable->spaceUsed = 0; hashtable->spacePeak = 0; hashtable->spaceAllowed = space_allowed; + hashtable->workmem = ExecGetWorkMem(&state->ps); + hashtable->hash_mem_limit = ExecGetHashMemoryLimit(&state->ps); hashtable->spaceUsedSkew = 0; hashtable->spaceAllowedSkew = hashtable->spaceAllowed * SKEW_HASH_MEM_PERCENT / 100; @@ -683,6 +686,7 @@ void ExecChooseHashTableSize(double ntuples, int tupwidth, bool useskew, bool try_combined_hash_mem, int parallel_workers, + size_t hash_mem_limit, size_t *space_allowed, int *numbuckets, int *numbatches, @@ -712,9 +716,9 @@ ExecChooseHashTableSize(double ntuples, int tupwidth, bool useskew, inner_rel_bytes = ntuples * tupsize; /* - * Compute in-memory hashtable size limit from GUCs. + * In-memory hashtable size limit, as given by the caller. */ - hash_table_bytes = get_hash_memory_limit(); + hash_table_bytes = hash_mem_limit; /* * Parallel Hash tries to use the combined hash_mem of all workers to @@ -832,7 +836,7 @@ ExecChooseHashTableSize(double ntuples, int tupwidth, bool useskew, if (try_combined_hash_mem) { ExecChooseHashTableSize(ntuples, tupwidth, useskew, - false, parallel_workers, + false, parallel_workers, hash_mem_limit, space_allowed, numbuckets, numbatches, @@ -1266,7 +1270,7 @@ ExecParallelHashIncreaseNumBatches(HashJoinTable hashtable) * to switch from one large combined memory budget to the * regular hash_mem budget. */ - pstate->space_allowed = get_hash_memory_limit(); + pstate->space_allowed = hashtable->hash_mem_limit; /* * The combined hash_mem of all participants wasn't @@ -2807,7 +2811,7 @@ ExecHashBuildNullTupleStore(HashJoinTable hashtable) * consumption too much. */ oldcxt = MemoryContextSwitchTo(hashtable->hashCxt); - tstore = tuplestore_begin_heap(false, false, work_mem / 16); + tstore = tuplestore_begin_heap(false, false, hashtable->workmem / 16); MemoryContextSwitchTo(oldcxt); return tstore; } @@ -3678,11 +3682,21 @@ ExecParallelHashTuplePrealloc(HashJoinTable hashtable, int batchno, size_t size) */ size_t get_hash_memory_limit(void) +{ + return compute_hash_memory_limit(work_mem, hash_mem_multiplier); +} + +/* + * Like get_hash_memory_limit(), for the given work_mem (in kB) and + * multiplier rather than the current settings. + */ +size_t +compute_hash_memory_limit(int workmem, double multiplier) { double mem_limit; /* Do initial calculation in double arithmetic */ - mem_limit = (double) work_mem * hash_mem_multiplier * 1024.0; + mem_limit = (double) workmem * multiplier * 1024.0; /* Clamp in case it doesn't fit in size_t */ mem_limit = Min(mem_limit, (double) SIZE_MAX); diff --git a/src/backend/executor/nodeIncrementalSort.c b/src/backend/executor/nodeIncrementalSort.c index 613d7d1c349..d0d1b983dd8 100644 --- a/src/backend/executor/nodeIncrementalSort.c +++ b/src/backend/executor/nodeIncrementalSort.c @@ -310,7 +310,7 @@ switchToPresortedPrefixMode(PlanState *pstate) &(plannode->sort.sortOperators[nPresortedCols]), &(plannode->sort.collations[nPresortedCols]), &(plannode->sort.nullsFirst[nPresortedCols]), - work_mem, + ExecGetWorkMem(&node->ss.ps), NULL, node->bounded ? TUPLESORT_ALLOWBOUNDED : TUPLESORT_NONE); node->prefixsort_state = prefixsort_state; @@ -599,7 +599,7 @@ ExecIncrementalSort(PlanState *pstate) plannode->sort.sortOperators, plannode->sort.collations, plannode->sort.nullsFirst, - work_mem, + ExecGetWorkMem(&node->ss.ps), NULL, node->bounded ? TUPLESORT_ALLOWBOUNDED : diff --git a/src/backend/executor/nodeMaterial.c b/src/backend/executor/nodeMaterial.c index e5f387612bc..d3d7ac7f2aa 100644 --- a/src/backend/executor/nodeMaterial.c +++ b/src/backend/executor/nodeMaterial.c @@ -62,7 +62,8 @@ ExecMaterial(PlanState *pstate) */ if (tuplestorestate == NULL && node->eflags != 0) { - tuplestorestate = tuplestore_begin_heap(true, false, work_mem); + tuplestorestate = tuplestore_begin_heap(true, false, + ExecGetWorkMem(&node->ss.ps)); tuplestore_set_eflags(tuplestorestate, node->eflags); if (node->eflags & EXEC_FLAG_MARK) { diff --git a/src/backend/executor/nodeMemoize.c b/src/backend/executor/nodeMemoize.c index a51dab1803f..272e0dd122b 100644 --- a/src/backend/executor/nodeMemoize.c +++ b/src/backend/executor/nodeMemoize.c @@ -1074,7 +1074,7 @@ ExecInitMemoize(Memoize *node, EState *estate, int eflags) mstate->mem_used = 0; /* Limit the total memory consumed by the cache to this */ - mstate->mem_limit = get_hash_memory_limit(); + mstate->mem_limit = ExecGetHashMemoryLimit(&mstate->ss.ps); /* A memory context dedicated for the cache */ mstate->tableContext = AllocSetContextCreate(CurrentMemoryContext, diff --git a/src/backend/executor/nodeRecursiveunion.c b/src/backend/executor/nodeRecursiveunion.c index 7166397e59b..d6753dacbfe 100644 --- a/src/backend/executor/nodeRecursiveunion.c +++ b/src/backend/executor/nodeRecursiveunion.c @@ -202,8 +202,10 @@ ExecInitRecursiveUnion(RecursiveUnion *node, EState *estate, int eflags) /* initialize processing state */ rustate->recursing = false; rustate->intermediate_empty = true; - rustate->working_table = tuplestore_begin_heap(false, false, work_mem); - rustate->intermediate_table = tuplestore_begin_heap(false, false, work_mem); + rustate->working_table = tuplestore_begin_heap(false, false, + ExecGetWorkMem(&rustate->ps)); + rustate->intermediate_table = tuplestore_begin_heap(false, false, + ExecGetWorkMem(&rustate->ps)); /* * If hashing, we need a per-tuple memory context for comparisons, and a diff --git a/src/backend/executor/nodeSort.c b/src/backend/executor/nodeSort.c index 2398bdabf0d..a0529c50b85 100644 --- a/src/backend/executor/nodeSort.c +++ b/src/backend/executor/nodeSort.c @@ -98,7 +98,7 @@ ExecSort(PlanState *pstate) plannode->sortOperators[0], plannode->collations[0], plannode->nullsFirst[0], - work_mem, + ExecGetWorkMem(&node->ss.ps), NULL, tuplesortopts); else @@ -108,7 +108,7 @@ ExecSort(PlanState *pstate) plannode->sortOperators, plannode->collations, plannode->nullsFirst, - work_mem, + ExecGetWorkMem(&node->ss.ps), NULL, tuplesortopts); if (node->bounded) diff --git a/src/backend/executor/nodeTableFuncscan.c b/src/backend/executor/nodeTableFuncscan.c index 9394156f405..fe02afa98aa 100644 --- a/src/backend/executor/nodeTableFuncscan.c +++ b/src/backend/executor/nodeTableFuncscan.c @@ -277,7 +277,8 @@ tfuncFetchRows(TableFuncScanState *tstate, ExprContext *econtext) /* build tuplestore for the result */ oldcxt = MemoryContextSwitchTo(econtext->ecxt_per_query_memory); - tstate->tupstore = tuplestore_begin_heap(false, false, work_mem); + tstate->tupstore = tuplestore_begin_heap(false, false, + ExecGetWorkMem(&tstate->ss.ps)); /* * Each call to fetch a new set of rows - of which there may be very many diff --git a/src/backend/executor/nodeWindowAgg.c b/src/backend/executor/nodeWindowAgg.c index b86dcbba055..63a934ae388 100644 --- a/src/backend/executor/nodeWindowAgg.c +++ b/src/backend/executor/nodeWindowAgg.c @@ -1153,7 +1153,8 @@ prepare_tuplestore(WindowAggState *winstate) Assert(winstate->buffer == NULL); /* Create new tuplestore */ - winstate->buffer = tuplestore_begin_heap(false, false, work_mem); + winstate->buffer = tuplestore_begin_heap(false, false, + ExecGetWorkMem(&winstate->ss.ps)); /* * Set up read pointers for the tuplestore. The current pointer doesn't diff --git a/src/backend/optimizer/path/costsize.c b/src/backend/optimizer/path/costsize.c index 7bbddb8bee4..0843574979e 100644 --- a/src/backend/optimizer/path/costsize.c +++ b/src/backend/optimizer/path/costsize.c @@ -2939,7 +2939,8 @@ cost_agg(Path *path, PlannerInfo *root, hashentrysize = hash_agg_entry_size(list_length(root->aggtransinfos), input_width, aggcosts->transitionSpace); - hash_agg_set_limits(hashentrysize, numGroups, 0, &mem_limit, + hash_agg_set_limits(hashentrysize, numGroups, 0, + get_hash_memory_limit(), &mem_limit, &ngroups_limit, &num_partitions); nbatches = Max((numGroups * hashentrysize) / mem_limit, @@ -4398,6 +4399,7 @@ initial_cost_hashjoin(PlannerInfo *root, JoinCostWorkspace *workspace, true, /* useskew */ parallel_hash, /* try_combined_hash_mem */ outer_path->parallel_workers, + get_hash_memory_limit(), &space_allowed, &numbuckets, &numbatches, diff --git a/src/backend/utils/adt/jsonfuncs.c b/src/backend/utils/adt/jsonfuncs.c index adbafd0ecd3..e9278127781 100644 --- a/src/backend/utils/adt/jsonfuncs.c +++ b/src/backend/utils/adt/jsonfuncs.c @@ -4127,7 +4127,9 @@ populate_recordset_worker(FunctionCallInfo fcinfo, const char *funcname, old_cxt = MemoryContextSwitchTo(rsi->econtext->ecxt_per_query_memory); state->tuple_store = tuplestore_begin_heap(rsi->allowedModes & SFRM_Materialize_Random, - false, work_mem); + false, + rsi->econtext->ecxt_workmem > 0 ? + rsi->econtext->ecxt_workmem : work_mem); MemoryContextSwitchTo(old_cxt); state->function_name = funcname; diff --git a/src/backend/utils/adt/orderedsetaggs.c b/src/backend/utils/adt/orderedsetaggs.c index 657fc4969a1..e59a95600aa 100644 --- a/src/backend/utils/adt/orderedsetaggs.c +++ b/src/backend/utils/adt/orderedsetaggs.c @@ -298,7 +298,7 @@ ordered_set_startup(FunctionCallInfo fcinfo, bool use_tuples) qstate->sortOperators, qstate->sortCollations, qstate->sortNullsFirsts, - work_mem, + AggGetWorkMem(fcinfo), NULL, tuplesortopt); else @@ -306,7 +306,7 @@ ordered_set_startup(FunctionCallInfo fcinfo, bool use_tuples) qstate->sortOperator, qstate->sortCollation, qstate->sortNullsFirst, - work_mem, + AggGetWorkMem(fcinfo), NULL, tuplesortopt); diff --git a/src/backend/utils/fmgr/funcapi.c b/src/backend/utils/fmgr/funcapi.c index 5c56002ba0e..89e210b5993 100644 --- a/src/backend/utils/fmgr/funcapi.c +++ b/src/backend/utils/fmgr/funcapi.c @@ -115,7 +115,9 @@ InitMaterializedSRF(FunctionCallInfo fcinfo, uint32 flags) random_access = (rsinfo->allowedModes & SFRM_Materialize_Random) != 0; - tupstore = tuplestore_begin_heap(random_access, false, work_mem); + tupstore = tuplestore_begin_heap(random_access, false, + rsinfo->econtext->ecxt_workmem > 0 ? + rsinfo->econtext->ecxt_workmem : work_mem); rsinfo->returnMode = SFRM_Materialize; rsinfo->setResult = tupstore; rsinfo->setDesc = stored_tupdesc; diff --git a/src/include/executor/executor.h b/src/include/executor/executor.h index 152a1dfa568..124e2fbca31 100644 --- a/src/include/executor/executor.h +++ b/src/include/executor/executor.h @@ -652,6 +652,10 @@ extern EState *CreateExecutorState(void); extern void FreeExecutorState(EState *estate); extern ExprContext *CreateExprContext(EState *estate); extern ExprContext *CreateWorkExprContext(EState *estate); +extern ExprContext *CreateWorkExprContextExtended(EState *estate, int workmem); +extern int ExecGetWorkMem(PlanState *ps); +extern size_t ExecGetHashMemoryLimit(PlanState *ps); + extern ExprContext *CreateStandaloneExprContext(void); extern void FreeExprContext(ExprContext *econtext, bool isCommit); extern void ReScanExprContext(ExprContext *econtext); diff --git a/src/include/executor/hashjoin.h b/src/include/executor/hashjoin.h index 4d342174b9a..72dcfc23b2c 100644 --- a/src/include/executor/hashjoin.h +++ b/src/include/executor/hashjoin.h @@ -361,6 +361,8 @@ typedef struct HashJoinTableData Size spaceUsed; /* memory space currently used by tuples */ Size spaceAllowed; /* upper limit for space used */ + int workmem; /* the Hash node's work_mem, in kB */ + Size hash_mem_limit; /* the Hash node's hash memory limit */ Size spacePeak; /* peak space used */ Size spaceUsedSkew; /* skew hash table's current space usage */ Size spaceAllowedSkew; /* upper limit for skew hashtable */ diff --git a/src/include/executor/nodeAgg.h b/src/include/executor/nodeAgg.h index 1e1be9666ae..c04e5e1b003 100644 --- a/src/include/executor/nodeAgg.h +++ b/src/include/executor/nodeAgg.h @@ -329,7 +329,8 @@ extern void ExecReScanAgg(AggState *node); extern Size hash_agg_entry_size(int numTrans, Size tupleWidth, Size transitionSpace); extern void hash_agg_set_limits(double hashentrysize, double input_groups, - int used_bits, Size *mem_limit, + int used_bits, Size hash_mem_limit, + Size *mem_limit, uint64 *ngroups_limit, int *num_partitions); /* parallel instrumentation support */ diff --git a/src/include/executor/nodeHash.h b/src/include/executor/nodeHash.h index 9ff493b627a..648276eeba6 100644 --- a/src/include/executor/nodeHash.h +++ b/src/include/executor/nodeHash.h @@ -59,6 +59,7 @@ extern void ExecHashTableResetMatchFlags(HashJoinTable hashtable); extern void ExecChooseHashTableSize(double ntuples, int tupwidth, bool useskew, bool try_combined_hash_mem, int parallel_workers, + size_t hash_mem_limit, size_t *space_allowed, int *numbuckets, int *numbatches, diff --git a/src/include/fmgr.h b/src/include/fmgr.h index 04b7914095f..e0d6397d88d 100644 --- a/src/include/fmgr.h +++ b/src/include/fmgr.h @@ -823,6 +823,7 @@ extern int AggCheckCallContext(FunctionCallInfo fcinfo, extern Aggref *AggGetAggref(FunctionCallInfo fcinfo); extern MemoryContext AggGetTempMemoryContext(FunctionCallInfo fcinfo); extern bool AggStateIsShared(FunctionCallInfo fcinfo); +extern int AggGetWorkMem(FunctionCallInfo fcinfo); extern void AggRegisterCallback(FunctionCallInfo fcinfo, ExprContextCallbackFunction func, Datum arg); diff --git a/src/include/miscadmin.h b/src/include/miscadmin.h index 8d6aacc4d5a..2cc9cb96a46 100644 --- a/src/include/miscadmin.h +++ b/src/include/miscadmin.h @@ -556,5 +556,6 @@ extern void RestoreClientConnectionInfo(char *conninfo); /* in executor/nodeHash.c */ extern size_t get_hash_memory_limit(void); +extern size_t compute_hash_memory_limit(int workmem, double multiplier); #endif /* MISCADMIN_H */ diff --git a/src/include/nodes/execnodes.h b/src/include/nodes/execnodes.h index 91bb0bd2e13..74ddc673739 100644 --- a/src/include/nodes/execnodes.h +++ b/src/include/nodes/execnodes.h @@ -329,6 +329,9 @@ typedef struct ExprContext /* Functions to call back when ExprContext is shut down or rescanned */ ExprContext_CB *ecxt_callbacks; + + /* working memory (kB) of the plan node, or 0 to use work_mem */ + int ecxt_workmem; } ExprContext; /* diff --git a/src/include/nodes/plannodes.h b/src/include/nodes/plannodes.h index 09a1ec73180..e3aabb37cc0 100644 --- a/src/include/nodes/plannodes.h +++ b/src/include/nodes/plannodes.h @@ -231,6 +231,8 @@ typedef struct Plan */ /* unique across entire final plan tree */ int plan_node_id; + /* working memory (kB) for this node, or 0 to use work_mem */ + int workmem; /* target list to be computed at this node */ List *targetlist; /* implicitly-ANDed qual conditions */ diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile index 71a2e65ad70..43c790205ad 100644 --- a/src/test/modules/Makefile +++ b/src/test/modules/Makefile @@ -43,6 +43,7 @@ SUBDIRS = \ test_parser \ test_pg_dump \ test_plan_advice \ + test_plan_workmem \ test_predtest \ test_radixtree \ test_rbtree \ diff --git a/src/test/modules/meson.build b/src/test/modules/meson.build index 77e1a2810e5..c533448a49d 100644 --- a/src/test/modules/meson.build +++ b/src/test/modules/meson.build @@ -44,6 +44,7 @@ subdir('test_oat_hooks') subdir('test_parser') subdir('test_pg_dump') subdir('test_plan_advice') +subdir('test_plan_workmem') subdir('test_predtest') subdir('test_radixtree') subdir('test_rbtree') diff --git a/src/test/modules/test_plan_workmem/.gitignore b/src/test/modules/test_plan_workmem/.gitignore new file mode 100644 index 00000000000..0c53fe590c6 --- /dev/null +++ b/src/test/modules/test_plan_workmem/.gitignore @@ -0,0 +1,3 @@ +results/ +log/ +tmp_check/ diff --git a/src/test/modules/test_plan_workmem/Makefile b/src/test/modules/test_plan_workmem/Makefile new file mode 100644 index 00000000000..37c8988c102 --- /dev/null +++ b/src/test/modules/test_plan_workmem/Makefile @@ -0,0 +1,21 @@ +# src/test/modules/test_plan_workmem/Makefile + +PGFILEDESC = "test_plan_workmem - test the per-node working-memory limit" + +MODULE_big = test_plan_workmem +OBJS = \ + $(WIN32RES) \ + test_plan_workmem.o + +REGRESS = test_plan_workmem + +ifdef USE_PGXS +PG_CONFIG = pg_config +PGXS := $(shell $(PG_CONFIG) --pgxs) +include $(PGXS) +else +subdir = src/test/modules/test_plan_workmem +top_builddir = ../../../.. +include $(top_builddir)/src/Makefile.global +include $(top_srcdir)/contrib/contrib-global.mk +endif diff --git a/src/test/modules/test_plan_workmem/expected/test_plan_workmem.out b/src/test/modules/test_plan_workmem/expected/test_plan_workmem.out new file mode 100644 index 00000000000..f0f56a678ed --- /dev/null +++ b/src/test/modules/test_plan_workmem/expected/test_plan_workmem.out @@ -0,0 +1,285 @@ +-- +-- The executor enforces a plan node's own working-memory limit when one is +-- set, and work_mem otherwise. Each query below exercises one kind of node +-- that uses working memory, and wm_explain() reports, per node, whether it +-- stayed in memory, without the sizes. +-- +LOAD 'test_plan_workmem'; +CREATE TABLE wm_tab AS + SELECT g AS a, g % 100 AS b, g % 10 AS c, repeat('x', 100) AS pad + FROM generate_series(1, 20000) g; +CREATE INDEX wm_tab_a ON wm_tab (a); +CREATE INDEX wm_tab_b ON wm_tab (b); +CREATE INDEX wm_tab_c ON wm_tab (c); +ANALYZE wm_tab; +CREATE FUNCTION wm_explain(query text) RETURNS SETOF text +LANGUAGE plpgsql AS +$$ +DECLARE + ln text; + node text; +BEGIN + FOR ln IN EXECUTE + 'EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF, BUFFERS OFF) ' || query + LOOP + IF ln ~ '\(actual' THEN + node := substring(ln FROM '^\s*(?:->\s*)?([A-Za-z ]+?)(?: on | using | \()'); + ELSIF ln ~ 'Sort Method: ' THEN + RETURN NEXT node || ': ' || + CASE WHEN ln ~ 'external' THEN 'disk' ELSE 'memory' END; + ELSIF ln ~ 'Storage: ' THEN + RETURN NEXT node || ': ' || + CASE WHEN ln ~ 'Storage: Disk' THEN 'disk' ELSE 'memory' END; + ELSIF ln ~ 'Batches: ' THEN + RETURN NEXT node || ': ' || + CASE WHEN ln ~ 'Batches: 1( |$)' THEN 'one batch' ELSE 'several batches' END; + ELSIF ln ~ 'Evictions: ' THEN + RETURN NEXT node || ': ' || + CASE WHEN ln ~ 'Evictions: 0 ' THEN 'no evictions' ELSE 'evictions' END; + ELSIF ln ~ 'Heap Blocks: ' THEN + RETURN NEXT node || ': ' || + CASE WHEN ln ~ 'lossy' THEN 'lossy' ELSE 'exact' END; + END IF; + END LOOP; +END; +$$; +-- More pages, so that a bitmap of them does not fit in 64kB +CREATE TABLE wm_big AS + SELECT g AS a, g % 10 AS c, repeat('x', 100) AS pad + FROM generate_series(1, 150000) g; +CREATE INDEX wm_big_a ON wm_big (a); +CREATE INDEX wm_big_c ON wm_big (c); +ANALYZE wm_big; +-- Set-returning functions and ordered-set aggregates keep their data in a +-- tuplestore or tuplesort that EXPLAIN does not report on: look for a +-- temporary file while the query's rows are read. (EXECUTE, so that the +-- query is planned with the current node limit.) +CREATE FUNCTION wm_spills(label text, query text) RETURNS text +LANGUAGE plpgsql AS +$$ +DECLARE + spilled bool; +BEGIN + -- q must be used, or the planner can leave the query out entirely + EXECUTE 'SELECT (SELECT count(*) > 0 FROM pg_ls_tmpdir()) FROM (' + || query || ') q WHERE length(q::text) > 0 LIMIT 1' + INTO spilled; + RETURN label || ': ' || CASE WHEN spilled THEN 'disk' ELSE 'memory' END; +END; +$$; +CREATE FUNCTION wm_run_all() RETURNS SETOF text +LANGUAGE plpgsql AS +$$ +BEGIN + SET LOCAL max_parallel_workers_per_gather = 0; + + RETURN QUERY SELECT wm_explain('SELECT * FROM wm_tab ORDER BY pad, a'); + + SET LOCAL enable_sort = off; + RETURN QUERY SELECT wm_explain( + 'SELECT * FROM wm_tab WHERE c < 2 ORDER BY c, pad'); + SET LOCAL enable_indexscan = off; + SET LOCAL enable_bitmapscan = off; + RETURN QUERY SELECT wm_explain( + 'SELECT a, count(*) FROM wm_tab GROUP BY a'); + RESET enable_sort; + RESET enable_indexscan; + RESET enable_bitmapscan; + + SET LOCAL enable_mergejoin = off; + SET LOCAL enable_nestloop = off; + RETURN QUERY SELECT wm_explain( + 'SELECT count(*) FROM wm_tab t1 JOIN wm_tab t2 USING (a)'); + RESET enable_mergejoin; + RESET enable_nestloop; + + SET LOCAL enable_hashjoin = off; + SET LOCAL enable_mergejoin = off; + RETURN QUERY SELECT wm_explain( + 'SELECT count(*) FROM generate_series(1, 3) g, + (SELECT * FROM wm_tab OFFSET 0) t WHERE t.a > g'); + RESET enable_hashjoin; + RESET enable_mergejoin; + + RETURN QUERY SELECT wm_explain( + 'WITH c AS MATERIALIZED (SELECT * FROM wm_tab) + SELECT count(*) FROM c c1, c c2 WHERE c1.a = c2.a'); + RETURN QUERY SELECT wm_explain( + 'SELECT count(*) OVER () FROM wm_tab'); + RETURN QUERY SELECT wm_explain( + 'WITH RECURSIVE r AS (SELECT a, pad, 1 AS n FROM wm_tab + UNION ALL + SELECT a, pad, n + 1 FROM r WHERE n < 2) + SELECT count(*) FROM r'); + RETURN QUERY SELECT wm_explain( + 'SELECT count(*) FROM JSON_TABLE( + (SELECT jsonb_agg(pad) FROM wm_tab), ''$[*]'' + COLUMNS (v text PATH ''$''))'); + + SET LOCAL enable_hashjoin = off; + SET LOCAL enable_mergejoin = off; + SET LOCAL enable_material = off; + RETURN QUERY SELECT wm_explain( + 'SELECT count(*) FROM wm_tab t1 JOIN wm_tab t2 ON t2.b = t1.b + WHERE t1.a <= 2000'); + RESET enable_hashjoin; + RESET enable_mergejoin; + RESET enable_material; + + SET LOCAL enable_seqscan = off; + SET LOCAL enable_indexscan = off; + RETURN QUERY SELECT wm_explain( + 'SELECT count(*) FROM wm_big WHERE c < 5'); + RETURN QUERY SELECT wm_explain( + 'SELECT count(*) FROM wm_big WHERE c = 1 OR a < 50000'); + RESET enable_seqscan; + RESET enable_indexscan; + + RETURN NEXT wm_spills('Function Scan (value per call)', + 'SELECT * FROM generate_series(1, 100000)'); + RETURN NEXT wm_spills('Function Scan (materialize mode)', + 'SELECT * FROM jsonb_each((SELECT jsonb_object_agg(a, pad) FROM wm_tab))'); + RETURN NEXT wm_spills('Function Scan (jsonb_to_recordset)', + 'SELECT * FROM jsonb_to_recordset((SELECT jsonb_agg(to_jsonb(t)) FROM wm_tab t)) + AS x(a int, pad text)'); + RETURN NEXT wm_spills('ProjectSet', + 'SELECT jsonb_each((SELECT jsonb_object_agg(a, pad) FROM wm_tab))'); + RETURN NEXT wm_spills('Aggregate (ordered-set)', + 'SELECT percentile_disc(0.5) WITHIN GROUP (ORDER BY pad || a) FROM wm_tab'); +END; +$$; +-- EXPLAIN VERBOSE shows a node's own limit, and nothing when it has none +EXPLAIN (VERBOSE, COSTS OFF) SELECT * FROM wm_tab ORDER BY pad; + QUERY PLAN +--------------------------------- + Sort + Output: a, b, c, pad + Sort Key: wm_tab.pad + -> Seq Scan on public.wm_tab + Output: a, b, c, pad +(5 rows) + +SET test_plan_workmem.node_limit = '1MB'; +EXPLAIN (VERBOSE, COSTS OFF) SELECT * FROM wm_tab ORDER BY pad; + QUERY PLAN +--------------------------------- + Sort + Output: a, b, c, pad + Work Mem: 1024 kB + Sort Key: wm_tab.pad + -> Seq Scan on public.wm_tab + Output: a, b, c, pad + Work Mem: 1024 kB +(7 rows) + +-- including a node the planner adds last: the Materialize on top of the +-- plan of a scrollable cursor +EXPLAIN (VERBOSE, COSTS OFF) + DECLARE wm_cur SCROLL CURSOR FOR SELECT b, count(*) FROM wm_tab GROUP BY b; + QUERY PLAN +------------------------------------------------------------- + Materialize + Output: b, (count(*)) + Work Mem: 1024 kB + -> GroupAggregate + Output: b, count(*) + Work Mem: 1024 kB + Group Key: wm_tab.b + -> Index Only Scan using wm_tab_b on public.wm_tab + Output: b + Work Mem: 1024 kB +(10 rows) + +RESET test_plan_workmem.node_limit; +-- No node limit: the nodes use work_mem +SET work_mem = '64MB'; +SELECT wm_run_all(); + wm_run_all +-------------------------------------------- + Sort: memory + Incremental Sort: memory + Incremental Sort: memory + HashAggregate: one batch + Hash: one batch + Materialize: memory + CTE Scan: memory + Hash: one batch + CTE Scan: memory + WindowAgg: memory + Recursive Union: memory + CTE Scan on r: memory + Table Function Scan: memory + Memoize: no evictions + Bitmap Heap Scan: exact + Bitmap Heap Scan: exact + Function Scan (value per call): memory + Function Scan (materialize mode): memory + Function Scan (jsonb_to_recordset): memory + ProjectSet: memory + Aggregate (ordered-set): memory +(21 rows) + +-- A node limit below work_mem is enforced +SET test_plan_workmem.node_limit = '64kB'; +SELECT wm_run_all(); + wm_run_all +------------------------------------------ + Sort: disk + Incremental Sort: memory + Incremental Sort: disk + HashAggregate: several batches + Hash: several batches + Materialize: disk + CTE Scan: disk + Hash: several batches + CTE Scan: disk + WindowAgg: disk + Recursive Union: disk + CTE Scan on r: disk + Table Function Scan: disk + Memoize: evictions + Bitmap Heap Scan: lossy + Bitmap Heap Scan: lossy + Function Scan (value per call): disk + Function Scan (materialize mode): disk + Function Scan (jsonb_to_recordset): disk + ProjectSet: disk + Aggregate (ordered-set): disk +(21 rows) + +-- So is a node limit above work_mem +SET work_mem = '64kB'; +SET test_plan_workmem.node_limit = '64MB'; +SELECT wm_run_all(); + wm_run_all +-------------------------------------------- + Sort: memory + Incremental Sort: memory + Incremental Sort: memory + HashAggregate: one batch + Hash: one batch + Materialize: memory + CTE Scan: memory + Hash: one batch + CTE Scan: memory + WindowAgg: memory + Recursive Union: memory + CTE Scan on r: memory + Table Function Scan: memory + Memoize: no evictions + Bitmap Heap Scan: exact + Bitmap Heap Scan: exact + Function Scan (value per call): memory + Function Scan (materialize mode): memory + Function Scan (jsonb_to_recordset): memory + ProjectSet: memory + Aggregate (ordered-set): memory +(21 rows) + +RESET test_plan_workmem.node_limit; +RESET work_mem; +DROP FUNCTION wm_run_all(); +DROP FUNCTION wm_spills(text, text); +DROP FUNCTION wm_explain(text); +DROP TABLE wm_big; +DROP TABLE wm_tab; diff --git a/src/test/modules/test_plan_workmem/meson.build b/src/test/modules/test_plan_workmem/meson.build new file mode 100644 index 00000000000..87a2e4e1a46 --- /dev/null +++ b/src/test/modules/test_plan_workmem/meson.build @@ -0,0 +1,28 @@ +# Copyright (c) 2026, PostgreSQL Global Development Group + +test_plan_workmem_sources = files( + 'test_plan_workmem.c', +) + +if host_system == 'windows' + test_plan_workmem_sources += rc_lib_gen.process(win32ver_rc, extra_args: [ + '--NAME', 'test_plan_workmem', + '--FILEDESC', 'test_plan_workmem - test the per-node working-memory limit',]) +endif + +test_plan_workmem = shared_module('test_plan_workmem', + test_plan_workmem_sources, + kwargs: pg_test_mod_args, +) +test_install_libs += test_plan_workmem + +tests += { + 'name': 'test_plan_workmem', + 'sd': meson.current_source_dir(), + 'bd': meson.current_build_dir(), + 'regress': { + 'sql': [ + 'test_plan_workmem', + ], + }, +} diff --git a/src/test/modules/test_plan_workmem/sql/test_plan_workmem.sql b/src/test/modules/test_plan_workmem/sql/test_plan_workmem.sql new file mode 100644 index 00000000000..5dd22c0a892 --- /dev/null +++ b/src/test/modules/test_plan_workmem/sql/test_plan_workmem.sql @@ -0,0 +1,186 @@ +-- +-- The executor enforces a plan node's own working-memory limit when one is +-- set, and work_mem otherwise. Each query below exercises one kind of node +-- that uses working memory, and wm_explain() reports, per node, whether it +-- stayed in memory, without the sizes. +-- +LOAD 'test_plan_workmem'; + +CREATE TABLE wm_tab AS + SELECT g AS a, g % 100 AS b, g % 10 AS c, repeat('x', 100) AS pad + FROM generate_series(1, 20000) g; +CREATE INDEX wm_tab_a ON wm_tab (a); +CREATE INDEX wm_tab_b ON wm_tab (b); +CREATE INDEX wm_tab_c ON wm_tab (c); +ANALYZE wm_tab; + +CREATE FUNCTION wm_explain(query text) RETURNS SETOF text +LANGUAGE plpgsql AS +$$ +DECLARE + ln text; + node text; +BEGIN + FOR ln IN EXECUTE + 'EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF, BUFFERS OFF) ' || query + LOOP + IF ln ~ '\(actual' THEN + node := substring(ln FROM '^\s*(?:->\s*)?([A-Za-z ]+?)(?: on | using | \()'); + ELSIF ln ~ 'Sort Method: ' THEN + RETURN NEXT node || ': ' || + CASE WHEN ln ~ 'external' THEN 'disk' ELSE 'memory' END; + ELSIF ln ~ 'Storage: ' THEN + RETURN NEXT node || ': ' || + CASE WHEN ln ~ 'Storage: Disk' THEN 'disk' ELSE 'memory' END; + ELSIF ln ~ 'Batches: ' THEN + RETURN NEXT node || ': ' || + CASE WHEN ln ~ 'Batches: 1( |$)' THEN 'one batch' ELSE 'several batches' END; + ELSIF ln ~ 'Evictions: ' THEN + RETURN NEXT node || ': ' || + CASE WHEN ln ~ 'Evictions: 0 ' THEN 'no evictions' ELSE 'evictions' END; + ELSIF ln ~ 'Heap Blocks: ' THEN + RETURN NEXT node || ': ' || + CASE WHEN ln ~ 'lossy' THEN 'lossy' ELSE 'exact' END; + END IF; + END LOOP; +END; +$$; + +-- More pages, so that a bitmap of them does not fit in 64kB +CREATE TABLE wm_big AS + SELECT g AS a, g % 10 AS c, repeat('x', 100) AS pad + FROM generate_series(1, 150000) g; +CREATE INDEX wm_big_a ON wm_big (a); +CREATE INDEX wm_big_c ON wm_big (c); +ANALYZE wm_big; + +-- Set-returning functions and ordered-set aggregates keep their data in a +-- tuplestore or tuplesort that EXPLAIN does not report on: look for a +-- temporary file while the query's rows are read. (EXECUTE, so that the +-- query is planned with the current node limit.) +CREATE FUNCTION wm_spills(label text, query text) RETURNS text +LANGUAGE plpgsql AS +$$ +DECLARE + spilled bool; +BEGIN + -- q must be used, or the planner can leave the query out entirely + EXECUTE 'SELECT (SELECT count(*) > 0 FROM pg_ls_tmpdir()) FROM (' + || query || ') q WHERE length(q::text) > 0 LIMIT 1' + INTO spilled; + RETURN label || ': ' || CASE WHEN spilled THEN 'disk' ELSE 'memory' END; +END; +$$; + +CREATE FUNCTION wm_run_all() RETURNS SETOF text +LANGUAGE plpgsql AS +$$ +BEGIN + SET LOCAL max_parallel_workers_per_gather = 0; + + RETURN QUERY SELECT wm_explain('SELECT * FROM wm_tab ORDER BY pad, a'); + + SET LOCAL enable_sort = off; + RETURN QUERY SELECT wm_explain( + 'SELECT * FROM wm_tab WHERE c < 2 ORDER BY c, pad'); + SET LOCAL enable_indexscan = off; + SET LOCAL enable_bitmapscan = off; + RETURN QUERY SELECT wm_explain( + 'SELECT a, count(*) FROM wm_tab GROUP BY a'); + RESET enable_sort; + RESET enable_indexscan; + RESET enable_bitmapscan; + + SET LOCAL enable_mergejoin = off; + SET LOCAL enable_nestloop = off; + RETURN QUERY SELECT wm_explain( + 'SELECT count(*) FROM wm_tab t1 JOIN wm_tab t2 USING (a)'); + RESET enable_mergejoin; + RESET enable_nestloop; + + SET LOCAL enable_hashjoin = off; + SET LOCAL enable_mergejoin = off; + RETURN QUERY SELECT wm_explain( + 'SELECT count(*) FROM generate_series(1, 3) g, + (SELECT * FROM wm_tab OFFSET 0) t WHERE t.a > g'); + RESET enable_hashjoin; + RESET enable_mergejoin; + + RETURN QUERY SELECT wm_explain( + 'WITH c AS MATERIALIZED (SELECT * FROM wm_tab) + SELECT count(*) FROM c c1, c c2 WHERE c1.a = c2.a'); + RETURN QUERY SELECT wm_explain( + 'SELECT count(*) OVER () FROM wm_tab'); + RETURN QUERY SELECT wm_explain( + 'WITH RECURSIVE r AS (SELECT a, pad, 1 AS n FROM wm_tab + UNION ALL + SELECT a, pad, n + 1 FROM r WHERE n < 2) + SELECT count(*) FROM r'); + RETURN QUERY SELECT wm_explain( + 'SELECT count(*) FROM JSON_TABLE( + (SELECT jsonb_agg(pad) FROM wm_tab), ''$[*]'' + COLUMNS (v text PATH ''$''))'); + + SET LOCAL enable_hashjoin = off; + SET LOCAL enable_mergejoin = off; + SET LOCAL enable_material = off; + RETURN QUERY SELECT wm_explain( + 'SELECT count(*) FROM wm_tab t1 JOIN wm_tab t2 ON t2.b = t1.b + WHERE t1.a <= 2000'); + RESET enable_hashjoin; + RESET enable_mergejoin; + RESET enable_material; + + SET LOCAL enable_seqscan = off; + SET LOCAL enable_indexscan = off; + RETURN QUERY SELECT wm_explain( + 'SELECT count(*) FROM wm_big WHERE c < 5'); + RETURN QUERY SELECT wm_explain( + 'SELECT count(*) FROM wm_big WHERE c = 1 OR a < 50000'); + RESET enable_seqscan; + RESET enable_indexscan; + + RETURN NEXT wm_spills('Function Scan (value per call)', + 'SELECT * FROM generate_series(1, 100000)'); + RETURN NEXT wm_spills('Function Scan (materialize mode)', + 'SELECT * FROM jsonb_each((SELECT jsonb_object_agg(a, pad) FROM wm_tab))'); + RETURN NEXT wm_spills('Function Scan (jsonb_to_recordset)', + 'SELECT * FROM jsonb_to_recordset((SELECT jsonb_agg(to_jsonb(t)) FROM wm_tab t)) + AS x(a int, pad text)'); + RETURN NEXT wm_spills('ProjectSet', + 'SELECT jsonb_each((SELECT jsonb_object_agg(a, pad) FROM wm_tab))'); + RETURN NEXT wm_spills('Aggregate (ordered-set)', + 'SELECT percentile_disc(0.5) WITHIN GROUP (ORDER BY pad || a) FROM wm_tab'); +END; +$$; + +-- EXPLAIN VERBOSE shows a node's own limit, and nothing when it has none +EXPLAIN (VERBOSE, COSTS OFF) SELECT * FROM wm_tab ORDER BY pad; +SET test_plan_workmem.node_limit = '1MB'; +EXPLAIN (VERBOSE, COSTS OFF) SELECT * FROM wm_tab ORDER BY pad; +-- including a node the planner adds last: the Materialize on top of the +-- plan of a scrollable cursor +EXPLAIN (VERBOSE, COSTS OFF) + DECLARE wm_cur SCROLL CURSOR FOR SELECT b, count(*) FROM wm_tab GROUP BY b; +RESET test_plan_workmem.node_limit; + +-- No node limit: the nodes use work_mem +SET work_mem = '64MB'; +SELECT wm_run_all(); + +-- A node limit below work_mem is enforced +SET test_plan_workmem.node_limit = '64kB'; +SELECT wm_run_all(); + +-- So is a node limit above work_mem +SET work_mem = '64kB'; +SET test_plan_workmem.node_limit = '64MB'; +SELECT wm_run_all(); + +RESET test_plan_workmem.node_limit; +RESET work_mem; +DROP FUNCTION wm_run_all(); +DROP FUNCTION wm_spills(text, text); +DROP FUNCTION wm_explain(text); +DROP TABLE wm_big; +DROP TABLE wm_tab; diff --git a/src/test/modules/test_plan_workmem/test_plan_workmem.c b/src/test/modules/test_plan_workmem/test_plan_workmem.c new file mode 100644 index 00000000000..99222236107 --- /dev/null +++ b/src/test/modules/test_plan_workmem/test_plan_workmem.c @@ -0,0 +1,117 @@ +/*------------------------------------------------------------------------- + * + * test_plan_workmem.c + * Test module for the per-node working-memory limit (Plan.workmem). + * + * A planner hook sets the workmem field of every node of the plan to + * test_plan_workmem.node_limit, so that tests can check that the executor + * enforces the node's limit rather than work_mem. + * + * Copyright (c) 2026, PostgreSQL Global Development Group + * + * IDENTIFICATION + * src/test/modules/test_plan_workmem/test_plan_workmem.c + * + *------------------------------------------------------------------------- + */ +#include "postgres.h" + +#include + +#include "fmgr.h" +#include "nodes/plannodes.h" +#include "optimizer/planner.h" +#include "utils/guc.h" + +PG_MODULE_MAGIC; + +/* limit to set on every plan node, in kB; zero leaves the plan alone */ +static int node_limit = 0; + +static planner_hook_type prev_planner_hook = NULL; + +static void set_node_limit(Plan *plan); + +static void +set_node_limit_list(List *plans) +{ + ListCell *lc; + + foreach(lc, plans) + set_node_limit((Plan *) lfirst(lc)); +} + +static void +set_node_limit(Plan *plan) +{ + if (plan == NULL) + return; + + plan->workmem = node_limit; + + switch (nodeTag(plan)) + { + case T_Append: + set_node_limit_list(((Append *) plan)->appendplans); + break; + case T_MergeAppend: + set_node_limit_list(((MergeAppend *) plan)->mergeplans); + break; + case T_BitmapAnd: + set_node_limit_list(((BitmapAnd *) plan)->bitmapplans); + break; + case T_BitmapOr: + set_node_limit_list(((BitmapOr *) plan)->bitmapplans); + break; + case T_SubqueryScan: + set_node_limit(((SubqueryScan *) plan)->subplan); + break; + case T_CustomScan: + set_node_limit_list(((CustomScan *) plan)->custom_plans); + break; + default: + break; + } + set_node_limit(outerPlan(plan)); + set_node_limit(innerPlan(plan)); +} + +static PlannedStmt * +test_plan_workmem_planner(Query *parse, const char *query_string, + int cursorOptions, ParamListInfo boundParams, + ExplainState *es) +{ + PlannedStmt *result; + + if (prev_planner_hook) + result = prev_planner_hook(parse, query_string, cursorOptions, + boundParams, es); + else + result = standard_planner(parse, query_string, cursorOptions, + boundParams, es); + + if (node_limit > 0) + { + set_node_limit(result->planTree); + set_node_limit_list(result->subplans); + } + + return result; +} + +void +_PG_init(void) +{ + DefineCustomIntVariable("test_plan_workmem.node_limit", + "Working memory to set on every plan node.", + "Zero leaves the plan alone, so nodes use work_mem.", + &node_limit, + 0, 0, MAX_KILOBYTES, + PGC_USERSET, + GUC_UNIT_KB, + NULL, NULL, NULL); + MarkGUCPrefixReserved("test_plan_workmem"); + + prev_planner_hook = planner_hook; + planner_hook = test_plan_workmem_planner; +} -- 2.55.0