GROUP BY key batching runs out of memory: batch size is judged from the previous batch's finished keys

[bgmcmullen]

Summary

A GROUP BY runs out of memory when the group key is a small value computed from a large string, for example the first 150 characters of a cleaned-up 90KB text column. It fails on v0.16.6 even with a 4GB heap, on a table whose real data fits comfortably in memory.

The cause is in the adaptive batching added by #60. The batch size is chosen from the size of the finished keys of the previous batch. That can be fooled in two separate ways:

1. A function shrinks a long string into a short key. The engine measures the short result and never sees the large input each row had to read and copy. Row order does not matter. This is most of this issue.
2. Short rows come before long rows. Nothing is shrunk, but the engine grows the batch to 4,000 on the short rows and has already committed to a 4,000-row batch of long rows before it can measure them. See the second section below.

Both come from the same root: the cost of work about to start is judged from the size of results already finished. The first looks at the wrong thing, the second looks at the wrong time.

First way: a small key hides a large input

To group rows, the engine computes each row's key in batches, and all rows in a batch are in flight at the same time. After each batch it measures the finished keys and sizes the next batch to stay near a 16MB budget: small keys mean a big batch (up to 4,000 rows), big keys mean a small batch.

Now take this key:

substr(regexp_replace(system_text, '\s+', ' '), 1, 150)

The finished key is 150 characters, so the engine decides the work is cheap and takes 4,000 rows at once. But to make those 150 characters, every row first builds a cleaned-up copy of the entire 90KB string, and only then trims it. So 4,000 rows each hold a full-size copy at the same moment, and the heap fills before any of them finish.

An analogy: a kitchen decides how many orders to cook at once by looking at the size of the plates going out. The plates are small, so it starts 4,000 orders. But every order needs a whole turkey roasted to carve one slice. The plates were never the problem. The counter space for 4,000 turkeys was.

This is why #60's own benchmark shape passes. With GROUP BY regexp_replace(text, ...) and no wrapper, the finished key is the 90KB string, so measuring the key happens to measure the right thing. The substr wrapper hides the size.

One thing makes it worse

The copy each row holds is much bigger than 90KB. V8's String.prototype.replace with a global regex returns a rope (a tree of match fragments), not a flat string. For a 90KB input with about 13,000 whitespace runs, that rope weighs about 1.2MB until something flattens it:

200 replace outputs, as returned:          238 MB  (1190 KB each)
200 replace outputs, after flattening:      17 MB  (  81 KB each)

So the real in-flight cost is roughly 4,000 rows x 1.2MB.

Reproduce

mkdir sq-repro && cd sq-repro && npm init -y && npm install squirreling

repro.mjs:

import { collect, executeSql } from 'squirreling'

// 20,000 rows that share 700 distinct 90KB strings, like a dictionary-encoded text column
const pool = Array.from({ length: 700 }, (_, i) => `prompt ${i} ` + 'lorem  ipsum\n\tdolor '.repeat(4500))
const messages = Array.from({ length: 20000 }, (_, i) => ({ id: i, system_text: pool[i % 700] }))

const t0 = performance.now()
const rows = await collect(executeSql({ tables: { messages }, query: process.argv[2] }))
console.log(`OK rows=${rows.length} ${Math.round(performance.now() - t0)}ms`)

Run any query below with node --max-old-space-size=512 repro.mjs "<query>". The whole table is 63MB of string data.

Example queries (v0.16.6, 512MB heap unless noted)

Fails. Small key, large input. Also fails with a 4GB heap.

SELECT substr(regexp_replace(system_text, '\s+', ' '), 1, 150) AS k, count(*) AS n
FROM messages
GROUP BY substr(regexp_replace(system_text, '\s+', ' '), 1, 150)

Fails. Same shape with a different function, so it is not specific to regex.

SELECT substr(lower(system_text), 1, 150) AS k, count(*) AS n
FROM messages
GROUP BY substr(lower(system_text), 1, 150)

Works (23s). No wrapper, so the key is the large string and #60 shrinks the batch.

SELECT regexp_replace(system_text, '\s+', ' ') AS k, count(*) AS n
FROM messages
GROUP BY regexp_replace(system_text, '\s+', ' ')

Works (0.2s). Rewrite: trim the string before the function, so nothing large is ever built. Not quite the same answer: two strings that differ only after a very long whitespace run inside the first 600 characters will merge.

SELECT substr(regexp_replace(substr(system_text, 1, 600), '\s+', ' '), 1, 150) AS k, count(*) AS n
FROM messages
GROUP BY substr(regexp_replace(substr(system_text, 1, 600), '\s+', ' '), 1, 150)

