feat: Implement IPC RecordBatch body buffer compression

package.json

@@ -71,6 +71,7 @@
     "@types/benchmark": "2.1.5",
     "@types/glob": "8.1.0",
     "@types/jest": "29.5.14",
+    "@types/lz4js": "0.2.1",
     "@types/multistream": "4.1.3",
     "async-done": "2.0.0",
     "benny": "3.7.1",
@@ -96,6 +97,7 @@
     "ix": "7.0.0",
     "jest": "29.7.0",
     "jest-silent-reporter": "0.6.0",
+    "lz4js": "0.2.0",
     "memfs": "4.17.2",
     "mkdirp": "3.0.1",
     "multistream": "4.1.0",
@@ -114,7 +116,8 @@
     "webpack": "5.99.9",
     "webpack-bundle-analyzer": "4.10.2",
     "webpack-stream": "7.0.0",
-    "xml2js": "0.6.2"
+    "xml2js": "0.6.2",
+    "zstd-codec": "0.1.5"
   },
   "engines": {
     "node": ">=12.0"

src/Arrow.dom.ts

@@ -76,6 +76,7 @@ export {
     RecordBatch,
     util,
     Builder, makeBuilder, builderThroughIterable, builderThroughAsyncIterable,
+    compressionRegistry, CompressionType,
 } from './Arrow.js';
 
 export {

src/Arrow.ts

@@ -16,6 +16,7 @@
 // under the License.
 
 export { MessageHeader } from './fb/message-header.js';
+export { CompressionType } from './fb/compression-type.js';
 
 export {
     Type,
@@ -92,6 +93,7 @@ export type { ReadableSource, WritableSink } from './io/stream.js';
 export { RecordBatchReader, RecordBatchFileReader, RecordBatchStreamReader, AsyncRecordBatchFileReader, AsyncRecordBatchStreamReader } from './ipc/reader.js';
 export { RecordBatchWriter, RecordBatchFileWriter, RecordBatchStreamWriter, RecordBatchJSONWriter } from './ipc/writer.js';
 export { tableToIPC, tableFromIPC } from './ipc/serialization.js';
+export { compressionRegistry } from './ipc/compression/registry.js';
 export { MessageReader, AsyncMessageReader, JSONMessageReader } from './ipc/message.js';
 export { Message } from './ipc/metadata/message.js';
 export { RecordBatch } from './recordbatch.js';

src/ipc/compression/constants.ts

@@ -0,0 +1,19 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+export const LENGTH_NO_COMPRESSED_DATA = -1;
+export const COMPRESS_LENGTH_PREFIX = 8;

src/ipc/compression/registry.ts

@@ -0,0 +1,46 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+import { CompressionType } from '../../fb/compression-type.js';
+import { compressionValidators } from './validators.js';
+
+export interface Codec {
+    encode?(data: Uint8Array): Uint8Array;
+    decode?(data: Uint8Array): Uint8Array;
+}
+
+class _CompressionRegistry {
+    protected declare registry: { [key in CompressionType]?: Codec };
+
+    constructor() {
+        this.registry = {};
+    }
+
+    set(compression: CompressionType, codec: Codec) {
+        if (codec?.encode && typeof codec.encode === 'function' && !compressionValidators[compression].isValidCodecEncode(codec)) {
+            throw new Error(`Encoder for ${CompressionType[compression]} is not valid.`);
+        }
+        this.registry[compression] = codec;
+    }
+
+    get(compression: CompressionType): Codec | null {
+        return this.registry?.[compression] || null;
+    }
+
+}
+
+export const compressionRegistry = new _CompressionRegistry();

src/ipc/compression/validators.ts

@@ -0,0 +1,92 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+import type { Codec } from './registry.ts';
+import { CompressionType } from '../../fb/compression-type.js';
+
+export interface CompressionValidator {
+    isValidCodecEncode(codec: Codec): boolean;
+}
+
+class Lz4FrameValidator implements CompressionValidator {
+    private readonly LZ4_FRAME_MAGIC = new Uint8Array([4, 34, 77, 24]);
+    private readonly MIN_HEADER_LENGTH = 7; // 4 (magic) + 2 (FLG + BD) + 1 (header checksum) = 7 min bytes
+
+    isValidCodecEncode(codec: Codec): boolean {
+        const testData = new Uint8Array([1, 2, 3, 4, 5, 6, 7, 8]);
+        const compressed = codec.encode!(testData);
+        return this._isValidCompressed(compressed);
+    }
+
+    private _isValidCompressed(buffer: Uint8Array): boolean {
+        return (
+            this._hasMinimumLength(buffer) &&
+            this._hasValidMagicNumber(buffer) &&
+            this._hasValidVersion(buffer)
+        );
+    }
+
+    private _hasMinimumLength(buffer: Uint8Array): boolean {
+        return buffer.length >= this.MIN_HEADER_LENGTH;
+    }
+
+    private _hasValidMagicNumber(buffer: Uint8Array): boolean {
+        return this.LZ4_FRAME_MAGIC.every(
+            (byte, i) => buffer[i] === byte
+        );
+    }
+
+    private _hasValidVersion(buffer: Uint8Array): boolean {
+        const flg = buffer[4];
+        const versionBits = (flg & 0xC0) >> 6;
+        return versionBits === 1;
+    }
+
+}
+
+class ZstdValidator implements CompressionValidator {
+    private readonly ZSTD_MAGIC = new Uint8Array([40, 181, 47, 253]);
+    private readonly MIN_HEADER_LENGTH = 6; // 4 (magic) + 2 (min Frame_Header) = 6 min bytes
+
+    isValidCodecEncode(codec: Codec): boolean {
+        const testData = new Uint8Array([1, 2, 3, 4, 5, 6, 7, 8]);
+        const compressed = codec.encode!(testData);
+        return this._isValidCompressed(compressed);
+    }
+
+    private _isValidCompressed(buffer: Uint8Array): boolean {
+        return (
+            this._hasMinimumLength(buffer) &&
+            this._hasValidMagicNumber(buffer)
+        );
+    }
+
+    private _hasMinimumLength(buffer: Uint8Array): boolean {
+        return buffer.length >= this.MIN_HEADER_LENGTH;
+    }
+
+    private _hasValidMagicNumber(buffer: Uint8Array): boolean {
+        return this.ZSTD_MAGIC.every(
+            (byte, i) => buffer[i] === byte
+        );
+    }
+}
+
+export const compressionValidators: Record<CompressionType, CompressionValidator> = {
+    [CompressionType.LZ4_FRAME]: new Lz4FrameValidator(),
+    [CompressionType.ZSTD]: new ZstdValidator(),
+};

src/ipc/metadata/json.ts

@@ -40,7 +40,8 @@ export function recordBatchFromJSON(b: any) {
     return new RecordBatch(
         b['count'],
         fieldNodesFromJSON(b['columns']),
-        buffersFromJSON(b['columns'])
+        buffersFromJSON(b['columns']),
+        null
     );
 }
 

src/ipc/metadata/message.ts

@@ -40,6 +40,9 @@ import { FixedSizeBinary as _FixedSizeBinary } from '../../fb/fixed-size-binary.
 import { FixedSizeList as _FixedSizeList } from '../../fb/fixed-size-list.js';
 import { Map as _Map } from '../../fb/map.js';
 import { Message as _Message } from '../../fb/message.js';
+import { CompressionType as _CompressionType } from '../../fb/compression-type.js';
+import { BodyCompression as _BodyCompression } from '../../fb/body-compression.js';
+import { BodyCompressionMethod as _BodyCompressionMethod } from '../../fb/body-compression-method.js';
 
 import { Schema, Field } from '../../schema.js';
 import { toUint8Array } from '../../util/buffer.js';
@@ -122,9 +125,11 @@ export class Message<T extends MessageHeader = any> {
     protected _headerType: T;
     protected _bodyLength: number;
     protected _version: MetadataVersion;
+    protected _compression: BodyCompression | null;
     public get type() { return this.headerType; }
     public get version() { return this._version; }
     public get headerType() { return this._headerType; }
+    public get compression() { return this._compression; }
     public get bodyLength() { return this._bodyLength; }
     declare protected _createHeader: MessageHeaderDecoder;
     public header() { return this._createHeader<T>(); }
@@ -136,6 +141,7 @@ export class Message<T extends MessageHeader = any> {
         this._version = version;
         this._headerType = headerType;
         this.body = new Uint8Array(0);
+        this._compression = header?.compression;
         header && (this._createHeader = () => header);
         this._bodyLength = bigIntToNumber(bodyLength);
     }
@@ -149,13 +155,21 @@ export class RecordBatch {
     protected _length: number;
     protected _nodes: FieldNode[];
     protected _buffers: BufferRegion[];
+    protected _compression: BodyCompression | null;
     public get nodes() { return this._nodes; }
     public get length() { return this._length; }
     public get buffers() { return this._buffers; }
-    constructor(length: bigint | number, nodes: FieldNode[], buffers: BufferRegion[]) {
+    public get compression() { return this._compression; }
+    constructor(
+        length: bigint | number,
+        nodes: FieldNode[],
+        buffers: BufferRegion[],
+        compression: BodyCompression | null
+    ) {
         this._nodes = nodes;
         this._buffers = buffers;
         this._length = bigIntToNumber(length);
+        this._compression = compression;
     }
 }
 
@@ -208,6 +222,19 @@ export class FieldNode {
     }
 }
 
+/**
+ * @ignore
+ * @private
+ **/
+export class BodyCompression {
+    public type: _CompressionType;
+    public method: _BodyCompressionMethod;
+    constructor(type: _CompressionType, method: _BodyCompressionMethod = _BodyCompressionMethod.BUFFER) {
+        this.type = type;
+        this.method = method;
+    }
+}
+
 /** @ignore */
 function messageHeaderFromJSON(message: any, type: MessageHeader) {
     return (() => {
@@ -254,6 +281,9 @@ FieldNode['decode'] = decodeFieldNode;
 BufferRegion['encode'] = encodeBufferRegion;
 BufferRegion['decode'] = decodeBufferRegion;
 
+BodyCompression['encode'] = encodeBodyCompression;
+BodyCompression['decode'] = decodeBodyCompression;
+
 declare module '../../schema' {
     namespace Field {
         export { encodeField as encode };
@@ -286,6 +316,10 @@ declare module './message' {
         export { encodeBufferRegion as encode };
         export { decodeBufferRegion as decode };
     }
+    namespace BodyCompression {
+        export { encodeBodyCompression as encode };
+        export { decodeBodyCompression as decode };
+    }
 }
 
 /** @ignore */
@@ -296,10 +330,13 @@ function decodeSchema(_schema: _Schema, dictionaries: Map<number, DataType> = ne
 
 /** @ignore */
 function decodeRecordBatch(batch: _RecordBatch, version = MetadataVersion.V5) {
-    if (batch.compression() !== null) {
-        throw new Error('Record batch compression not implemented');
-    }
-    return new RecordBatch(batch.length(), decodeFieldNodes(batch), decodeBuffers(batch, version));
+    const recordBatch = new RecordBatch(
+        batch.length(),
+        decodeFieldNodes(batch),
+        decodeBuffers(batch, version),
+        decodeBodyCompression(batch.compression())
+    );
+    return recordBatch;
 }
 
 /** @ignore */
@@ -491,6 +528,11 @@ function decodeFieldType(f: _Field, children?: Field[]): DataType<any> {
     throw new Error(`Unrecognized type: "${Type[typeId]}" (${typeId})`);
 }
 
+/** @ignore */
+function decodeBodyCompression(b: _BodyCompression | null) {
+    return b ? new BodyCompression(b.codec(), b.method()) : null;
+}
+
 /** @ignore */
 function encodeSchema(b: Builder, schema: Schema) {
 
@@ -583,13 +625,29 @@ function encodeRecordBatch(b: Builder, recordBatch: RecordBatch) {
 
     const buffersVectorOffset = b.endVector();
 
+    let bodyCompressionOffset = null;
+    if (recordBatch.compression !== null) {
+        bodyCompressionOffset = encodeBodyCompression(b, recordBatch.compression);
+    }
+
     _RecordBatch.startRecordBatch(b);
     _RecordBatch.addLength(b, BigInt(recordBatch.length));
     _RecordBatch.addNodes(b, nodesVectorOffset);
     _RecordBatch.addBuffers(b, buffersVectorOffset);
+    if (recordBatch.compression !== null && bodyCompressionOffset) {
+        _RecordBatch.addCompression(b, bodyCompressionOffset);
+    }
     return _RecordBatch.endRecordBatch(b);
 }
 
+/** @ignore */
+function encodeBodyCompression(b: Builder, node: BodyCompression) {
+    _BodyCompression.startBodyCompression(b);
+    _BodyCompression.addCodec(b, node.type);
+    _BodyCompression.addMethod(b, node.method);
+    return _BodyCompression.endBodyCompression(b);
+}
+
 /** @ignore */
 function encodeDictionaryBatch(b: Builder, dictionaryBatch: DictionaryBatch) {
     const dataOffset = RecordBatch.encode(b, dictionaryBatch.data);