Support format-version 3 in icebergRewrite (preserve row lineage)

src/write/rewrite.js

@@ -14,7 +14,7 @@ import { computeColumnStats } from './stats.js'
 import { checkWriteFormat, newSnapshotId, resolveParquetCodec } from './stage.js'
 
 /**
- * @import {Manifest, Resolver, Snapshot, StagedUpdate, TableMetadata} from '../../src/types.js'
+ * @import {Manifest, Resolver, Schema, Snapshot, StagedUpdate, TableMetadata} from '../../src/types.js'
  */
 
 /**
@@ -30,8 +30,16 @@ import { checkWriteFormat, newSnapshotId, resolveParquetCodec } from './stage.js
  * boundary). Row contents and counts are preserved (modulo deleted rows and
  * order); deletes are consumed, so the new snapshot has no delete files.
  *
- * v2 only for now: a v3 rewrite would have to preserve `_row_id` row lineage
- * rather than let `assignFirstRowIds` renumber the rewritten rows.
+ * On v3 tables, row lineage is preserved: each surviving row's `_row_id` and
+ * `_last_updated_sequence_number` are materialized as explicit columns in the
+ * rewritten files (the global sort breaks positional derivation, so stored
+ * values are required). When every live row carries lineage, the new manifest's
+ * `first_row_id` is pinned to the minimum carried `_row_id` so no new row ids
+ * are consumed (`next-row-id` does not advance). When some rows lack lineage
+ * (a table upgraded from v2 whose pre-upgrade rows were never assigned ids),
+ * the manifest is left for commit-time assignment per the spec: stored ids
+ * still win on read, null rows get derived ids, and `next-row-id` advances by
+ * the manifest's row count.
  *
  * @param {object} options
  * @param {string} options.tableUrl
@@ -48,10 +56,10 @@ export async function icebergStageRewrite({
   if (!tableUrl) throw new Error('tableUrl is required')
   if (!resolver?.writer) throw new Error('resolver.writer is required')
   const writerFn = resolver.writer
-  const formatVersion = metadata['format-version']
-  if (formatVersion !== 2) {
-    throw new Error(`icebergRewrite supports format-version 2 only (got ${formatVersion}); v3 row lineage is not yet handled`)
+  if (metadata['format-version'] !== 2 && metadata['format-version'] !== 3) {
+    throw new Error(`unsupported format-version: ${metadata['format-version']}`)
   }
+  const formatVersion = /** @type {2|3} */ (metadata['format-version'])
   if (targetFileRows !== undefined && !(targetFileRows > 0)) {
     throw new Error('targetFileRows must be a positive number')
   }