Works with a 1GB heap (1.1s), fails at 512MB. Rewrite: collapse duplicates first, then run the function once per distinct value. Exactly the same answer as the original. Only helps when values repeat. This is #61's idea written by hand.

SELECT substr(regexp_replace(system_text, '\s+', ' '), 1, 150) AS k, sum(n) AS n
FROM (SELECT system_text, count(*) AS n FROM messages GROUP BY system_text)
GROUP BY substr(regexp_replace(system_text, '\s+', ' '), 1, 150)

Note that the last one still needs 1GB to process 700 rows of 90KB text. That is the same bug showing through: all 700 rows are in one batch, each holding a 1.2MB rope.

Evidence that the batch size is the cause

Changing only MAX_CHUNK_ROWS in src/execute/fold.js, with the failing query not rewritten:

MAX_CHUNK_ROWS	Failing query above, 512MB	Same query, 6,000 all-distinct 90KB values, 1.5GB	Ordinary GROUP BY upper(name), 200k rows of 40 chars
4000 (shipped)	OOM	OOM	0.19s
64	OK, 28s	OK	0.20s
8	OK, 5.3s	OK	0.52s

The all-distinct column matters. #61 concluded that case needs a cardinality hint from the data source. It does not: with the batch bounded it passes, because the finished keys are tiny and only the in-flight copies were ever the problem.

A fixed smaller cap is not the right fix though. 8 slows every ordinary query, and 64 only barely survives and would fail again on wider strings.

Second way: short rows first

Here the key is not shortened. This is the exact shape #60 fixed and benchmarked:

SELECT regexp_replace(system_text, '\s+', ' ') AS k, count(*) AS n
FROM messages
GROUP BY regexp_replace(system_text, '\s+', ' ')

It works when the large rows come first, because the first 64-row batch returns 90KB keys and the next batch shrinks to about 90 rows. It fails when small rows come first. The first batches return tiny keys, which really are tiny, so the batch grows to 4,000. When the large rows begin, 4,000 of them are dispatched in one batch. The engine would shrink the batch after measuring those results, but the heap is already gone.

skew.mjs, same table as above but half the rows are short:

import { collect, executeSql } from 'squirreling'

const pool = Array.from({ length: 700 }, (_, i) => `prompt ${i} ` + 'lorem  ipsum\n\tdolor '.repeat(4500))
const short = Array.from({ length: 10000 }, (_, i) => ({ id: i, system_text: `short ${i % 50}` }))
const fat = Array.from({ length: 10000 }, (_, i) => ({ id: 10000 + i, system_text: pool[i % 700] }))
const messages = process.env.ORDER === 'fat-first' ? [...fat, ...short] : [...short, ...fat]

const t0 = performance.now()
const rows = await collect(executeSql({ tables: { messages }, query: process.argv[2] }))
console.log(`OK rows=${rows.length} ${Math.round(performance.now() - t0)}ms`)

Same 20,000 rows, same query, 512MB heap, only the order differs:

Row order	v0.16.6	Prototype patch below
ORDER=fat-first	OK, 12s	OK, 2.6s
short rows first (default)	OOM	OOM

This layout is common in real data. In the table behind the production query, about 90% of rows have an empty system_text, so a long run of short rows followed by a cluster of wide ones is the normal case.

Note the prototype patch does not fix this one. It also resizes after the fact, so it has the same blind spot.

Proposed fix

Size the batch from what goes in as well as what comes out, and use whichever is larger. Expressions read their inputs only through row.cells, so foldEvaluatedRows can wrap each row's cells, tally the bytes read during the batch, and feed max(bytes read, result bytes) into the existing sizing. The 4,000 cap stays, so small-input queries keep full overlap.

Separately, flatten the regexp_replace result before returning it. Reading one character is enough to make V8 flatten a rope in place.

I prototyped both against master (v0.16.6). About 22 lines of code in two files. All 2,032 tests pass, tsc and eslint clean.

Case	master	patched
Failing query, repeated values, 512MB	OOM	OK, 12s
Failing query, 6,000 all-distinct values, 1.5GB	OOM	OK, 1.6s
substr(lower(...)) variant, 512MB	OOM	OK, 0.2s
Unwrapped GROUP BY regexp_replace(...), 512MB	OK, 24s	OK, 14s
Ordinary small-key GROUP BY, 200k rows	0.19s	0.245s

Each change alone: input-byte sizing alone fixes the failing query (21s). Flattening alone does not fix it. Together, 12s.

Known costs and gaps in the prototype:

