apacheGH-23576: [JS] Support DictionaryMessage replacement (apache#41965

)

This PR makes the stream reader/writer support replacement dictionary
messages.

Closes apache#23576
Closes apache#41683

* GitHub Issue: apache#23576

js/gulp/argv.js

@@ -22,6 +22,7 @@ export const argv = args([
  { name: `target`, type: String, defaultValue: `` },
  { name: `module`, type: String, defaultValue: `` },
  { name: `coverage`, type: Boolean, defaultValue: false },
+ { name: `tests`, type: String, multiple: true, defaultValue: [`test/unit/`] },
  { name: `targets`, alias: `t`, type: String, multiple: true, defaultValue: [] },
  { name: `modules`, alias: `m`, type: String, multiple: true, defaultValue: [] },
 ], { partial: true });

js/gulp/test-task.js

@@ -36,19 +36,6 @@ import { createRequire } from 'node:module';
 const require = createRequire(import.meta.url);
 
 const jestArgv = [];
-const testFiles = [
- `test/unit/`,
- // `test/unit/bit-tests.ts`,
- // `test/unit/int-tests.ts`,
- // `test/unit/bn-tests.ts`,
- // `test/unit/math-tests.ts`,
- // `test/unit/table-tests.ts`,
- // `test/unit/generated-data-tests.ts`,
- // `test/unit/builders/`,
- // `test/unit/recordbatch/`,
- // `test/unit/table/`,
- // `test/unit/ipc/`,
-];
 
 if (argv.verbose) {
  jestArgv.push(`--verbose`);
@@ -80,8 +67,10 @@ export const testTask = ((cache, execArgv, testOptions) => memoizeTask(cache, fu
  args.push(`-c`, `jestconfigs/jest.coverage.config.js`);
  } else {
  const cfgname = [target, format].filter(Boolean).join('.');
- args.push(`-c`, `jestconfigs/jest.${cfgname}.config.js`, ...testFiles);
+ args.push(`-c`, `jestconfigs/jest.${cfgname}.config.js`);
  }
+ args.push(...(argv._unknown || []).filter((x) => x !== 'test'));
+ args.push(...argv.tests);
  opts.env = {
  ...opts.env,
  TEST_TARGET: target,

js/src/ipc/reader.ts

@@ -363,14 +363,11 @@ abstract class RecordBatchReaderImpl<T extends TypeMap = any> implements RecordB
  const { id, isDelta } = header;
  const { dictionaries, schema } = this;
  const dictionary = dictionaries.get(id);
- if (isDelta || !dictionary) {
- const type = schema.dictionaries.get(id)!;
- const data = this._loadVectors(header.data, body, [type]);
- return (dictionary && isDelta ? dictionary.concat(
- new Vector(data)) :
- new Vector(data)).memoize() as Vector;
- }
- return dictionary.memoize();
+ const type = schema.dictionaries.get(id)!;
+ const data = this._loadVectors(header.data, body, [type]);
+ return (dictionary && isDelta ? dictionary.concat(
+ new Vector(data)) :
+ new Vector(data)).memoize() as Vector;
  }
  protected _loadVectors(header: metadata.RecordBatch, body: any, types: (Field | DataType)[]) {
  return new VectorLoader(body, header.nodes, header.buffers, this.dictionaries, this.schema.metadataVersion).visitMany(types);

js/src/ipc/writer.ts

@@ -84,6 +84,7 @@ export class RecordBatchWriter<T extends TypeMap = any> extends ReadableInterop<
  protected _schema: Schema | null = null;
  protected _dictionaryBlocks: FileBlock[] = [];
  protected _recordBatchBlocks: FileBlock[] = [];
+ protected _seenDictionaries = new Map<number, Vector>();
  protected _dictionaryDeltaOffsets = new Map<number, number>();
 
  public toString(sync: true): string;
@@ -144,6 +145,7 @@ export class RecordBatchWriter<T extends TypeMap = any> extends ReadableInterop<
  this._started = false;
  this._dictionaryBlocks = [];
  this._recordBatchBlocks = [];
+ this._seenDictionaries = new Map();
  this._dictionaryDeltaOffsets = new Map();
 
  if (!schema || !(compareSchemas(schema, this._schema))) {
@@ -259,7 +261,6 @@ export class RecordBatchWriter<T extends TypeMap = any> extends ReadableInterop<
  }
 
  protected _writeDictionaryBatch(dictionary: Data, id: number, isDelta = false) {
- this._dictionaryDeltaOffsets.set(id, dictionary.length + (this._dictionaryDeltaOffsets.get(id) || 0));
  const { byteLength, nodes, bufferRegions, buffers } = VectorAssembler.assemble(new Vector([dictionary]));
  const recordBatch = new metadata.RecordBatch(dictionary.length, nodes, bufferRegions);
  const dictionaryBatch = new metadata.DictionaryBatch(recordBatch, id, isDelta);
@@ -284,14 +285,21 @@ export class RecordBatchWriter<T extends TypeMap = any> extends ReadableInterop<
  }
 
  protected _writeDictionaries(batch: RecordBatch<T>) {
- for (let [id, dictionary] of batch.dictionaries) {
- let offset = this._dictionaryDeltaOffsets.get(id) || 0;
- if (offset === 0 || (dictionary = dictionary?.slice(offset)).length > 0) {
- for (const data of dictionary.data) {
- this._writeDictionaryBatch(data, id, offset > 0);
- offset += data.length;
- }
+ for (const [id, dictionary] of batch.dictionaries) {
+ const chunks = dictionary?.data ?? [];
+ const prevDictionary = this._seenDictionaries.get(id);
+ const offset = this._dictionaryDeltaOffsets.get(id) ?? 0;
+ // * If no previous dictionary was written, write an initial DictionaryMessage.
+ // * If the current dictionary does not share chunks with the previous dictionary, write a replacement DictionaryMessage.
+ if (!prevDictionary || prevDictionary.data[0] !== chunks[0]) {
+ // * If `index > 0`, then `isDelta` is true.
+ // * If `index = 0`, then `isDelta` is false, because this is either the initial or a replacement DictionaryMessage.
+ for (const [index, chunk] of chunks.entries()) this._writeDictionaryBatch(chunk, id, index > 0);
+ } else if (offset < chunks.length) {
+ for (const chunk of chunks.slice(offset)) this._writeDictionaryBatch(chunk, id, true);
  }
+ this._seenDictionaries.set(id, dictionary);
+ this._dictionaryDeltaOffsets.set(id, chunks.length);
  }
  return this;
  }
@@ -342,6 +350,13 @@ export class RecordBatchFileWriter<T extends TypeMap = any> extends RecordBatchW
  return this._writeMagic()._writePadding(2);
  }
 
+ protected _writeDictionaryBatch(dictionary: Data, id: number, isDelta = false) {
+ if (!isDelta && this._seenDictionaries.has(id)) {
+ throw new Error('The Arrow File format does not support replacement dictionaries. ');
+ }
+ return super._writeDictionaryBatch(dictionary, id, isDelta);
+ }
+
  protected _writeFooter(schema: Schema<T>) {
  const buffer = Footer.encode(new Footer(
  schema, MetadataVersion.V5,
@@ -369,13 +384,13 @@ export class RecordBatchJSONWriter<T extends TypeMap = any> extends RecordBatchW
  }
 
  private _recordBatches: RecordBatch[];
- private _dictionaries: RecordBatch[];
+ private _recordBatchesWithDictionaries: RecordBatch[];
 
  constructor() {
  super();
  this._autoDestroy = true;
  this._recordBatches = [];
- this._dictionaries = [];
+ this._recordBatchesWithDictionaries = [];
  }
 
  protected _writeMessage() { return this; }
@@ -386,12 +401,11 @@ export class RecordBatchJSONWriter<T extends TypeMap = any> extends RecordBatchW
  }
  protected _writeDictionaries(batch: RecordBatch<T>) {
  if (batch.dictionaries.size > 0) {
- this._dictionaries.push(batch);
+ this._recordBatchesWithDictionaries.push(batch);
  }
  return this;
  }
  protected _writeDictionaryBatch(dictionary: Data, id: number, isDelta = false) {
- this._dictionaryDeltaOffsets.set(id, dictionary.length + (this._dictionaryDeltaOffsets.get(id) || 0));
  this._write(this._dictionaryBlocks.length === 0 ? ` ` : `,\n `);
  this._write(dictionaryBatchToJSON(dictionary, id, isDelta));
  this._dictionaryBlocks.push(new FileBlock(0, 0, 0));
@@ -403,9 +417,9 @@ export class RecordBatchJSONWriter<T extends TypeMap = any> extends RecordBatchW
  return this;
  }
  public close() {
- if (this._dictionaries.length > 0) {
+ if (this._recordBatchesWithDictionaries.length > 0) {
  this._write(`,\n "dictionaries": [\n`);
- for (const batch of this._dictionaries) {
+ for (const batch of this._recordBatchesWithDictionaries) {
  super._writeDictionaries(batch);
  }
  this._write(`\n ]`);
@@ -424,7 +438,7 @@ export class RecordBatchJSONWriter<T extends TypeMap = any> extends RecordBatchW
  this._write(`\n}`);
  }
 
- this._dictionaries = [];
+ this._recordBatchesWithDictionaries = [];
  this._recordBatches = [];
 
  return super.close();

js/test/unit/ipc/message-reader-tests.ts

@@ -24,11 +24,16 @@ for (const table of generateRandomTables([10, 20, 30])) {
 
  const io = ArrowIOTestHelper.stream(table);
  const name = `[\n ${table.schema.fields.join(',\n ')}\n]`;
- const numMessages = table.batches.reduce((numMessages, batch) => {
- return numMessages +
- /* recordBatch message */ 1 +
- /* dictionary messages */ batch.dictionaries.size;
- }, /* schema message */ 1);
+
+ const numDictionaries = table.batches.reduce((dictionaries, batch) => {
+ return [...batch.dictionaries.values()]
+ .flatMap((dictionary) => dictionary.data)
+ .reduce((dictionaries, data) => dictionaries.add(data), dictionaries);
+ }, new Set()).size;
+
+ const numMessages = /* schema message */ 1 +
+ /* recordBatch messages */ table.batches.length +
+ /* dictionary messages */ numDictionaries;
 
  const validate = validateMessageReader.bind(0, numMessages);
  const validateAsync = validateAsyncMessageReader.bind(0, numMessages);

js/test/unit/ipc/writer/file-writer-tests.ts

@@ -18,9 +18,20 @@
 import {
  generateDictionaryTables, generateRandomTables
 } from '../../../data/tables.js';
+import * as generate from '../../../generate-test-data.js';
 import { validateRecordBatchIterator } from '../validate.js';
 
-import { RecordBatchFileWriter, RecordBatchReader, Table } from 'apache-arrow';
+import {
+ builderThroughIterable,
+ Dictionary,
+ Int32,
+ RecordBatch,
+ RecordBatchFileWriter,
+ RecordBatchReader,
+ Table,
+ Uint32,
+ Vector
+} from 'apache-arrow';
 
 describe('RecordBatchFileWriter', () => {
  for (const table of generateRandomTables([10, 20, 30])) {
@@ -29,6 +40,55 @@ describe('RecordBatchFileWriter', () => {
  for (const table of generateDictionaryTables([10, 20, 30])) {
  testFileWriter(table, `${table.schema.fields[0]}`);
  }
+
+ it('should throw if attempting to write replacement dictionary batches', async () => {
+ const type = new Dictionary<Uint32, Int32>(new Uint32, new Int32, 0);
+ const writer = new RecordBatchFileWriter();
+ writer.write(new RecordBatch({
+ // Clone the data with the original Dictionary type so the cloned chunk has id 0
+ dictionary_encoded_uint32: generate.dictionary(50, 20, new Uint32, new Int32).vector.data[0].clone(type)
+ }));
+ expect(() => {
+ writer.write(new RecordBatch({
+ // Clone the data with the original Dictionary type so the cloned chunk has id 0
+ dictionary_encoded_uint32: generate.dictionary(50, 20, new Uint32, new Int32).vector.data[0].clone(type)
+ }));
+ }).toThrow();
+ });
+
+ it('should write delta dictionary batches', async () => {
+
+ const name = 'dictionary_encoded_uint32';
+ const resultChunks: Vector<Dictionary<Uint32, Int32>>[] = [];
+ const {
+ vector: sourceVector, values: sourceValues,
+ } = generate.dictionary(1000, 20, new Uint32(), new Int32());
+
+ const writer = RecordBatchFileWriter.writeAll((function* () {
+ const transform = builderThroughIterable({
+ type: sourceVector.type, nullValues: [null],
+ queueingStrategy: 'count', highWaterMark: 50,
+ });
+ for (const chunk of transform(sourceValues())) {
+ resultChunks.push(chunk);
+ yield new RecordBatch({ [name]: chunk.data[0] });
+ }
+ })());
+
+ expect(new Vector(resultChunks)).toEqualVector(sourceVector);
+
+ type T = { [name]: Dictionary<Uint32, Int32> };
+ const sourceTable = new Table({ [name]: sourceVector });
+ const resultTable = new Table(RecordBatchReader.from<T>(await writer.toUint8Array()));
+
+ const child = resultTable.getChild(name)!;
+ const dicts = child.data.map(({ dictionary }) => dictionary!);
+ const dictionary = dicts[child.data.length - 1];
+
+ expect(resultTable).toEqualTable(sourceTable);
+ expect(dictionary).toBeInstanceOf(Vector);
+ expect(dictionary.data).toHaveLength(20);
+ });
 });
 
 function testFileWriter(table: Table, name: string) {

js/test/unit/ipc/writer/stream-writer-tests.ts

@@ -25,6 +25,7 @@ import { validateRecordBatchIterator } from '../validate.js';
 import type { RecordBatchStreamWriterOptions } from 'apache-arrow/ipc/writer';
 import {
  builderThroughIterable,
+ Data,
  Dictionary,
  Field,
  Int32,
@@ -81,10 +82,46 @@ describe('RecordBatchStreamWriter', () => {
  await validate;
  });
 
+ it('should write replacement dictionary batches', async () => {
+
+ const name = 'dictionary_encoded_uint32';
+ const type = new Dictionary<Uint32, Int32>(new Uint32, new Int32, 0);
+ const sourceChunks: Data<Dictionary<Uint32, Int32>>[] = [];
+ const resultChunks: Data<Dictionary<Uint32, Int32>>[] = [];
+
+ const writer = RecordBatchStreamWriter.writeAll((function* () {
+ for (let i = 0; i < 1000; i += 50) {
+ const { vector: { data: [chunk] } } = generate.dictionary(50, 20, type.dictionary, type.indices);
+ sourceChunks.push(chunk);
+ // Clone the data with the original Dictionary type so the cloned chunk has id 0
+ resultChunks.push(chunk.clone(type));
+ yield new RecordBatch({ [name]: resultChunks.at(-1)! });
+ }
+ })());
+
+ expect(new Vector(resultChunks)).toEqualVector(new Vector(sourceChunks));
+
+ type T = { [name]: Dictionary<Uint32, Int32> };
+ const sourceTable = new Table({ [name]: new Vector(sourceChunks) });
+ const resultTable = new Table(RecordBatchReader.from<T>(await writer.toUint8Array()));
+
+ // 1000 / 50 = 20
+ expect(sourceTable.batches).toHaveLength(20);
+ expect(resultTable.batches).toHaveLength(20);
+ expect(resultTable).toEqualTable(sourceTable);
+
+ for (const batch of resultTable.batches) {
+ for (const [_, dictionary] of batch.dictionaries) {
+ expect(dictionary).toBeInstanceOf(Vector);
+ expect(dictionary.data).toHaveLength(1);
+ }
+ }
+ });
+
  it('should write delta dictionary batches', async () => {
 
  const name = 'dictionary_encoded_uint32';
- const chunks: Vector<Dictionary<Uint32, Int32>>[] = [];
+ const resultChunks: Vector<Dictionary<Uint32, Int32>>[] = [];
  const {
  vector: sourceVector, values: sourceValues,
  } = generate.dictionary(1000, 20, new Uint32(), new Int32());
@@ -95,12 +132,12 @@ describe('RecordBatchStreamWriter', () => {
  queueingStrategy: 'count', highWaterMark: 50,
  });
  for (const chunk of transform(sourceValues())) {
- chunks.push(chunk);
+ resultChunks.push(chunk);
  yield new RecordBatch({ [name]: chunk.data[0] });
  }
  })());
 
- expect(new Vector(chunks)).toEqualVector(sourceVector);
+ expect(new Vector(resultChunks)).toEqualVector(sourceVector);
 
  type T = { [name]: Dictionary<Uint32, Int32> };
  const sourceTable = new Table({ [name]: sourceVector });