Add scanColumn streaming primitive to the iceberg SQL data source

src/sql/icebergDataSource.js

@@ -7,7 +7,7 @@ import { fileMightMatch, partitionMightMatch } from '../prune.js'
 import { whereToParquetFilter } from './whereFilter.js'
 
 /**
- * @import {AsyncDataSource} from 'squirreling'
+ * @import {AsyncDataSource, SqlPrimitive} from 'squirreling'
  * @import {Lister, Resolver, TableMetadata} from '../../src/types.js'
  */
 
@@ -192,5 +192,111 @@ export async function icebergDataSource({ tableUrl, metadataFileName, metadata,
         },
       }
     },
+    /**
+     * Streams a single column's values in row order as an async iterable of
+     * chunks (one chunk per parquet row group), so peak memory is bounded by a
+     * single row group's worth of values rather than the whole table. This is
+     * squirreling's optional `AsyncDataSource.scanColumn` hook: its
+     * `tryColumnScanAggregate` consumes it to compute a scalar aggregate
+     * (`COUNT`/`MIN`/`MAX`/`SUM`/`AVG`, low-cardinality `COUNT(DISTINCT …)`) in
+     * O(1)/O(cardinality) state; without the hook the engine falls back to
+     * buffering every scanned row.
+     *
+     * `scanColumn` is never given a WHERE (the engine only takes the streaming
+     * aggregate path when the scan has no filter), so there is no file or
+     * row-group pruning here, and there is no `appliedLimitOffset` flag to defer
+     * to: the source must fully honor LIMIT/OFFSET itself. Without deletes,
+     * record_count is exact, so whole files are skipped for OFFSET and the
+     * per-file read is bounded for LIMIT. With deletes, record_count is
+     * pre-delete, so OFFSET/LIMIT are applied over the post-delete value stream
+     * instead. `signal` aborts between chunks, mirroring `scan`.
+     *
+     * @param {object} options
+     * @param {string} options.column - Name of the single column to stream.
+     * @param {number} [options.limit] - Max number of values to yield.
+     * @param {number} [options.offset] - Number of leading values to skip.
+     * @param {AbortSignal} [options.signal] - Aborts the stream between chunks.
+     * @returns {AsyncIterable<ArrayLike<SqlPrimitive>>} Chunks of column values.
+     */
+    scanColumn({ column, limit, offset, signal }) {
+      const wantedColumns = [column]
+      const skip = offset ?? 0
+      const take = limit ?? Infinity
+      return {
+        async *[Symbol.asyncIterator]() {
+          if (signal?.aborted) throw new DOMException('Aborted', 'AbortError')
+          if (take === 0 || dataEntries.length === 0) return
+
+          const { positionDeletesMap, equalityDeleteGroups } = await deleteMapsPromise
+
+          let remainingSkip = skip
+          let remaining = take
+          for (const entry of dataEntries) {
+            if (remaining <= 0) break
+            if (signal?.aborted) throw new DOMException('Aborted', 'AbortError')
+            const recordCount = Number(entry.data_file.record_count)
+            // Without deletes record_count counts visible rows, so skip whole
+            // files for OFFSET and bound the per-file end for LIMIT. With
+            // deletes it is pre-delete, so read each file whole and slice the
+            // post-delete batches below.
+            let fileRowStart = 0
+            if (!hasDeletes && remainingSkip > 0) {
+              if (remainingSkip >= recordCount) {
+                remainingSkip -= recordCount
+                continue
+              }
+              fileRowStart = remainingSkip
+              remainingSkip = 0
+            }
+            const fileRowEnd = !hasDeletes && remaining !== Infinity
+              ? Math.min(recordCount, fileRowStart + remaining)
+              : recordCount
+
+            let stop = false
+            for await (const batch of readDataFile({
+              dataEntry: entry,
+              fileRowStart,
+              fileRowEnd,
+              schema,
+              metadata: tableMetadata,
+              resolver: fetchResolver,
+              rowLineage,
+              positionDeletesMap,
+              equalityDeleteGroups,
+              wantedColumns,
+              signal,
+            })) {
+              if (signal?.aborted) throw new DOMException('Aborted', 'AbortError')
+              // Apply any OFFSET still owed in post-delete coordinates. Without
+              // deletes this is already spent at fileRowStart, so the guard is
+              // only reached when deletes are present.
+              let start = 0
+              if (remainingSkip > 0) {
+                if (remainingSkip >= batch.length) {
+                  remainingSkip -= batch.length
+                  continue
+                }
+                start = remainingSkip
+                remainingSkip = 0
+              }
+              let end = batch.length
+              if (remaining !== Infinity && end - start > remaining) {
+                end = start + remaining
+                stop = true
+              }
+              /** @type {SqlPrimitive[]} */
+              const chunk = []
+              for (let i = start; i < end; i++) chunk.push(batch[i][column])
+              if (chunk.length > 0) {
+                yield chunk
+                remaining -= chunk.length
+              }
+              if (stop || remaining <= 0) break
+            }
+            if (stop || remaining <= 0) break
+          }
+        },
+      }
+    },
   }
 }