@@ -77,10 +85,21 @@ export async function icebergStageRewrite({
   checkWriteFormat(metadata.properties?.['write.format.default'])
   const codec = resolveParquetCodec(metadata.properties?.['write.parquet.compression-codec'])
 
-  // Read every live row (deletes applied), then sort globally.
+  // Read every live row (deletes applied), then sort globally. For v3 tables
+  // the rows carry `_row_id` / `_last_updated_sequence_number` (derived or
+  // stored by the read path).
   const liveRows = await icebergRead({ tableUrl, metadata, resolver })
   const sortedRows = comparator ? [...liveRows].sort(comparator) : liveRows
 
+  const rowLineage = formatVersion >= 3
+  // Preserve mode requires complete lineage: rows from pre-upgrade v2
+  // snapshots read with null ids and need commit-time assignment instead.
+  const allLineage = rowLineage && liveRows.length > 0 &&
+    liveRows.every(r => r._row_id != null && r._last_updated_sequence_number != null)
+  const minRowId = allLineage
+    ? liveRows.reduce((min, r) => r._row_id < min ? r._row_id : min, liveRows[0]._row_id)
+    : undefined
+
   // Regroup under the target partition spec (re-derives tuples from values, so
   // files written under an older spec are rewritten under the new one).
   const groups = partitionSpec.fields.length
@@ -90,6 +109,22 @@ export async function icebergStageRewrite({
   const snapshotId = newSnapshotId(metadata)
   const manifestUuid = uuid4()
 
+  // For v3, materialize the carried lineage as explicit columns in the
+  // rewritten files (reserved field ids per spec). The extended schema is
+  // passed to the parquet writer only: stats, partitioning, and the manifest's
+  // embedded schema stay user-only.
+  /** @type {Schema} */
+  const writeSchema = rowLineage
+    ? {
+      ...schema,
+      fields: [
+        ...schema.fields,
+        { id: 2147483540, name: '_row_id', required: false, type: 'long' },
+        { id: 2147483539, name: '_last_updated_sequence_number', required: false, type: 'long' },
+      ],
+    }
+    : schema
+
   /** @type {{ partition: Record<string, any>, dataFile: any, path: string }[]} */
   const writtenDataFiles = []
   for (const group of groups) {
@@ -98,7 +133,7 @@ export async function icebergStageRewrite({
       if (chunk.length === 0) continue
       const dataPath = `${tableUrl}/data/${uuid4()}.parquet`
       const dataWriter = writerFn(dataPath)
-      await writeParquet({ writer: dataWriter, schema, records: chunk, codec })
+      await writeParquet({ writer: dataWriter, schema: writeSchema, records: chunk, codec })
       const stats = computeColumnStats(chunk, schema)
       writtenDataFiles.push({
         partition: group.partition,
@@ -160,6 +195,14 @@ export async function icebergStageRewrite({
     deleted_rows_count: 0n,
     partitions,
   }
+  if (allLineage) {
+    // Every rewritten row carries a materialized `_row_id`, so the manifest
+    // does not need a fresh id range from `assignFirstRowIds`. Pin its
+    // first_row_id to the smallest carried id: ids are distinct and below
+    // next-row-id, so min + row count never exceeds next-row-id and no new
+    // ids are consumed (`next-row-id` does not advance for a pure rewrite).
+    newManifest.first_row_id = minRowId
+  }
 
   // Supersede every prior manifest (data + delete) for the rewritten snapshot.
   const priorManifests = await loadPriorManifests(metadata, resolver)

test/write/rewrite.test.js

@@ -1,11 +1,13 @@
 import { describe, expect, it, vi } from 'vitest'
 import { fileCatalog } from '../../src/catalog/file.js'
+import { fetchAvroRecords } from '../../src/fetch.js'
 import { fileCatalogCommit } from '../../src/write/commit.js'
 import { icebergCreate } from '../../src/create.js'
 import { icebergManifests, splitManifestEntries } from '../../src/manifest.js'
 import { icebergRead } from '../../src/read.js'
 import { icebergAppend, icebergRewrite } from '../../src/write/write.js'
 import { icebergStageAppend } from '../../src/write/stage.js'
+import { icebergStageDeletionVector } from '../../src/write/stage-deletion-vector.js'
 import { icebergStagePositionDelete } from '../../src/write/stage-position-delete.js'
 import { icebergStageRewrite } from '../../src/write/rewrite.js'
 import { deserializeValue } from '../../src/write/serde.js'
@@ -54,17 +56,27 @@ function longBound(entry, fieldId, side) {
   return bytes ? deserializeValue(bytes, 'long') : undefined
 }
 
+/**
+ * Map each row's id to its lineage pair, for before/after comparison.
+ * @param {Record<string, any>[]} rows
+ * @returns {Map<any, {rowId: any, lusn: any}>}
+ */
+function lineageById(rows) {
+  return new Map(rows.map(r => [r.id, { rowId: r._row_id, lusn: r._last_updated_sequence_number }]))
+}
+
 /**
  * Create a sorted, multi-file table and return its committed metadata.
  * @param {object} [opts]
  * @param {SortOrder} [opts.sortOrder]
+ * @param {2 | 3} [opts.formatVersion]
  * @returns {Promise<{ tableUrl: string, resolver: import('../../src/types.js').Resolver, metadata: TableMetadata }>}
  */
-async function makeMultiFileTable({ sortOrder } = {}) {
+async function makeMultiFileTable({ sortOrder, formatVersion } = {}) {
   vi.spyOn(Date, 'now').mockReturnValue(1700000000000)
   const tableUrl = 'mem://rewrite'
   const { resolver } = memResolver()
-  let metadata = await icebergCreate({ tableUrl, resolver, schema, sortOrder })
+  let metadata = await icebergCreate({ tableUrl, resolver, schema, sortOrder, formatVersion })
   const batches = [
     [{ id: 5n, name: 'e' }, { id: 2n, name: 'b' }],
     [{ id: 1n, name: 'a' }, { id: 6n, name: 'f' }],
@@ -209,16 +221,155 @@ describe('icebergRewrite — safety', () => {
     expect(rows.map(r => r.id).sort()).toEqual([1n, 2n])
   })
 
-  it('rejects format-version 3 (row lineage not yet handled)', async () => {
+  it('rejects unsupported format versions', async () => {
+    const { tableUrl, resolver, metadata } = await makeMultiFileTable()
+    await expect(() => icebergStageRewrite({
+      tableUrl, metadata: { ...metadata, 'format-version': 1 }, resolver,
+    })).rejects.toThrow(/unsupported format-version: 1/)
+  })
+})
+
+describe('icebergRewrite — v3 row lineage', () => {
+  it('preserves _row_id and _last_updated_sequence_number across a sorted rewrite', async () => {
+    const { tableUrl, resolver, metadata } = await makeMultiFileTable({ sortOrder: sortById, formatVersion: 3 })
+    expect(metadata['next-row-id']).toBe(6)
+    const before = await icebergRead({ tableUrl, metadata, resolver })
+    expect(new Set(before.map(r => r._row_id))).toEqual(new Set([0n, 1n, 2n, 3n, 4n, 5n]))
+
+    const staged = await icebergStageRewrite({ tableUrl, metadata, resolver })
+    // A pure rewrite carries existing rows, so it assigns no new row ids.
+    expect(staged.snapshot['first-row-id']).toBe(6)
+    expect(staged.snapshot['added-rows']).toBe(0)
+    expect(staged.requirements).toContainEqual({ type: 'assert-next-row-id', 'next-row-id': 6 })
+
+    const after = await fileCatalogCommit({ tableUrl, metadata, staged, resolver })
+    expect(after['next-row-id']).toBe(6)
+
+    const rows = await icebergRead({ tableUrl, metadata: after, resolver })
+    expect(rows.map(r => r.id)).toEqual([1n, 2n, 3n, 4n, 5n, 6n])
+    // The global sort scrambled row positions, so preserved values must come
+    // from the materialized _row_id column, not positional derivation.
+    expect(rows.map(r => r._row_id)).not.toEqual([0n, 1n, 2n, 3n, 4n, 5n])
+    expect(lineageById(rows)).toEqual(lineageById(before))
+
+    // The rewritten manifest is pinned to the carried range: min(_row_id).
+    const snapshot = after.snapshots?.find(s => s['snapshot-id'] === after['current-snapshot-id'])
+    const manifests = await fetchAvroRecords(snapshot?.['manifest-list'] ?? '', resolver)
+    expect(manifests.length).toBe(1)
+    expect(manifests[0].first_row_id).toBe(0n)
+  })
+
+  it('preserves lineage across split output files', async () => {
+    const { tableUrl, resolver, metadata } = await makeMultiFileTable({ sortOrder: sortById, formatVersion: 3 })
+    const before = await icebergRead({ tableUrl, metadata, resolver })
+
+    const staged = await icebergStageRewrite({ tableUrl, metadata, resolver, targetFileRows: 2 })
+    const after = await fileCatalogCommit({ tableUrl, metadata, staged, resolver })
+
+    const entries = splitManifestEntries(await icebergManifests({ metadata: after, resolver })).dataEntries
+    expect(entries.length).toBe(3)
+    const rows = await icebergRead({ tableUrl, metadata: after, resolver })
+    expect(lineageById(rows)).toEqual(lineageById(before))
+    expect(after['next-row-id']).toBe(6)
+  })
+
+  it('consumes deletion vectors and keeps surviving rows\' lineage', async () => {
     vi.spyOn(Date, 'now').mockReturnValue(1700000000000)
-    const tableUrl = 'mem://rewrite-v3'
+    const tableUrl = 'mem://rewrite-v3-dv'
     const { resolver } = memResolver()
     const created = await icebergCreate({ tableUrl, resolver, schema, formatVersion: 3 })
-    const appended = await icebergStageAppend({ tableUrl, metadata: created, records: [{ id: 1n, name: 'a' }], resolver })
+    const records = [{ id: 1n, name: 'a' }, { id: 2n, name: 'b' }, { id: 3n, name: 'c' }]
+    const appended = await icebergStageAppend({ tableUrl, metadata: created, records, resolver })
     const afterAppend = await fileCatalogCommit({ tableUrl, metadata: created, staged: appended, resolver })
+    const dataPath = appended.writtenFiles[0]
 
-    await expect(() => icebergStageRewrite({ tableUrl, metadata: afterAppend, resolver }))
-      .rejects.toThrow(/format-version 2 only/)
+    const delStaged = await icebergStageDeletionVector({
+      tableUrl, metadata: afterAppend, deletes: [{ file_path: dataPath, pos: 1n }], resolver,
+    })
+    const afterDelete = await fileCatalogCommit({ tableUrl, metadata: afterAppend, staged: delStaged, resolver })
+
+    const staged = await icebergStageRewrite({ tableUrl, metadata: afterDelete, resolver })
+    const afterRewrite = await fileCatalogCommit({ tableUrl, metadata: afterDelete, staged, resolver })
+
+    const rows = await icebergRead({ tableUrl, metadata: afterRewrite, resolver })
+    expect(rows.map(r => ({ id: r.id, _row_id: r._row_id, _last_updated_sequence_number: r._last_updated_sequence_number }))).toEqual([
+      { id: 1n, _row_id: 0n, _last_updated_sequence_number: 1n },
+      { id: 3n, _row_id: 2n, _last_updated_sequence_number: 1n },
+    ])
+    const { dataEntries, deleteEntries } = splitManifestEntries(await icebergManifests({ metadata: afterRewrite, resolver }))
+    expect(deleteEntries.length).toBe(0)
+    expect(dataEntries.length).toBe(1)
+    // The deleted row's id (1n) is retired with it; next-row-id is unchanged.
+    expect(afterRewrite['next-row-id']).toBe(3)
+  })
+
+  it('appends after a rewrite continue from the unchanged next-row-id', async () => {
+    const { tableUrl, resolver, metadata } = await makeMultiFileTable({ sortOrder: sortById, formatVersion: 3 })
+    const staged = await icebergStageRewrite({ tableUrl, metadata, resolver })
+    const after = await fileCatalogCommit({ tableUrl, metadata, staged, resolver })
+
+    const appended = await icebergStageAppend({
+      tableUrl, metadata: after, records: [{ id: 7n, name: 'g' }, { id: 8n, name: 'h' }], resolver,
+    })
+    const final = await fileCatalogCommit({ tableUrl, metadata: after, staged: appended, resolver })
+    expect(final['next-row-id']).toBe(8)
+
+    const rows = await icebergRead({ tableUrl, metadata: final, resolver })
+    // New rows take fresh ids 6n and 7n; no collision with the preserved 0n-5n.
+    expect(new Set(rows.map(r => r._row_id))).toEqual(new Set([0n, 1n, 2n, 3n, 4n, 5n, 6n, 7n]))
+    expect(rows.find(r => r.id === 7n)?._row_id).toBe(6n)
+    expect(rows.find(r => r.id === 8n)?._row_id).toBe(7n)
+  })
+
+  it('keeps lineage stable through a rewrite of a rewrite', async () => {
+    const { tableUrl, resolver, metadata } = await makeMultiFileTable({ sortOrder: sortById, formatVersion: 3 })
+    const before = await icebergRead({ tableUrl, metadata, resolver })
+
+    const staged1 = await icebergStageRewrite({ tableUrl, metadata, resolver })
+    const after1 = await fileCatalogCommit({ tableUrl, metadata, staged: staged1, resolver })
+    // Source files now carry materialized lineage columns.
+    const staged2 = await icebergStageRewrite({ tableUrl, metadata: after1, resolver, targetFileRows: 3 })
+    const after2 = await fileCatalogCommit({ tableUrl, metadata: after1, staged: staged2, resolver })
+
+    const rows = await icebergRead({ tableUrl, metadata: after2, resolver })
+    expect(lineageById(rows)).toEqual(lineageById(before))
+    expect(after2['next-row-id']).toBe(6)
+  })
+
+  it('assigns fresh row ids when rewriting an upgraded table whose rows lack lineage', async () => {
+    vi.spyOn(Date, 'now').mockReturnValue(1700000000000)
+    const tableUrl = 'mem://rewrite-v3-upgrade'
+    const { resolver } = memResolver()
+    const created = await icebergCreate({ tableUrl, resolver, schema })
+    const records = [{ id: 1n, name: 'a' }, { id: 2n, name: 'b' }, { id: 3n, name: 'c' }]
+    const appended = await icebergStageAppend({ tableUrl, metadata: created, records, resolver })
+    const afterAppend = await fileCatalogCommit({ tableUrl, metadata: created, staged: appended, resolver })
+    /** @type {TableMetadata} */
+    const upgraded = { ...afterAppend, 'format-version': 3, 'next-row-id': 0 }
+
+    // Pre-upgrade rows read with null lineage.
+    const before = await icebergRead({ tableUrl, metadata: upgraded, resolver })
+    expect(before.every(r => r._row_id === null)).toBe(true)
+
+    const staged = await icebergStageRewrite({ tableUrl, metadata: upgraded, resolver })
+    // No carried ids, so the rewritten manifest gets a fresh range at commit.
+    expect(staged.snapshot['first-row-id']).toBe(0)
+    expect(staged.snapshot['added-rows']).toBe(3)
+
+    const after = await fileCatalogCommit({ tableUrl, metadata: upgraded, staged, resolver })
+    expect(after['next-row-id']).toBe(3)
+    const rows = await icebergRead({ tableUrl, metadata: after, resolver })
+    expect(rows.map(r => r._row_id)).toEqual([0n, 1n, 2n])
+  })
+
+  it('detects a concurrent v3 commit via assert-next-row-id', async () => {
+    const { tableUrl, resolver, metadata } = await makeMultiFileTable({ sortOrder: sortById, formatVersion: 3 })
+    const staged = await icebergStageRewrite({ tableUrl, metadata, resolver })
+    // Committing against metadata whose next-row-id moved (a concurrent append
+    // that somehow passed the ref check) must fail rather than risk id reuse.
+    await expect(fileCatalogCommit({
+      tableUrl, metadata: { ...metadata, 'next-row-id': 8 }, staged, resolver,
+    })).rejects.toThrow(/next-row-id expected 6, got 8/)
   })
 })