Add native prepared Parquet batches

.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

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"

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<AsyncBatch>}
+ */
+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<number, {name?: string, constant?: SqlPrimitive}>} */
+  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<ColumnVector>} | 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<DecodedArray>}
+ */
+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