From c0387fd7845b60b1f4759f34593aa30db2b21444 Mon Sep 17 00:00:00 2001 From: Greg Burd Date: Fri, 23 Oct 2026 07:36:15 -0400 Subject: [PATCH v1 1/2] Use quickselect instead of a heap for large 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 most tuples get in, and the heap is slower than an unbounded quicksort. Instead, keep a buffer of up to 2 * N tuples. Whenever it fills, use quickselect to keep the N smallest and free the rest; that costs O(N) per 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 buffer never holds more tuples than were in memory when the heap would have been built, so peak memory is unchanged. If wider tuples later push us over work_mem, prune to N and finish with the heap as before. Bounds below 8 keep using the heap: there the buffer refills every few tuples and the heap is faster. 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 | 185 ++++++++++++++++++++++++ src/test/regress/expected/tuplesort.out | 51 +++++++ src/test/regress/sql/tuplesort.sql | 41 ++++++ 3 files changed, 277 insertions(+) diff --git a/src/backend/utils/sort/tuplesort.c b/src/backend/utils/sort/tuplesort.c index 37c40763ee0..80189302101 100644 --- a/src/backend/utils/sort/tuplesort.c +++ b/src/backend/utils/sort/tuplesort.c @@ -155,6 +155,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 +462,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 +1140,24 @@ 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 that many and prune + * back to bound with quickselect whenever the buffer fills. + * Below a bound of 8 the buffer refills too often to pay for + * itself, and the heap wins. + */ + if (state->bound >= 8 && !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 +1211,37 @@ 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; + } + state->memtuples[state->memtupcount++] = *tuple; + if (state->memtupcount > state->bound * 2) + 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 +1371,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 +2626,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..35925a11e30 100644 --- a/src/test/regress/expected/tuplesort.out +++ b/src/test/regress/expected/tuplesort.out @@ -703,3 +703,54 @@ EXPLAIN (COSTS OFF) :qry; (10 rows) COMMIT; +---- +-- Bounded sorts. Once 2 * bound tuples are in memory, bounds of at least +-- 8 use quickselect and smaller ones use 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, 7, 8, 9, 100, 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) + +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..997b12ab7c3 100644 --- a/src/test/regress/sql/tuplesort.sql +++ b/src/test/regress/sql/tuplesort.sql @@ -305,3 +305,44 @@ EXPLAIN (COSTS OFF) :qry; :qry; COMMIT; + +---- +-- Bounded sorts. Once 2 * bound tuples are in memory, bounds of at least +-- 8 use quickselect and smaller ones use 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, 7, 8, 9, 100, 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; + +DROP FUNCTION bounded_sort_ok(text, int); -- 2.54.0