draft(ipc/writer): try to add compression to file writer

Author: Djjanks

src/ipc/writer.ts

@@ -283,8 +283,8 @@ export class RecordBatchWriter<T extends TypeMap = any> extends ReadableInterop<
             ._writeBodyBuffers(buffers);
     }
 
-    protected _assembleRecordBatch(batch: RecordBatch<T>) {
-        let { byteLength, nodes, bufferRegions, buffers } = VectorAssembler.assemble(batch);
+    protected _assembleRecordBatch<T extends Vector | RecordBatch>(...args: (T | T[])[]) {
+        let { byteLength, nodes, bufferRegions, buffers } = VectorAssembler.assemble(...args);
         if (this._compression != null) {
             ({ byteLength, bufferRegions, buffers } = this._compressBodyBuffers(buffers));
         }
@@ -337,7 +337,7 @@ export class RecordBatchWriter<T extends TypeMap = any> extends ReadableInterop<
     }
 
     protected _writeDictionaryBatch(dictionary: Data, id: number, isDelta = false) {
-        const { byteLength, nodes, bufferRegions, buffers } = VectorAssembler.assemble(new Vector([dictionary]));
+        const { byteLength, nodes, bufferRegions, buffers } = this._assembleRecordBatch(new Vector([dictionary]));
         const recordBatch = new metadata.RecordBatch(dictionary.length, nodes, bufferRegions, this._compression);
         const dictionaryBatch = new metadata.DictionaryBatch(recordBatch, id, isDelta);
         const message = Message.from(dictionaryBatch, byteLength);

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

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+import { RecordBatchStreamWriterOptions } from 'apache-arrow/ipc/writer.js';
 import {
     generateDictionaryTables, generateRandomTables
 } from '../../../data/tables.js';
@@ -23,6 +24,7 @@ import { validateRecordBatchIterator } from '../validate.js';
 
 import {
     builderThroughIterable,
+    CompressionType,
     Dictionary,
     Int32,
     RecordBatch,
@@ -32,6 +34,8 @@ import {
     Uint32,
     Vector
 } from 'apache-arrow';
+import { Codec, compressionRegistry } from 'apache-arrow/ipc/compression/registry';
+import * as lz4js from 'lz4js';
 
 describe('RecordBatchFileWriter', () => {
     for (const table of generateRandomTables([10, 20, 30])) {
@@ -41,6 +45,24 @@ describe('RecordBatchFileWriter', () => {
         testFileWriter(table, `${table.schema.fields[0]}`);
     }
 
+    const compressionTypes = [CompressionType.LZ4_FRAME/*, CompressionType.ZSTD*/];
+
+    const lz4Codec: Codec = {
+        encode(data: Uint8Array): Uint8Array {
+            return lz4js.compress(data);
+        },
+        decode(data: Uint8Array): Uint8Array {
+            return lz4js.decompress(data);
+        }
+    };
+    compressionRegistry.set(CompressionType.LZ4_FRAME, lz4Codec);
+
+    const table = generate.table([10, 20, 30]).table;
+    for (const compressionType of compressionTypes) {
+        const testName = `[${table.schema.fields.join(', ')}] - ${CompressionType[compressionType]} compressed`;
+        testFileWriter(table, testName, { compressionType });
+    }
+
     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();
@@ -91,14 +113,14 @@ describe('RecordBatchFileWriter', () => {
     });
 });
 
-function testFileWriter(table: Table, name: string) {
+function testFileWriter(table: Table, name: string, options?: RecordBatchStreamWriterOptions) {
     describe(`should write the Arrow IPC file format (${name})`, () => {
-        test(`Table`, validateTable.bind(0, table));
+        test(`Table`, validateTable.bind(0, table, options));
     });
 }
 
-async function validateTable(source: Table) {
-    const writer = RecordBatchFileWriter.writeAll(source);
+async function validateTable(source: Table, options?: RecordBatchStreamWriterOptions) {
+    const writer = RecordBatchFileWriter.writeAll(source, options);
     const result = new Table(RecordBatchReader.from(await writer.toUint8Array()));
     validateRecordBatchIterator(3, source.batches);
     expect(result).toEqualTable(source);