From 1cee9ec4d8487d277fd15f6827fbc610c3a2fccb Mon Sep 17 00:00:00 2001 From: Greg Burd Date: Fri, 23 Oct 2026 07:36:15 -0400 Subject: [PATCH v2 1/3] Use quickselect instead of a heap for bounded sorts A bounded sort (ORDER BY ... LIMIT N) switches to a heap of N tuples once it has 2 * N tuples in memory, and pays O(log N) comparisons for every tuple that gets into the heap. When N is large, or the input arrives in an order that admits most tuples, the heap is slow; at large N it can lose to an unbounded quicksort. Instead, keep a buffer of up to Max(2 * N, 1000) tuples. Whenever it fills, use quickselect to keep the N smallest and free the rest. That costs O(buffer) per (buffer - N) admitted tuples, so O(1) amortized per tuple. An incoming tuple is still rejected with one comparison, against the largest survivor of the last prune, as the heap does with its root. At the end the survivors are sorted with tuplesort_sort_memtuples(), which picks radix sort or the specialized qsort the same way an unbounded sort does. The 1000-tuple floor keeps small N from pruning every few tuples; it fits in the initial memtuples array, so the array never grows past what the sort already allocated. If the buffer fills work_mem before it fills up, prune early; if N tuples alone no longer fit, prune to N and finish with the existing heap code, which never holds more than N. Quickselect uses median-of-3 (ninther above 40 elements) pivots and Hoare partitioning, which splits runs of equal keys evenly. If it has not converged after 2 * log2(n) rounds it sorts the remaining range, so the worst case is O(n log n) per prune. --- src/backend/utils/sort/tuplesort.c | 203 ++++++++++++++++++++++++ src/test/regress/expected/tuplesort.out | 73 +++++++++ src/test/regress/sql/tuplesort.sql | 59 +++++++ 3 files changed, 335 insertions(+) diff --git a/src/backend/utils/sort/tuplesort.c b/src/backend/utils/sort/tuplesort.c index da3c41d7084..7650c8ef806 100644 --- a/src/backend/utils/sort/tuplesort.c +++ b/src/backend/utils/sort/tuplesort.c @@ -120,6 +120,15 @@ #define INITIAL_MEMTUPSIZE Max(1024, \ ALLOCSET_SEPARATE_THRESHOLD / sizeof(SortTuple) + 1) +/* + * A quickselect bounded sort buffers at least this many tuples between + * prunes, so that a small LIMIT does not prune every few tuples. It must fit + * in the initial memtuples array, which is never shrunk. + */ +#define TOPN_SELECT_MIN_BUFFER 1000 +StaticAssertDecl(TOPN_SELECT_MIN_BUFFER < INITIAL_MEMTUPSIZE, + "TOPN_SELECT_MIN_BUFFER must fit in the initial memtuples array"); + /* GUC variables */ bool trace_sort = false; @@ -155,6 +164,8 @@ typedef enum { TSS_INITIAL, /* Loading tuples; still within memory limit */ TSS_BOUNDED, /* Loading tuples into bounded-size heap */ + TSS_BOUNDED_SELECT, /* Loading tuples into a 2 * bound buffer that + * quickselect prunes back to bound */ TSS_BUILDRUNS, /* Loading tuples; writing to tape */ TSS_SORTEDINMEM, /* Sort completed entirely in memory */ TSS_SORTEDONTAPE, /* Sort completed, final run is on tape */ @@ -460,6 +471,7 @@ static bool mergereadnext(Tuplesortstate *state, LogicalTape *srcTape, SortTuple static void dumptuples(Tuplesortstate *state, bool alltuples); static void make_bounded_heap(Tuplesortstate *state); static void sort_bounded_heap(Tuplesortstate *state); +static void prune_bounded_select(Tuplesortstate *state); static void tuplesort_sort_memtuples(Tuplesortstate *state); static void tuplesort_heap_insert(Tuplesortstate *state, SortTuple *tuple); static void tuplesort_heap_replace_top(Tuplesortstate *state, SortTuple *tuple); @@ -1137,6 +1149,22 @@ tuplesort_puttuple_common(Tuplesortstate *state, SortTuple *tuple, (state->memtupcount > state->bound * 2 || (state->memtupcount > state->bound && LACKMEM(state)))) { + /* + * If 2 * bound tuples fit in memory, keep a buffer of them and + * prune it back to bound with quickselect whenever it fills. + */ + if (state->bound > 0 && !LACKMEM(state)) + { + if (trace_sort) + elog(LOG, "switching to bounded quickselect at %d tuples: %s", + state->memtupcount, + pg_rusage_show(&state->ru_start)); + prune_bounded_select(state); + state->status = TSS_BOUNDED_SELECT; + MemoryContextSwitchTo(oldcontext); + return; + } + if (trace_sort) elog(LOG, "switching to bounded heapsort at %d tuples: %s", state->memtupcount, @@ -1190,6 +1218,48 @@ tuplesort_puttuple_common(Tuplesortstate *state, SortTuple *tuple, } break; + case TSS_BOUNDED_SELECT: + + /* + * memtuples[bound - 1] is the largest of the bound smallest + * tuples seen so far, so nothing >= it can make the cut. + */ + if (COMPARETUP(state, tuple, &state->memtuples[state->bound - 1]) >= 0) + { + free_sort_tuple(state, tuple); + CHECK_FOR_INTERRUPTS(); + break; + } + /* the buffer never outgrows the array; see TOPN_SELECT_MIN_BUFFER */ + Assert(state->memtupcount < state->memtupsize); + state->memtuples[state->memtupcount++] = *tuple; + if (state->memtupcount > Max(state->bound * 2, TOPN_SELECT_MIN_BUFFER)) + prune_bounded_select(state); + else if (LACKMEM(state) && state->memtupcount > state->bound * 2) + { + /* + * The floor, not the bound, filled work_mem. More than bound + * tuples arrived since the last prune, so pruning early still + * costs O(1) per tuple. + */ + prune_bounded_select(state); + } + else if (LACKMEM(state)) + { + /* + * Tuples got wider than when we started. Keep just bound of + * them and finish with the heap, which never holds more. + */ + if (trace_sort) + elog(LOG, "switching from bounded quickselect to heapsort at %d tuples: %s", + state->memtupcount, + pg_rusage_show(&state->ru_start)); + prune_bounded_select(state); + state->status = TSS_INITIAL; + make_bounded_heap(state); + } + break; + case TSS_BUILDRUNS: /* @@ -1319,6 +1389,18 @@ tuplesort_performsort(Tuplesortstate *state) state->markpos_eof = false; break; + case TSS_BOUNDED_SELECT: + if (state->memtupcount > state->bound) + prune_bounded_select(state); + tuplesort_sort_memtuples(state); + state->status = TSS_SORTEDINMEM; + state->boundUsed = true; + state->current = 0; + state->eof_reached = false; + state->markpos_offset = 0; + state->markpos_eof = false; + break; + case TSS_BUILDRUNS: /* @@ -2562,6 +2644,127 @@ sort_bounded_heap(Tuplesortstate *state) state->boundUsed = true; } +/* Index of the median of memtuples[i], [j], [k] */ +static inline int +med3_index(Tuplesortstate *state, int i, int j, int k) +{ + SortTuple *a = state->memtuples; + + return COMPARETUP(state, &a[i], &a[j]) < 0 ? + (COMPARETUP(state, &a[j], &a[k]) < 0 ? j : + (COMPARETUP(state, &a[i], &a[k]) < 0 ? k : i)) : + (COMPARETUP(state, &a[j], &a[k]) > 0 ? j : + (COMPARETUP(state, &a[i], &a[k]) < 0 ? i : k)); +} + +/* + * Rearrange memtuples[0 .. n-1] so that memtuples[k] holds the tuple it would + * hold if the array were sorted, with nothing larger before it and nothing + * smaller after it. + * + * Hoare partitioning stops on keys equal to the pivot, so runs of equal keys + * still split evenly. If pivots keep coming out badly, sort what is left. + */ +static void +tuplesort_select(Tuplesortstate *state, int n, int k) +{ + SortTuple *a = state->memtuples; + int lo = 0; + int hi = n - 1; + int budget = 2 * (pg_leftmost_one_pos32((uint32) n) + 1); + + Assert(k >= 0 && k < n); + + while (hi > lo) + { + int len = hi - lo + 1; + int mid = lo + len / 2; + int p; + int i; + int j; + SortTuple tmp; + + CHECK_FOR_INTERRUPTS(); + + if (--budget < 0) + { + qsort_tuple(a + lo, len, state->base.comparetup, state); + return; + } + + /* median of three, or pseudo-median of nine for larger ranges */ + if (len > 40) + { + int d = len / 8; + + p = med3_index(state, + med3_index(state, lo, lo + d, lo + 2 * d), + med3_index(state, mid - d, mid, mid + d), + med3_index(state, hi - 2 * d, hi - d, hi)); + } + else + p = med3_index(state, lo, mid, hi); + + tmp = a[lo]; + a[lo] = a[p]; + a[p] = tmp; + + i = lo; + j = hi + 1; + for (;;) + { + while (COMPARETUP(state, &a[++i], &a[lo]) < 0) + if (i == hi) + break; + while (COMPARETUP(state, &a[lo], &a[--j]) < 0) + if (j == lo) + break; + if (i >= j) + break; + tmp = a[i]; + a[i] = a[j]; + a[j] = tmp; + } + tmp = a[lo]; + a[lo] = a[j]; + a[j] = tmp; + + /* a[lo .. j-1] <= a[j] <= a[j+1 .. hi] */ + if (j == k) + return; + if (j > k) + hi = j - 1; + else + lo = j + 1; + } +} + +/* + * Keep only the bound smallest memtuples, leaving the largest of them at + * memtuples[bound - 1]. + */ +static void +prune_bounded_select(Tuplesortstate *state) +{ + Assert(state->bounded); + Assert(state->memtupcount > state->bound); + Assert(SERIAL(state)); + + tuplesort_select(state, state->memtupcount, state->bound - 1); +#ifdef USE_ASSERT_CHECKING + for (int i = 0; i < state->memtupcount; i++) + { + int c = COMPARETUP(state, &state->memtuples[i], + &state->memtuples[state->bound - 1]); + + Assert(i < state->bound ? c <= 0 : c >= 0); + } +#endif + for (int i = state->bound; i < state->memtupcount; i++) + free_sort_tuple(state, &state->memtuples[i]); + state->memtupcount = state->bound; +} + /* radix sort routines */ diff --git a/src/test/regress/expected/tuplesort.out b/src/test/regress/expected/tuplesort.out index fc1321bf443..448eb73ee99 100644 --- a/src/test/regress/expected/tuplesort.out +++ b/src/test/regress/expected/tuplesort.out @@ -703,3 +703,76 @@ EXPLAIN (COSTS OFF) :qry; (10 rows) COMMIT; +---- +-- Bounded sorts. Once 2 * bound tuples are in memory, the sort keeps a +-- buffer of them (at least TOPN_SELECT_MIN_BUFFER, 1000) and prunes it to +-- bound with quickselect. If work_mem runs out first it falls back to a +-- heap. Compare each with a full sort, which is what a bound above +-- INT_MAX / 2 gets. +---- +CREATE TEMP TABLE bounded_sort AS + SELECT g AS id, (g * 7919) % 1000 AS k, md5(g::text) AS t + FROM generate_series(1, 5000) g; +-- true if ORDER BY ... LIMIT returns the same rows, in the same order, as a +-- full sort cut to the same length +CREATE FUNCTION bounded_sort_ok(ord text, lim int) RETURNS bool +LANGUAGE plpgsql AS $$ +DECLARE + bounded int[]; + full_sort int[]; +BEGIN + EXECUTE format('SELECT array(SELECT id FROM bounded_sort ORDER BY %s LIMIT %s)', + ord, lim) INTO bounded; + EXECUTE format('SELECT array(SELECT id FROM (SELECT id FROM bounded_sort ORDER BY %s LIMIT 1073741824) f LIMIT %s)', + ord, lim) INTO full_sort; + RETURN bounded = full_sort; +END $$; +SELECT o, l +FROM unnest(ARRAY['k, id', 'k DESC, id', 't', 't DESC', 'id DESC']) o, + unnest(ARRAY[1, 2, 7, 100, 499, 500, 501, 1250, 2499, 2500]) l +WHERE NOT bounded_sort_ok(o, l); + o | l +---+--- +(0 rows) + +-- duplicates: the bound falls inside a run of equal keys +SELECT count(*), count(DISTINCT k), min(k), max(k) +FROM (SELECT k FROM bounded_sort ORDER BY k LIMIT 2000) s; + count | count | min | max +-------+-------+-----+----- + 2000 | 400 | 0 | 399 +(1 row) + +-- rescan with a different bound each time +SELECT l, max(s.k), count(*) +FROM generate_series(10, 50, 20) l, + LATERAL (SELECT k FROM bounded_sort ORDER BY k, id LIMIT l) s +GROUP BY l ORDER BY l; + l | max | count +----+-----+------- + 10 | 1 | 10 + 30 | 5 | 30 + 50 | 9 | 50 +(3 rows) + +-- work_mem runs out after the buffer starts: narrow rows with large keys +-- come first, then wide rows that sort before them, so the survivors get +-- wider until bound of them no longer fit and the sort falls back to a heap +CREATE TEMP TABLE bounded_sort_wide AS + SELECT g AS id, + CASE WHEN g > 3000 THEN g - 3000 ELSE 100000 + g END AS k, + CASE WHEN g > 3000 THEN repeat(md5(g::text), 16) ELSE 'x' END AS pad + FROM generate_series(1, 6000) g; +BEGIN; +SET LOCAL work_mem = '256kB'; +SELECT array(SELECT id FROM (SELECT id, pad FROM bounded_sort_wide + ORDER BY k, id LIMIT 400) s) = + array(SELECT id FROM (SELECT id FROM bounded_sort_wide + ORDER BY k, id LIMIT 1073741824) f LIMIT 400) AS ok; + ok +---- + t +(1 row) + +COMMIT; +DROP FUNCTION bounded_sort_ok(text, int); diff --git a/src/test/regress/sql/tuplesort.sql b/src/test/regress/sql/tuplesort.sql index 8476e594e6c..fa2da7f2f4f 100644 --- a/src/test/regress/sql/tuplesort.sql +++ b/src/test/regress/sql/tuplesort.sql @@ -305,3 +305,62 @@ EXPLAIN (COSTS OFF) :qry; :qry; COMMIT; + +---- +-- Bounded sorts. Once 2 * bound tuples are in memory, the sort keeps a +-- buffer of them (at least TOPN_SELECT_MIN_BUFFER, 1000) and prunes it to +-- bound with quickselect. If work_mem runs out first it falls back to a +-- heap. Compare each with a full sort, which is what a bound above +-- INT_MAX / 2 gets. +---- +CREATE TEMP TABLE bounded_sort AS + SELECT g AS id, (g * 7919) % 1000 AS k, md5(g::text) AS t + FROM generate_series(1, 5000) g; + +-- true if ORDER BY ... LIMIT returns the same rows, in the same order, as a +-- full sort cut to the same length +CREATE FUNCTION bounded_sort_ok(ord text, lim int) RETURNS bool +LANGUAGE plpgsql AS $$ +DECLARE + bounded int[]; + full_sort int[]; +BEGIN + EXECUTE format('SELECT array(SELECT id FROM bounded_sort ORDER BY %s LIMIT %s)', + ord, lim) INTO bounded; + EXECUTE format('SELECT array(SELECT id FROM (SELECT id FROM bounded_sort ORDER BY %s LIMIT 1073741824) f LIMIT %s)', + ord, lim) INTO full_sort; + RETURN bounded = full_sort; +END $$; + +SELECT o, l +FROM unnest(ARRAY['k, id', 'k DESC, id', 't', 't DESC', 'id DESC']) o, + unnest(ARRAY[1, 2, 7, 100, 499, 500, 501, 1250, 2499, 2500]) l +WHERE NOT bounded_sort_ok(o, l); + +-- duplicates: the bound falls inside a run of equal keys +SELECT count(*), count(DISTINCT k), min(k), max(k) +FROM (SELECT k FROM bounded_sort ORDER BY k LIMIT 2000) s; + +-- rescan with a different bound each time +SELECT l, max(s.k), count(*) +FROM generate_series(10, 50, 20) l, + LATERAL (SELECT k FROM bounded_sort ORDER BY k, id LIMIT l) s +GROUP BY l ORDER BY l; + +-- work_mem runs out after the buffer starts: narrow rows with large keys +-- come first, then wide rows that sort before them, so the survivors get +-- wider until bound of them no longer fit and the sort falls back to a heap +CREATE TEMP TABLE bounded_sort_wide AS + SELECT g AS id, + CASE WHEN g > 3000 THEN g - 3000 ELSE 100000 + g END AS k, + CASE WHEN g > 3000 THEN repeat(md5(g::text), 16) ELSE 'x' END AS pad + FROM generate_series(1, 6000) g; +BEGIN; +SET LOCAL work_mem = '256kB'; +SELECT array(SELECT id FROM (SELECT id, pad FROM bounded_sort_wide + ORDER BY k, id LIMIT 400) s) = + array(SELECT id FROM (SELECT id FROM bounded_sort_wide + ORDER BY k, id LIMIT 1073741824) f LIMIT 400) AS ok; +COMMIT; + +DROP FUNCTION bounded_sort_ok(text, int); -- 2.54.0