test/sql/icebergDataSource.test.js

@@ -4,7 +4,7 @@ import { icebergDataSource } from '../../src/sql/icebergDataSource.js'
 import { localResolver } from '../helpers.js'
 
 /**
- * @import {ExprNode} from 'squirreling'
+ * @import {AsyncDataSource, ExprNode} from 'squirreling'
  * @import {Resolver} from '../../src/types.js'
  */
 
@@ -468,3 +468,256 @@ describe.concurrent('icebergDataSource partition pruning', () => {
     expect(resolver.dataFilesRead()).toBe(2)
   })
 })
+
+describe.concurrent('icebergDataSource scanColumn', () => {
+  const tableUrl = 's3://hyperparam-iceberg/java/bunnies'
+  const resolver = localResolver('test/files')
+
+  // bunnies v2/v4 each resolve to a single data file, so the cross-file
+  // OFFSET/LIMIT and between-file abort branches need a multi-file table.
+  // spark/rename_column has three data files walked in record-count order
+  // (two single-row files then the two-row file; oracle id order [3,4,1,2]).
+  const renameTableUrl = 's3://hyperparam-iceberg/spark/rename_column'
+
+  /**
+   * Flatten a scanColumn stream into a single array of values, asserting each
+   * yielded chunk is an array-like batch (the streaming/bounded-memory shape).
+   *
+   * @param {AsyncIterable<ArrayLike<any>>} stream
+   * @returns {Promise<{ values: any[], chunks: number }>}
+   */
+  async function drain(stream) {
+    /** @type {any[]} */
+    const values = []
+    let chunks = 0
+    for await (const chunk of stream) {
+      expect(typeof chunk.length).toBe('number')
+      chunks++
+      for (let i = 0; i < chunk.length; i++) values.push(chunk[i])
+    }
+    return { values, chunks }
+  }
+
+  /**
+   * Read a single column's full value list via the row `scan()` path, the
+   * oracle scanColumn must agree with.
+   *
+   * @param {AsyncDataSource} source
+   * @param {string} column
+   * @returns {Promise<any[]>}
+   */
+  async function scanColumnOracle(source, column) {
+    /** @type {any[]} */
+    const out = []
+    for await (const row of source.scan({ columns: [column] }).rows()) {
+      out.push((row.resolved ?? {})[column])
+    }
+    return out
+  }
+
+  it('streams a column in row order, matching scan() for the same column', async () => {
+    const source = await icebergDataSource({ tableUrl, resolver, metadataFileName: 'v2.metadata.json' })
+    const expected = await scanColumnOracle(source, 'Popularity Rank')
+
+    const { scanColumn } = source
+    if (!scanColumn) throw new Error('scanColumn not implemented')
+    const { values, chunks } = await drain(scanColumn({ column: 'Popularity Rank' }))
+
+    expect(values).toEqual(expected)
+    expect(values).toHaveLength(21)
+    // At least one chunk was yielded (values arrive incrementally, not as one
+    // pre-materialized array the source built internally).
+    expect(chunks).toBeGreaterThanOrEqual(1)
+  })
+
+  it('respects LIMIT and OFFSET (no deletes)', async () => {
+    const source = await icebergDataSource({ tableUrl, resolver, metadataFileName: 'v2.metadata.json' })
+    const full = await scanColumnOracle(source, 'Popularity Rank')
+
+    const { scanColumn } = source
+    if (!scanColumn) throw new Error('scanColumn not implemented')
+    for (const [offset, limit] of [[0, 5], [2, 3], [18, 10], [0, 0], [5, 0]]) {
+      const { values } = await drain(scanColumn({ column: 'Popularity Rank', offset, limit }))
+      expect(values).toEqual(full.slice(offset, offset + limit))
+    }
+    // OFFSET past the end yields nothing.
+    const { values: past } = await drain(scanColumn({ column: 'Popularity Rank', offset: 100, limit: 5 }))
+    expect(past).toEqual([])
+  })
+
+  it('honors cross-file LIMIT/OFFSET against a multi-data-file table', async () => {
+    // The only fixture that exercises scanColumn's cross-file OFFSET whole-file
+    // skip-and-resume and cross-file LIMIT early-break (the off-by-one-prone
+    // sites). Asserted against the scan() oracle slice in every case.
+    const source = await icebergDataSource({ tableUrl: renameTableUrl, resolver, metadataFileName: 'v2.metadata.json' })
+    const full = await scanColumnOracle(source, 'id')
+    expect(full).toHaveLength(4)
+
+    const { scanColumn } = source
+    if (!scanColumn) throw new Error('scanColumn not implemented')
+
+    // Full-column read streams all three files, one chunk per row group/file.
+    const { values: all, chunks: allChunks } = await drain(scanColumn({ column: 'id' }))
+    expect(all).toEqual(full)
+    expect(allChunks).toBe(3)
+
+    // (a) OFFSET that crosses a data-file boundary: offset 1 skips the whole
+    // first file and resumes in the second; offset 2 skips the first two files;
+    // offset 3 lands inside the final file; offset+limit may overshoot the end.
+    for (const [offset, limit] of [[1, 2], [2, 2], [1, 10], [3, 1]]) {
+      const { values } = await drain(scanColumn({ column: 'id', offset, limit }))
+      expect(values).toEqual(full.slice(offset, offset + limit))
+    }
+
+    // (b) LIMIT satisfied before the last file: limit 2 is filled by the first
+    // two single-row files, so the final (two-row) file is never opened —
+    // proven by the chunk count dropping below the full-read 3.
+    const { values: capped, chunks: cappedChunks } = await drain(scanColumn({ column: 'id', limit: 2 }))
+    expect(capped).toEqual(full.slice(0, 2))
+    expect(cappedChunks).toBe(2)
+  })
+
+  it('applies LIMIT/OFFSET over post-delete values when deletes are present', async () => {
+    const source = await icebergDataSource({ tableUrl, resolver, metadataFileName: 'v4.metadata.json' })
+    const full = await scanColumnOracle(source, 'Breed Name')
+    expect(full).toHaveLength(15)
+
+    const { scanColumn } = source
+    if (!scanColumn) throw new Error('scanColumn not implemented')
+    // Whole column matches the post-delete scan.
+    const { values: all } = await drain(scanColumn({ column: 'Breed Name' }))
+    expect(all).toEqual(full)
+
+    // LIMIT/OFFSET are post-delete coordinates (record_count is pre-delete), so
+    // they must align with slices of the post-delete oracle, not raw files.
+    for (const [offset, limit] of [[0, 1], [4, 2], [8, 4], [13, 5]]) {
+      const { values } = await drain(scanColumn({ column: 'Breed Name', offset, limit }))
+      expect(values).toEqual(full.slice(offset, offset + limit))
+    }
+  })
+
+  it('aborts promptly on a pre-aborted signal', async () => {
+    const source = await icebergDataSource({ tableUrl, resolver, metadataFileName: 'v2.metadata.json' })
+    const { scanColumn } = source
+    if (!scanColumn) throw new Error('scanColumn not implemented')
+
+    const controller = new AbortController()
+    controller.abort()
+    await expect(drain(scanColumn({ column: 'Popularity Rank', signal: controller.signal })))
+      .rejects.toThrow('Aborted')
+  })
+
+  it('aborts mid-stream after consuming the first chunk', async () => {
+    // The pre-aborted case above only covers the entry guard. Use the
+    // multi-file fixture so there is more to read after the first chunk:
+    // consume the first file's chunk successfully, then abort, and the next
+    // pull must reject (the same error scan raises) at the between-file guard,
+    // honoring the JSDoc's "aborts between chunks" promise.
+    const source = await icebergDataSource({ tableUrl: renameTableUrl, resolver, metadataFileName: 'v2.metadata.json' })
+    const { scanColumn } = source
+    if (!scanColumn) throw new Error('scanColumn not implemented')
+
+    const controller = new AbortController()
+    const iterator = scanColumn({ column: 'id', signal: controller.signal })[Symbol.asyncIterator]()
+
+    const first = await iterator.next()
+    expect(first.done).toBe(false)
+    expect(first.value.length).toBeGreaterThanOrEqual(1)
+
+    controller.abort()
+    await expect(iterator.next()).rejects.toThrow('Aborted')
+  })
+
+  it('lights squirreling\'s streaming scalar-aggregate fast path', async () => {
+    const source = await icebergDataSource({ tableUrl, resolver, metadataFileName: 'v2.metadata.json' })
+    const baseScanColumn = source.scanColumn
+    if (!baseScanColumn) throw new Error('scanColumn not implemented')
+
+    // Spy on scanColumn so we can prove the engine took the streaming path
+    // (tryColumnScanAggregate) rather than the buffering aggregate fallback.
+    let scanColumnCalls = 0
+    /** @type {AsyncDataSource} */
+    const spied = {
+      ...source,
+      /** @type {NonNullable<AsyncDataSource['scanColumn']>} */
+      scanColumn(options) {
+        scanColumnCalls++
+        return baseScanColumn(options)
+      },
+    }
+
+    const result = await collect(executeSql({
+      tables: { bunnies: spied },
+      query: 'SELECT COUNT("Popularity Rank") AS c, MIN("Popularity Rank") AS mn, MAX("Popularity Rank") AS mx, SUM("Popularity Rank") AS s, AVG("Popularity Rank") AS a FROM bunnies',
+    }))
+
+    expect(scanColumnCalls).toBeGreaterThanOrEqual(1)
+    expect(result).toHaveLength(1)
+    const row = result[0]
+    expect(Number(row.c)).toBe(21)
+    expect(Number(row.mn)).toBe(1)
+    expect(Number(row.mx)).toBe(21)
+    expect(Number(row.s)).toBe(231)
+    expect(Number(row.a)).toBe(11)
+  })
+
+  it('serves a plain single-column SELECT with LIMIT/OFFSET through the hook', async () => {
+    // execute.js takes the scanColumn fast path for any single-column,
+    // WHERE-free scan, not just aggregates. Prove the hook is invoked and the
+    // streamed rows match the ordinary scan() path (hook removed).
+    const source = await icebergDataSource({ tableUrl, resolver, metadataFileName: 'v2.metadata.json' })
+    const baseScanColumn = source.scanColumn
+    if (!baseScanColumn) throw new Error('scanColumn not implemented')
+
+    let scanColumnCalls = 0
+    /** @type {AsyncDataSource} */
+    const spied = {
+      ...source,
+      /** @type {NonNullable<AsyncDataSource['scanColumn']>} */
+      scanColumn(options) {
+        scanColumnCalls++
+        return baseScanColumn(options)
+      },
+    }
+    // Same source with the hook removed forces the ordinary row scan() path.
+    /** @type {AsyncDataSource} */
+    const noHook = { ...source, scanColumn: undefined }
+
+    const query = 'SELECT "Popularity Rank" FROM bunnies LIMIT 5 OFFSET 2'
+    const viaHook = await collect(executeSql({ tables: { bunnies: spied }, query }))
+    const viaScan = await collect(executeSql({ tables: { bunnies: noHook }, query }))
+
+    expect(scanColumnCalls).toBe(1)
+    expect(viaHook).toEqual(viaScan)
+    expect(viaHook).toHaveLength(5)
+  })
+
+  it('characterizes: N aggregates on one column re-scan it N times (no coalescing)', async () => {
+    // Current behavior, pinned for visibility — NOT an endorsement. squirreling
+    // runs the streaming-aggregate fast path once per aggregate, so five
+    // aggregates over one column re-scan the column five times. Coalescing them
+    // into one pass is an upstream squirreling opportunity, not an icebird bug;
+    // this characterization makes a future fix a visible diff here.
+    const source = await icebergDataSource({ tableUrl, resolver, metadataFileName: 'v2.metadata.json' })
+    const baseScanColumn = source.scanColumn
+    if (!baseScanColumn) throw new Error('scanColumn not implemented')
+
+    let scanColumnCalls = 0
+    /** @type {AsyncDataSource} */
+    const spied = {
+      ...source,
+      /** @type {NonNullable<AsyncDataSource['scanColumn']>} */
+      scanColumn(options) {
+        scanColumnCalls++
+        return baseScanColumn(options)
+      },
+    }
+
+    await collect(executeSql({
+      tables: { bunnies: spied },
+      query: 'SELECT COUNT("Popularity Rank") AS c, MIN("Popularity Rank") AS mn, MAX("Popularity Rank") AS mx, SUM("Popularity Rank") AS s, AVG("Popularity Rank") AS a FROM bunnies',
+    }))
+
+    expect(scanColumnCalls).toBe(5)
+  })
+})