diff --git a/src/decode/decode-ipc.js b/src/decode/decode-ipc.js index 9c34a9f..098f96f 100644 --- a/src/decode/decode-ipc.js +++ b/src/decode/decode-ipc.js @@ -59,6 +59,13 @@ export function decodeIPCStream(data) { let schema; const records = []; const dictionaries = []; + // Parallel to `records`: for each record batch, how many dictionary batches + // were already in `dictionaries` when the record batch arrived. Lets + // `createTable` walk dict and record batches in their original stream + // order, which is required to honor non-delta dictionary REPLACEMENT + // messages (each replacement only applies to record batches that come + // AFTER it in the stream). + const dictsBeforeRecord = []; // consume each message in the stream for (const buf of stream) { @@ -80,6 +87,7 @@ export function decodeIPCStream(data) { break; case MessageHeader.RecordBatch: records.push(m.content); + dictsBeforeRecord.push(dictionaries.length); break; case MessageHeader.DictionaryBatch: dictionaries.push(m.content); @@ -89,7 +97,7 @@ export function decodeIPCStream(data) { } return /** @type {ArrowData} */ ( - { schema, dictionaries, records, metadata: null } + { schema, dictionaries, records, dictsBeforeRecord, metadata: null } ); } @@ -117,10 +125,20 @@ export function decodeIPCFile(data) { const dicts = get(8, decodeBlocks, []); const recs = get(10, decodeBlocks, []); + // Compute, for each record block, the number of dictionary blocks that + // appear BEFORE it in the file by byte offset. This preserves the + // original stream order for files that interleave dictionary and record + // blocks (and degrades to "all dicts before all records" when the file + // follows the conventional layout). + const dictsBeforeRecord = recs.map( + rec => dicts.filter(d => d.offset < rec.offset).length + ); + return /** @type {ArrowData} */ ({ schema: get(6, (buf, index) => decodeSchema(buf, index, version)), dictionaries: dicts.map(({ offset }) => decodeMessage(data, offset).content), records: recs.map(({ offset }) => decodeMessage(data, offset).content), + dictsBeforeRecord, metadata: get(12, decodeMetadata) }); } diff --git a/src/decode/table-from-ipc.js b/src/decode/table-from-ipc.js index a8c267a..376b2bb 100644 --- a/src/decode/table-from-ipc.js +++ b/src/decode/table-from-ipc.js @@ -2,7 +2,7 @@ * @import { ArrowData, BodyCompression, ExtractionOptions, Field, RecordBatch, Schema } from '../types.js' */ import { batchType } from '../batch-type.js'; -import { columnBuilder } from '../column.js'; +import { Column, columnBuilder } from '../column.js'; import { decompressBuffer, getCompressionCodec, missingCodec } from '../compression.js'; import { BodyCompressionMethod, Type, UnionMode, Version } from '../constants.js'; import { invalidDataType } from '../data-types.js'; @@ -40,8 +40,15 @@ export function tableFromIPC(data, options) { * @returns {Table} A Table instance. */ export function createTable(data, options = {}) { - const { schema = { fields: [] }, dictionaries, records } = data; + const { schema = { fields: [] }, dictionaries, records, dictsBeforeRecord } = data; const { version, fields } = schema; + // `dictionaryMap` is closed over by `context()` and read by `visit()` when + // it encounters a Dictionary-typed field on a record batch. Each record + // batch captures whatever Column reference is current in `dictionaryMap` + // at the time it is decoded — exactly mirroring how arrow-js's + // `_loadDictionaryBatch` calls `this.dictionaries.set(header.id, vector)` + // and how `_loadRecordBatch` then resolves dict references against the + // current `this.dictionaries` Map. const dictionaryMap = new Map; const context = contextGenerator(options, version, dictionaryMap); @@ -54,30 +61,61 @@ export function createTable(data, options = {}) { } }); - // decode dictionaries, build dictionary column map - const dicts = new Map; - for (const dict of dictionaries) { - const { id, data, isDelta, body } = dict; + /** + * Process a single dictionary batch by either replacing or extending the + * Column currently registered for its id in `dictionaryMap`. A fresh + * Column instance is always created (never mutated in place) so that any + * record batch that previously captured the old Column reference is + * unaffected. + * + * Mirrors arrow-js's `_loadDictionaryBatch`: + * return (dictionary && isDelta + * ? dictionary.concat(new Vector(data)) + * : new Vector(data)).memoize(); + * @param {import('../types.js').DictionaryBatch} dict + */ + const processDict = (dict) => { + const { id, data: dictData, isDelta, body } = dict; const type = dictionaryTypes.get(id); - const batch = visit(type, context({ ...data, body })); - if (!dicts.has(id)) { + const batch = visit(type, context({ ...dictData, body })); + const existing = dictionaryMap.get(id); + if (!existing) { if (isDelta) { throw new Error('Delta update can not be first dictionary batch.'); } - dicts.set(id, columnBuilder(type).add(batch)); + dictionaryMap.set(id, new Column([batch], type)); + } else if (isDelta) { + // delta — append to the existing dictionary column + dictionaryMap.set(id, new Column([...existing.data, batch], type)); } else { - const dict = dicts.get(id); - if (!isDelta) dict.clear(); - dict.add(batch); + // non-delta replacement — start fresh + dictionaryMap.set(id, new Column([batch], type)); } - } - dicts.forEach((value, key) => dictionaryMap.set(key, value.done())); + }; - // decode column fields + // Decode dictionary and record batches in their original stream order so + // that each record batch sees the dictionaries that were current at its + // position. `dictsBeforeRecord[i]` is the number of dictionary batches + // that preceded record batch `i` in the stream/file; if it is missing + // (legacy callers building `ArrowData` by hand), fall back to processing + // all dictionary batches before any record batch. const cols = fields.map(f => columnBuilder(f.type)); - for (const batch of records) { - const ctx = context(batch); - fields.forEach((f, i) => cols[i].add(visit(f.type, ctx))); + let dictIdx = 0; + for (let i = 0; i < records.length; i++) { + const target = dictsBeforeRecord ? dictsBeforeRecord[i] : dictionaries.length; + while (dictIdx < target) { + processDict(dictionaries[dictIdx++]); + } + + const ctx = context(records[i]); + fields.forEach((f, idx) => cols[idx].add(visit(f.type, ctx))); + } + // Drain any dictionary batches that come after the last record batch in + // the stream. These can't affect any already-decoded record batch, but + // processing them surfaces any malformed-dict errors that the caller + // would expect to see. + while (dictIdx < dictionaries.length) { + processDict(dictionaries[dictIdx++]); } return new Table(schema, cols.map(c => c.done()), options.useProxy); diff --git a/src/types.ts b/src/types.ts index 0ed5d4f..a8fb6b2 100644 --- a/src/types.ts +++ b/src/types.ts @@ -306,6 +306,16 @@ export interface ArrowData { schema?: Schema; dictionaries: DictionaryBatch[] | null; records: RecordBatch[] | null; + /** + * For each record batch (parallel to `records`), the number of dictionary + * batches that preceded it in the original IPC stream / file. Used by + * `createTable` to honor per-record-batch dictionary state when a stream + * carries dictionary REPLACEMENT messages (non-delta dictionary batches + * that overwrite a previously-transmitted dictionary for the same id). + * If null, all dictionary batches are treated as preceding all record + * batches (the legacy behavior). + */ + dictsBeforeRecord?: number[] | null; metadata: Metadata | null; } diff --git a/test/data/dict_replacement.arrows b/test/data/dict_replacement.arrows new file mode 100644 index 0000000..33f9e1a Binary files /dev/null and b/test/data/dict_replacement.arrows differ diff --git a/test/dict-replacement.test.js b/test/dict-replacement.test.js new file mode 100644 index 0000000..3d3c462 --- /dev/null +++ b/test/dict-replacement.test.js @@ -0,0 +1,73 @@ +// Repro for a dictionary-replacement decoding bug in `tableFromIPC`. +// +// In the Arrow IPC stream format, a single dictionary id can be associated +// with multiple non-delta dictionary batches: each subsequent non-delta +// batch REPLACES the previous one for the purposes of any record batches +// that come AFTER it in the stream. Record batches always reference the +// dictionary that is current at their position in the stream. +// +// `decode/table-from-ipc.js`'s `createTable` decodes all dictionary batches +// in one pass before any record batches, overwriting the dictionary stored +// for an id when it sees a non-delta replacement. After the loop, every +// record batch — including ones that were transmitted BEFORE a replacement +// — gets associated with the FINAL dictionary for the id, losing the +// dictionary that was current at their position. +// +// The fixture `test/data/dict_replacement.arrows` is an Arrow IPC stream +// (written by Apache Arrow's Rust `StreamWriter`, which Apache Arrow JS +// reads correctly) with the following logical contents: +// +// schema { id: Int64, name: Dict } +// dict batch (id=0, isDelta=false): values=["alice","bob","carol"] +// record batch 1: id=[1,2,3], name keys=[0,1,2] +// dict batch (id=0, isDelta=false): values=["dave","eve"] +// record batch 2: id=[4,5], name keys=[0,1] +// +// The expected decoded table is: +// +// id | name +// 1 | alice +// 2 | bob +// 3 | carol +// 4 | dave +// 5 | eve +// +// Today's flechette behavior: rows 0..2 are resolved against the LAST +// dictionary `["dave","eve"]` instead of the dictionary that was current at +// the position of record batch 1, so the decoded table is wrong: +// +// id | name +// 1 | dave (key 0 → wrong dict) +// 2 | eve (key 1 → wrong dict) +// 3 | undefined (key 2 → out of bounds in 2-entry dict) +// 4 | dave +// 5 | eve + +import { describe, it, expect } from 'vitest'; +import { readFile } from 'node:fs/promises'; +import { tableFromIPC } from '../src/index.js'; + +describe('tableFromIPC', () => { + it('honors per-record-batch dictionary replacements', async () => { + const bytes = new Uint8Array(await readFile('test/data/dict_replacement.arrows')); + const table = tableFromIPC(bytes); + + const ids = table.getChild('id'); + const names = table.getChild('name'); + + const expected = [ + { id: 1, name: 'alice' }, + { id: 2, name: 'bob' }, + { id: 3, name: 'carol' }, + { id: 4, name: 'dave' }, + { id: 5, name: 'eve' }, + ]; + + const actual = []; + for (let i = 0; i < table.numRows; i++) { + actual.push({ id: ids.get(i), name: names.get(i) }); + } + + expect(actual).toStrictEqual(expected); + }); +}); diff --git a/test/util/decimal.js b/test/util/decimal.js index 07f798f..ff9be2c 100644 --- a/test/util/decimal.js +++ b/test/util/decimal.js @@ -38,5 +38,8 @@ export function decimalDataDecoded() { record.body = record.buffers[0]; delete record.byteLength; delete record.buffers; + // `decodeIPC` populates one entry per record batch indicating how many + // dictionary batches preceded it in the stream/file. + data.dictsBeforeRecord = data.records.map(() => 0); return data; }