diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3662e54..51bcd9f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -9,6 +9,9 @@ jobs: timeout-minutes: 5 steps: - uses: actions/checkout@v6 + - uses: actions/setup-node@v6 + with: + node-version: 24 - run: npm i - run: npm run lint @@ -17,6 +20,9 @@ jobs: timeout-minutes: 5 steps: - uses: actions/checkout@v6 + - uses: actions/setup-node@v6 + with: + node-version: 24 - run: npm i - run: npx tsc @@ -25,5 +31,8 @@ jobs: timeout-minutes: 5 steps: - uses: actions/checkout@v6 + - uses: actions/setup-node@v6 + with: + node-version: 24 - run: npm i - run: npm run coverage diff --git a/package.json b/package.json index 1d7b538..731749c 100644 --- a/package.json +++ b/package.json @@ -56,7 +56,7 @@ "hyparquet": "1.28.2", "hyparquet-compressors": "1.1.1", "hyparquet-writer": "0.16.6", - "squirreling": "0.15.3" + "squirreling": "0.16.1" }, "peerDependencies": { "@aws-sdk/credential-providers": "^3.0.0" diff --git a/src/read.js b/src/read.js index c538095..f09008b 100644 --- a/src/read.js +++ b/src/read.js @@ -1,5 +1,8 @@ -import { cachedAsyncBuffer, parquetMetadataAsync, parquetReadObjects } from 'hyparquet' +import { cachedAsyncBuffer, flatten, parquetMetadataAsync, parquetReadObjects, parquetSchema } from 'hyparquet' import { compressors } from 'hyparquet-compressors' +import { parquetReadAsync } from 'hyparquet/src/read.js' +import { assembleAsync } from 'hyparquet/src/rowgroup.js' +import { selectVector } from 'squirreling' import { fetchDeleteMaps, urlResolver } from './fetch.js' import { icebergMetadata } from './metadata.js' import { icebergManifests, splitManifestEntries } from './manifest.js' @@ -461,6 +464,288 @@ export async function* readDataFile({ } } +/** + * Stream native squirreling batches from one delete-free parquet data file. + * The batch envelope follows parquet row-group boundaries, while every + * requested column remains deferred until the engine asks for it. Deferred + * reads receive the engine's effective row selection, allowing a selective + * predicate and LIMIT to avoid decoding unused values from wide columns. + * + * Delete-bearing tables intentionally use the legacy row scanner in + * `icebergDataSource`: position and equality deletes need coordinated access + * to multiple columns and row positions, which this independent-column reader + * does not attempt to reproduce. + * + * @import {AsyncBatch, ColumnVector, Field as BatchField, RowSelection, SqlPrimitive} from 'squirreling' + * @param {object} options + * @param {ManifestEntry} options.dataEntry + * @param {Schema} options.schema + * @param {TableMetadata} options.metadata + * @param {Resolver} options.resolver + * @param {readonly BatchField[]} options.fields - Requested fields in batch-column order. + * @param {AbortSignal} [options.signal] + * @returns {AsyncGenerator} + */ +export async function* readDataFileBatches({ + dataEntry, + schema, + metadata, + resolver, + fields, + signal, +}) { + const { data_file, partition_spec_id } = dataEntry + signal?.throwIfAborted() + + const partitionSpec = metadata['partition-specs'].find(s => s['spec-id'] === partition_spec_id) + const resolved = await resolver.reader(data_file.file_path, Number(data_file.file_size_in_bytes)) + const file = cachedAsyncBuffer(resolved) + const parquetMetadata = await parquetMetadataAsync(file) + signal?.throwIfAborted() + + const kv = parquetMetadata.key_value_metadata?.find(k => k.key === 'iceberg.schema') + /** @type {Schema} */ + let parquetIcebergSchema + if (kv?.value) { + parquetIcebergSchema = JSON.parse(kv.value) + } else if (parquetMetadata.schema.some(s => s.field_id !== undefined)) { + parquetIcebergSchema = parquetSchemaToIceberg(parquetMetadata.schema) + } else { + parquetIcebergSchema = schema + } + + const physicalNames = new Set(parquetSchema(parquetMetadata).children.map(child => child.element.name)) + /** @type {NameMapping[]} */ + const nameMappings = metadata.properties?.['schema.name-mapping.default'] + ? JSON.parse(metadata.properties['schema.name-mapping.default']) + : [] + + /** @type {Map} */ + const fieldSources = new Map() + for (const requested of fields) { + const field = schema.fields.find(candidate => candidate.id === requested.id) + if (!field) throw new Error(`Iceberg field id ${requested.id} not found`) + const parquetField = parquetIcebergSchema.fields.find(candidate => candidate.id === field.id) + let physicalName = parquetField && field.type !== 'unknown' + ? sanitize(parquetField.name) + : undefined + if (physicalName && !physicalNames.has(physicalName)) physicalName = undefined + + if (!physicalName) { + const mapping = nameMappingById(nameMappings, field.id) + physicalName = mapping?.names + .map(name => sanitize(name)) + .find(name => physicalNames.has(name)) + } + if (physicalName) { + fieldSources.set(requested.id, { name: physicalName }) + continue + } + + const partitionField = partitionSpec?.fields.find( + candidate => candidate['source-id'] === field.id && candidate.transform === 'identity') + let constant + if (partitionField && Object.hasOwn(data_file.partition, partitionField.name)) { + constant = data_file.partition[partitionField.name] ?? null + } else if (field['initial-default'] !== undefined) { + constant = field['initial-default'] + } else { + constant = null + } + fieldSources.set(requested.id, { constant: /** @type {SqlPrimitive} */ (constant) }) + } + + const schemaTree = parquetSchema(parquetMetadata) + let groupStart = 0 + for (const rowGroup of parquetMetadata.row_groups) { + signal?.throwIfAborted() + const groupRows = Number(rowGroup.num_rows) + if (groupRows === 0) continue + const batchStart = groupStart + /** @type {AsyncBatch} */ + const batch = { + selection: { type: 'all', length: groupRows }, + columns: fields.map(function batchColumn(field) { + const source = fieldSources.get(field.id) + if (!source) throw new Error(`Iceberg field id ${field.id} has no read source`) + if (!source.name) { + return { + type: 'constant', + value: source.constant ?? null, + length: groupRows, + } + } + const columnName = source.name + /** @type {{start: number, end: number, vector: Promise} | undefined} */ + let cachedRange + return { + async read({ selection, signal: readSignal }) { + readSignal?.throwIfAborted() + const count = selectedRows(selection) + if (count === 0) return { type: 'values', values: [], length: 0 } + const range = selectionRange(selection) + if (cachedRange && range.start >= cachedRange.start && range.end <= cachedRange.end) { + const vector = await cachedRange.vector + readSignal?.throwIfAborted() + return selectDecodedVector(vector, selection, cachedRange.start) + } + const values = readParquetColumnRange({ + file, + parquetMetadata, + schemaTree, + columnName, + rowStart: batchStart + range.start, + rowEnd: batchStart + range.end, + signal: readSignal, + }) + const vector = values.then(decodedColumnVector) + cachedRange = { ...range, vector } + return selectDecodedVector(await vector, selection, range.start) + }, + } + }), + } + yield batch + groupStart += groupRows + } +} + +/** + * Read one physical top-level parquet column over an absolute file row range. + * Offset indexes are used when available so a narrowed selection can skip + * unrelated data pages. + * + * @import {AsyncBuffer, DecodedArray, FileMetaData, SchemaTree} from 'hyparquet' + * @param {object} options + * @param {AsyncBuffer} options.file + * @param {FileMetaData} options.parquetMetadata + * @param {SchemaTree} options.schemaTree + * @param {string} options.columnName + * @param {number} options.rowStart + * @param {number} options.rowEnd + * @param {AbortSignal} [options.signal] + * @returns {Promise} + */ +async function readParquetColumnRange({ + file, + parquetMetadata, + schemaTree, + columnName, + rowStart, + rowEnd, + signal, +}) { + const asyncGroups = parquetReadAsync({ + file, + metadata: parquetMetadata, + columns: [columnName], + rowStart, + rowEnd, + compressors, + utf8: false, + useOffsetIndex: true, + }).map(group => assembleAsync(group, schemaTree)) + /** @type {DecodedArray[]} */ + const chunks = [] + for (const group of asyncGroups) { + signal?.throwIfAborted() + const column = group.asyncColumns.find(candidate => candidate.pathInSchema[0] === columnName) + if (!column) throw new Error(`parquet column not found: ${columnName}`) + const result = await column.data + const data = flatten(result.data) + const start = group.selectStart ?? Math.max(rowStart - group.groupStart, 0) + const end = group.selectEnd ?? Math.min(rowEnd - group.groupStart, group.groupRows) + const localStart = start - result.skipped + const localEnd = end - result.skipped + chunks.push(localStart === 0 && localEnd === data.length + ? data + : data.slice(localStart, localEnd)) + } + signal?.throwIfAborted() + return flatten(chunks) +} + +/** + * Preserve hyparquet's decoded arrays as native Squirreling vectors. Numeric + * typed arrays avoid boxing and ordinary arrays can be passed through without + * copying. + * + * @param {DecodedArray} values + * @returns {ColumnVector} + */ +function decodedColumnVector(values) { + if (Array.isArray(values)) { + return { + type: 'values', + values: /** @type {SqlPrimitive[]} */ (values), + length: values.length, + } + } + return { type: 'typed', values, length: values.length } +} + +/** + * Select rows from a decoded covering range. Selections use the enclosing + * batch's row coordinates, so translate them to the cached vector's local + * coordinates before creating a zero-copy selected view. + * + * @param {ColumnVector} vector + * @param {RowSelection} selection + * @param {number} rangeStart + * @returns {ColumnVector} + */ +function selectDecodedVector(vector, selection, rangeStart) { + if (selection.type === 'all') return vector + if (selection.type === 'range') { + if (selection.start === rangeStart && selection.end - selection.start === vector.length) { + return vector + } + return selectVector(vector, { + type: 'range', + start: selection.start - rangeStart, + end: selection.end - rangeStart, + length: vector.length, + }) + } + const indices = new Uint32Array(selection.indices.length) + for (let index = 0; index < indices.length; index++) { + indices[index] = selection.indices[index] - rangeStart + } + return selectVector(vector, { + type: 'indices', + indices, + length: vector.length, + }) +} + +/** + * @param {RowSelection} selection + * @returns {number} + */ +function selectedRows(selection) { + if (selection.type === 'all') return selection.length + if (selection.type === 'range') return selection.end - selection.start + return selection.indices.length +} + +/** + * Return the smallest local row range covering a selection. + * + * @param {RowSelection} selection + * @returns {{start: number, end: number}} + */ +function selectionRange(selection) { + if (selection.type === 'all') return { start: 0, end: selection.length } + if (selection.type === 'range') return { start: selection.start, end: selection.end } + let start = Infinity + let end = 0 + for (const index of selection.indices) { + start = Math.min(start, index) + end = Math.max(end, index + 1) + } + return selection.indices.length === 0 ? { start: 0, end: 0 } : { start, end } +} + /** * Recursively rewrite top-level column references in a parquet filter using * the provided mapping. Returns undefined if any referenced column has no diff --git a/src/sql/icebergDataSource.js b/src/sql/icebergDataSource.js index 97ac720..586b64b 100644 --- a/src/sql/icebergDataSource.js +++ b/src/sql/icebergDataSource.js @@ -1,17 +1,31 @@ -import { asyncRow } from 'squirreling' +import { asyncRow, rowsToBatches } from 'squirreling' import { fetchDeleteMaps, urlResolver } from '../fetch.js' import { icebergManifests, splitManifestEntries } from '../manifest.js' import { icebergMetadata } from '../metadata.js' -import { readDataFile } from '../read.js' +import { readDataFile, readDataFileBatches } from '../read.js' import { fileMightMatch, partitionMightMatch } from '../prune.js' import { whereToParquetFilter } from './whereFilter.js' /** - * @import {AsyncDataSource, ExprNode, SqlPrimitive} from 'squirreling' + * @import {AsyncDataSource, ExprNode, PrepareScan, RelationSchema, ScanResults, SqlPrimitive} from 'squirreling' * @import {ScanColumnResults} from 'squirreling/src/types.js' * @import {Lister, Resolver, TableMetadata} from '../../src/types.js' */ +/** + * Icebird keeps the legacy scan surface alongside prepared batches so callers + * written against older squirreling releases retain a statically callable + * `scan()` method. + * + * @typedef {object} IcebergAsyncDataSource + * @property {number} [numRows] + * @property {string[]} columns + * @property {RelationSchema} schema + * @property {(options: import('squirreling').ScanOptions) => ScanResults} scan + * @property {PrepareScan} prepareScan + * @property {NonNullable} scanColumn + */ + /** * Creates a squirreling AsyncDataSource backed by an Iceberg table that streams * rows lazily from the underlying parquet data files (row group by row group, @@ -47,7 +61,7 @@ import { whereToParquetFilter } from './whereFilter.js' * @param {number | bigint} [options.snapshotId] - Optional snapshot id for time travel; defaults to the current snapshot. * @param {Resolver} [options.resolver] - I/O resolver (defaults to `urlResolver()`). * @param {Lister} [options.lister] - Directory lister, used to discover the latest metadata. - * @returns {Promise} + * @returns {Promise} */ export async function icebergDataSource({ tableUrl, metadataFileName, metadata, snapshotId, resolver, lister }) { if (!tableUrl) throw new Error('tableUrl is required') @@ -64,6 +78,15 @@ export async function icebergDataSource({ tableUrl, metadataFileName, metadata, const schema = tableMetadata.schemas.find(s => s['schema-id'] === schemaId) if (!schema) throw new Error('schema not found in metadata') const columns = schema.fields.map(f => f.name) + /** @type {RelationSchema} */ + const relationSchema = { + fields: schema.fields.map(field => ({ + id: field.id, + name: field.name, + dataType: { type: 'unknown' }, + nullable: !field.required, + })), + } const rowLineage = tableMetadata['format-version'] >= 3 const manifestList = await icebergManifests({ metadata: tableMetadata, resolver: fetchResolver, snapshotId }) @@ -85,9 +108,59 @@ export async function icebergDataSource({ tableUrl, metadataFileName, metadata, } } - return { + /** @type {IcebergAsyncDataSource} */ + const thisSource = { numRows, columns, + schema: relationSchema, + prepareScan(request) { + const requestedFields = request.columns.map(demand => { + const field = relationSchema.fields.find(candidate => candidate.id === demand.field) + if (!field) throw new Error(`Prepared scan requested unknown field id ${demand.field}`) + return field + }) + const requestedNames = requestedFields.map(field => field.name) + const filter = whereToParquetFilter(request.filter) + const scanEntries = filter + ? dataEntries.filter(entry => + partitionMightMatch(filter, entry, schema, tableMetadata) && + fileMightMatch(filter, entry, schema)) + : dataEntries + let maxRows = 0 + for (const entry of scanEntries) maxRows += Number(entry.data_file.record_count) + + return { + schema: { fields: requestedFields }, + residual: { + filter: request.filter, + limit: request.limit, + offset: request.offset, + }, + properties: { + exactRows: request.filter || hasDeletes ? undefined : maxRows, + maxRows, + }, + async *batches({ signal } = {}) { + signal?.throwIfAborted() + if (hasDeletes) { + const legacy = thisSource.scan({ columns: requestedNames, signal }) + yield* rowsToBatches(legacy.rows(), requestedNames, { signal }) + return + } + for (const entry of scanEntries) { + signal?.throwIfAborted() + yield* readDataFileBatches({ + dataEntry: entry, + schema, + metadata: tableMetadata, + resolver: fetchResolver, + fields: requestedFields, + signal, + }) + } + }, + } + }, scan({ columns: scanColumns, where, limit, offset, signal }) { const rowColumns = scanColumns ?? columns // Convert the WHERE AST to a hyparquet filter; undefined means the @@ -317,4 +390,5 @@ export async function icebergDataSource({ tableUrl, metadataFileName, metadata, } }, } + return thisSource } diff --git a/test/sql/icebergDataSource.test.js b/test/sql/icebergDataSource.test.js index 32ba54a..ace7507 100644 --- a/test/sql/icebergDataSource.test.js +++ b/test/sql/icebergDataSource.test.js @@ -1,4 +1,4 @@ -import { collect, executeSql } from 'squirreling' +import { collect, executeSql, readBatchColumn, valueAt } from 'squirreling' import { describe, expect, it } from 'vitest' import { icebergDataSource } from '../../src/sql/icebergDataSource.js' import { localResolver } from '../helpers.js' @@ -38,6 +38,87 @@ describe.concurrent('icebergDataSource', () => { 'Temperament', 'Popularity Rank', ]) + expect(source.schema.fields.map(field => field.name)).toEqual(source.columns) + }) + + it('keeps parquet columns deferred and reads only a requested selection', async () => { + const source = await icebergDataSource({ + tableUrl, + resolver, + metadataFileName: 'v2.metadata.json', + }) + const prepared = source.prepareScan({ + columns: [ + { field: 1, phase: 1, purpose: 'output', mode: 'deferred' }, + { field: 8, phase: 0, purpose: 'filter', mode: 'required' }, + ], + }) + const iterator = prepared.batches()[Symbol.asyncIterator]() + const first = await iterator.next() + expect(first.done).toBe(false) + if (first.done) throw new Error('expected a parquet batch') + const batch = first.value + expect(batch.columns.every(column => 'read' in column)).toBe(true) + + const allNames = await readBatchColumn({ batch, columnIndex: 0 }) + expect(allNames.type).toBe('values') + const last = batch.selection.length - 1 + const selection = { + type: /** @type {const} */ ('indices'), + indices: new Uint32Array([1, last]), + length: batch.selection.length, + } + const selectedNames = await readBatchColumn({ batch, columnIndex: 0, selection }) + expect(selectedNames.type).toBe('selected') + if (selectedNames.type !== 'selected') throw new Error('expected selected vector') + expect(selectedNames.source).toBe(allNames) + expect(selectedNames.length).toBe(2) + expect(valueAt(selectedNames, 0)).toEqual(valueAt(allNames, 1)) + expect(valueAt(selectedNames, 1)).toEqual(valueAt(allNames, last)) + + const rankRange = { + type: /** @type {const} */ ('range'), + start: 1, + end: last, + length: batch.selection.length, + } + const rangeRanks = await readBatchColumn({ batch, columnIndex: 1, selection: rankRange }) + const rankSelection = { + type: /** @type {const} */ ('indices'), + indices: new Uint32Array([2, last - 1]), + length: batch.selection.length, + } + const selectedRanks = await readBatchColumn({ batch, columnIndex: 1, selection: rankSelection }) + expect(selectedRanks.type).toBe('selected') + if (selectedRanks.type !== 'selected') throw new Error('expected selected vector') + expect(selectedRanks.source).toBe(rangeRanks) + expect(valueAt(selectedRanks, 0)).toEqual(valueAt(rangeRanks, 1)) + expect(valueAt(selectedRanks, 1)).toEqual(valueAt(rangeRanks, rangeRanks.length - 1)) + await iterator.return?.() + }) + + it('executes through native prepared batches without calling legacy scan hooks', async () => { + const source = await icebergDataSource({ + tableUrl, + resolver, + metadataFileName: 'v2.metadata.json', + }) + /** @type {AsyncDataSource} */ + const preparedOnly = { + ...source, + scan() { + throw new Error('legacy scan should not run') + }, + scanColumn() { + throw new Error('legacy scanColumn should not run') + }, + } + const result = await collect(executeSql({ + tables: { bunnies: preparedOnly }, + query: 'SELECT "Breed Name" FROM bunnies WHERE "Popularity Rank" > 18 LIMIT 2', + })) + expect(result).toHaveLength(2) + expect(result.every(row => typeof row['Breed Name'] === 'string')).toBe(true) }) it('streams all rows via scan()', async () => { @@ -504,7 +585,7 @@ describe.concurrent('icebergDataSource scanColumn', () => { * Read a single column's full value list via the row `scan()` path, the * oracle scanColumn must agree with. * - * @param {AsyncDataSource} source + * @param {Awaited>} source * @param {string} column * @returns {Promise} */ @@ -653,7 +734,7 @@ describe.concurrent('icebergDataSource scanColumn', () => { * Filtered oracle: the same single column read via scan() with a WHERE, which * icebird already prunes and matches per row. scanColumn's pushdown must agree. * - * @param {AsyncDataSource} source + * @param {Awaited>} source * @param {string} column * @param {ExprNode} where * @returns {Promise} @@ -763,9 +844,9 @@ describe.concurrent('icebergDataSource scanColumn', () => { }) 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). + // Prepared scans supersede the legacy scanColumn hook. Disable prepareScan + // explicitly here to retain coverage of the backwards-compatible path and + // prove its streamed rows still match ordinary scan(). const source = await icebergDataSource({ tableUrl, resolver, metadataFileName: 'v2.metadata.json' }) const baseScanColumn = source.scanColumn if (!baseScanColumn) throw new Error('scanColumn not implemented') @@ -774,6 +855,7 @@ describe.concurrent('icebergDataSource scanColumn', () => { /** @type {AsyncDataSource} */ const spied = { ...source, + prepareScan: undefined, /** @type {NonNullable} */ scanColumn(options) { scanColumnCalls++ @@ -782,7 +864,7 @@ describe.concurrent('icebergDataSource scanColumn', () => { } // Same source with the hook removed forces the ordinary row scan() path. /** @type {AsyncDataSource} */ - const noHook = { ...source, scanColumn: undefined } + const noHook = { ...source, prepareScan: undefined, scanColumn: undefined } const query = 'SELECT "Popularity Rank" FROM bunnies LIMIT 5 OFFSET 2' const viaHook = await collect(executeSql({ tables: { bunnies: spied }, query }))