• It does not fix the second way (short rows first). Any scheme that only resizes after a batch finishes will overshoot when the data changes character. A complete fix also needs a limit during dispatch: keep a running total of input bytes in flight and stop starting new rows once it passes the budget. Reading the input cell is the first thing each row's evaluation does, before any expensive function runs, so for already-loaded data the size is known early enough to act on. I have not built this part.
• Ordinary grouped queries get about 28% slower, from wrapping every row in a Proxy. Metering a sample of rows per batch instead of all of them should recover most of that.
• Object inputs still count as 16 bytes in valueBytes, so CAST(json_column AS VARCHAR) behind a small key is still invisible to the sizing.
• Sorting is out of scope, as it was in Evaluate group and partition keys in byte-bounded adaptive chunks #60. ORDER BY upper(system_text) still retains every row's full key.

Prototype patch

diff --git a/src/execute/fold.js b/src/execute/fold.js
index e675d65..73d005c 100644
--- a/src/execute/fold.js
+++ b/src/execute/fold.js
@@ -31,6 +31,29 @@ function valueBytes(value) {
   return 16
 }
 
+/**
+ * Wraps a row so the bytes of every cell read through it are added to a
+ * shared tally. Expressions read inputs only through row.cells, so the tally
+ * is the input side of whatever the evaluation holds in flight.
+ *
+ * @param {AsyncRow} row
+ * @param {{ bytes: number }} read
+ * @returns {AsyncRow}
+ */
+function meteredRow(row, read) {
+  const cells = new Proxy(row.cells, {
+    get(target, name) {
+      const cell = Reflect.get(target, name)
+      if (typeof cell !== 'function') return cell
+      return () => cell().then((/** @type {unknown} */ value) => {
+        read.bytes += valueBytes(value)
+        return value
+      })
+    },
+  })
+  return { ...row, cells }
+}
+
 /**
  * Evaluates a value for every row and folds each result in row order, holding
  * at most one adaptively sized chunk of evaluated values. Rows in a chunk are
@@ -58,15 +81,22 @@ export async function foldEvaluatedRows({ rows, evaluate, fold, signal }) {
     }
     const end = Math.min(start + chunkSize, rows.length)
     pending.length = end - start
+    const read = { bytes: 0 }
     for (let i = start; i < end; i++) {
-      pending[i - start] = evaluate(rows[i], i)
+      pending[i - start] = evaluate(meteredRow(rows[i], read), i)
     }
     const values = await Promise.all(pending)
-    let bytes = 0
+    // Size from the larger of what the chunk read and what it produced: a
+    // small result can hide a fat intermediate, as in
+    // substr(regexp_replace(text, ...), 1, 150), where every in-flight row
+    // holds a full-length copy of its input that the 150-char key never shows.
+    let { bytes } = read
+    let resultBytes = 0
     for (let j = 0; j < values.length; j++) {
-      bytes += valueBytes(values[j])
+      resultBytes += valueBytes(values[j])
       fold(values[j], start + j)
     }
+    if (resultBytes > bytes) bytes = resultBytes
     start = end
     const bytesPerRow = Math.max(1, bytes / values.length)
     chunkSize = Math.min(MAX_CHUNK_ROWS, Math.max(1, Math.floor(CHUNK_BYTE_BUDGET / bytesPerRow)))
diff --git a/src/expression/regexp.js b/src/expression/regexp.js
index 57de0d6..504c49b 100644
--- a/src/expression/regexp.js
+++ b/src/expression/regexp.js
@@ -137,8 +137,12 @@ export function evaluateRegexpFunc({ funcName, node, args, rowIndex }) {
     const searchStr = strVal.substring(position - 1)
 
     if (occurrence === 0) {
-      // Replace all occurrences
-      return prefix + searchStr.replace(regex, replacementStr)
+      // Replace all occurrences. V8 returns a rope of match fragments that can
+      // weigh over 10x the flat string (about 1.2MB for a 90KB input with 13k
+      // matches); reading the last char flattens it in place.
+      const replaced = prefix + searchStr.replace(regex, replacementStr)
+      replaced.charCodeAt(replaced.length - 1)
+      return replaced
     }
 
     // Replace only the nth occurrence

Related

• Evaluate group and partition keys in byte-bounded adaptive chunks #60 added the adaptive batching this issue is about. It fixes the unwrapped shape.
• Memoize regexp functions over repeated large inputs #57 / Memoize regexp functions over repeated large inputs #58 (result memoization) and Dedup-first eager aggregation for GROUP BY over string functions of one column #61 (dedup-first plan rewrite) targeted the same production query and were closed unmerged. Dedup-first eager aggregation for GROUP BY over string functions of one column #61 reduces how many rows reach the oversized batch but leaves the sizing as is, so it cannot help when values do not repeat.
• hyparam/hypaware-server#395 is the caller-side guard that refuses these query shapes today. It can narrow once this is fixed.