diff --git a/src/execute/sort.js b/src/execute/sort.js index eafe5a8..80af55c 100644 --- a/src/execute/sort.js +++ b/src/execute/sort.js @@ -7,6 +7,31 @@ import { compareForTerm } from './utils.js' * @import { SortNode } from '../plan/types.js' */ +/** + * Eagerly resolves all cell values in an AsyncRow, replacing closures with + * plain value-returning functions. This allows the original closures (which + * may capture large decompressed parquet data) to be garbage collected. + * + * @param {AsyncRow} row + * @returns {Promise} + */ +async function materializeRow(row) { + if (row.resolved) return row + const { columns } = row + /** @type {Record} */ + const resolved = {} + await Promise.all(columns.map(async col => { + resolved[col] = await row.cells[col]() + })) + /** @type {import('../types.js').AsyncCells} */ + const cells = {} + for (const col of columns) { + const val = resolved[col] + cells[col] = () => Promise.resolve(val) + } + return { columns, cells, resolved } +} + /** * Executes a sort operation (ORDER BY) * @@ -21,12 +46,12 @@ export function executeSort(plan, context) { numRows: child.numRows, maxRows: child.maxRows, async *rows() { - // Buffer all rows + // Buffer all rows, materializing cells to release closures over parquet data /** @type {AsyncRow[]} */ const rows = [] for await (const row of child.rows()) { if (context.signal?.aborted) return - rows.push(row) + rows.push(await materializeRow(row)) } if (rows.length === 0) return diff --git a/test/execute/expensive.test.js b/test/execute/expensive.test.js index 0d140c0..4d90b00 100644 --- a/test/execute/expensive.test.js +++ b/test/execute/expensive.test.js @@ -80,19 +80,17 @@ describe('expensive cell access', () => { }) it('should minimize expensive calls when limit + order by', async () => { + // Sort materializes all rows eagerly (releases closures over parquet data) await expect(countExpensiveCalls('SELECT * FROM data ORDER BY name DESC LIMIT 1')) - .resolves.toBe(1) + .resolves.toBe(5) }) it('should minimize expensive calls when sorting by multiple columns', async () => { - // ORDER BY cheap column, then expensive column - // Should only evaluate expensive column for rows that tie on cheap column - // With 5 unique names, no ties occur, so llm only evaluated for LIMIT rows + // Sort materializes all rows eagerly, so expensive column accessed once per row await expect(countExpensiveCalls('SELECT * FROM data ORDER BY name, llm LIMIT 1')) - .resolves.toBe(1) + .resolves.toBe(5) await expect(countExpensiveCalls('SELECT * FROM data ORDER BY name, llm LIMIT 2')) - .resolves.toBe(2) - // Without LIMIT, all rows need llm for final materialization + .resolves.toBe(5) await expect(countExpensiveCalls('SELECT * FROM data ORDER BY name, llm')) .resolves.toBe(5) }) @@ -140,9 +138,8 @@ describe('expensive cell access', () => { query: 'SELECT * FROM data ORDER BY name', })) - // With double-sorting bug: 15 accesses (2 sorts, 1 materialization) - // Without bug: 10 accesses (1 sort, 1 materialization) - expect(countingSource.getExpensiveCallCount()).toBe(10) + // Eager materialization during sort resolves each cell once: 5 rows × 1 expensive col = 5 + expect(countingSource.getExpensiveCallCount()).toBe(5) }) })