Add compression support. (#63)

* feat: Add decompression support for decode.

* feat: Add compressed encoding.

* docs: Update tableFromIPC docs.

* Update src/encode/table-to-ipc.js

Author: jheer

docs/api/index.md

@@ -13,19 +13,22 @@ title: API Reference
 * [columnFromArray](#columnFromArray)
 * [columnFromValues](#columnFromValues)
 * [tableFromColumns](#tableFromColumns)
+* [setCompressionCodec](#setCompressionCodec)
 
 <hr/><a id="tableFromIPC" href="#tableFromIPC">#</a>
 <b>tableFromIPC</b>(<i>data</i>[, <i>options</i>])
 
 Decode [Apache Arrow IPC data](https://arrow.apache.org/docs/format/Columnar.html#serialization-and-interprocess-communication-ipc) and return a new [`Table`](table). The input binary data may be either an `ArrayBuffer` or `Uint8Array`. For Arrow data in the [IPC 'stream' format](https://arrow.apache.org/docs/format/Columnar.html#ipc-streaming-format), an array of `Uint8Array` values is also supported.
 
+By default Flechette assumes input data is uncompressed. If input IPC data contains compressed buffers, an appropriate compression codec (for `CompressionType.LZ4_FRAME` or `CompressionType.ZSTD`) must be registered ahead of time using the [`setCompressionCodec`](#setCompressionCodec) method. Otherwise, an error is thrown.
+
 * *data* (`ArrayBuffer` \| `Uint8Array` \| `Uint8Array[]`): The source byte buffer, or an array of buffers. If an array, each byte array may contain one or more self-contained messages. Messages may NOT span multiple byte arrays.
 * *options* (`ExtractionOptions`): Options for controlling how values are transformed when extracted from an Arrow binary representation.
-  * *useBigInt* (`boolean`): If true, extract 64-bit integers as JavaScript `BigInt` values. Otherwise, coerce long integers to JavaScript number values (default `false`).
+  * *useBigInt* (`boolean`): If true, extract 64-bit integers as JavaScript `BigInt` values. Otherwise, coerce long integers to JavaScript number values (default `false`), raising an error if the integer can not be represented as a double precision floating point number.
   * *useDate* (`boolean`): If true, extract dates and timestamps as JavaScript `Date` objects. Otherwise, return numerical timestamp values (default `false`).
-  * *useDecimalInt* (`boolean`): If true, extract decimal-type data as scaled integer values, where fractional digits are scaled to integer positions. Returned integers are `BigInt` values for decimal bit widths of 64 bits or higher and 32-bit integers (as JavaScript `number`) otherwise. If false, decimals are converted to floating-point numbers (default).
+  * *useDecimalInt* (`boolean`): If true, extract decimal-type data as scaled integer values, where fractional digits are scaled to integer positions. Returned integers are `BigInt` values for decimal bit widths of 64 bits or higher and 32-bit integers (as JavaScript `number`) otherwise. If false, decimals are lossily converted to floating-point numbers (default).
   * *useMap* (`boolean`): If true, extract Arrow 'Map' values as JavaScript `Map` instances. Otherwise, return an array of [key, value] pairs compatible with both `Map` and `Object.fromEntries` (default `false`).
-  * *useProxy* (`boolean`): If true, extract Arrow 'Struct' values and table row objects using zero-copy proxy objects that extract data from underlying Arrow batches. The proxy objects can improve performance and reduce memory usage, but do not support property enumeration (`Object.keys`, `Object.values`, `Object.entries`) or spreading (`{ ...object }`). Otherwise, use standard JS objects for structs and table rows (default `false`).
+  * *useProxy* (`boolean`): If true, extract Arrow 'Struct' values and table row objects using zero-copy proxy objects that extract data from underlying Arrow batches. The proxy objects can improve performance and reduce memory usage, but do not support property enumeration (`Object.keys`, `Object.values`, `Object.entries`) or spreading (`{ ...object }`). Otherwise (default `false`), use standard JS objects for structs and table rows.
 
 ```js
 import { tableFromIPC } from '@uwdata/flechette';
@@ -37,15 +40,50 @@ const table = tableFromIPC(ipc);
 <hr/><a id="tableToIPC" href="#tableToIPC">#</a>
 <b>tableToIPC</b>(<i>table</i>[, <i>options</i>])
 
-Encode an Arrow table into Arrow IPC binary format and return the result as a `Uint8Array`. Both the IPC ['stream'](https://arrow.apache.org/docs/format/Columnar.html#ipc-streaming-format) and ['file'](https://arrow.apache.org/docs/format/Columnar.html#ipc-file-format) formats are supported.
+Encode an Arrow table into Arrow IPC binary format and return the result as a `Uint8Array`. Both the IPC ['stream'](https://arrow.apache.org/docs/format/Columnar.html#ipc-streaming-format) and ['file'](https://arrow.apache.org/docs/format/Columnar.html#ipc-file-format) formats are supported. By default Flechette encodes uncompressed data. To perform compression, register a compression codec plugin using [`setCompressionCodec`](#setCompressionCodec) and then set the *codec* option.
 
 * *table* (`Table`): The Arrow table to encode.
 * *options* (`object`): Encoding options object.
-  * *format* (`string`): Arrow `'stream'` (the default) or `'file'` format.
+  * *format* (`string`): Arrow `'stream'` (the default) or `'file'` (a.k.a. [Feather V2](https://arrow.apache.org/docs/python/feather.html)) format.
+  * *codec* (`number`): The compression codec type to apply. By default no compression is applied. If specified, this option must be one of the `CompressionType` values, and a corresponding codec plugin must already be registered using the [`setCompressionCodec`](#setCompressionCodec) method.
 
 ```js
 import { tableToIPC } from '@uwdata/flechette';
-const bytes = tableFromIPC(table, { format: 'stream' });
+const bytes = tableToIPC(table, { format: 'stream' });
+```
+
+```js
+import { CompressionType, setCompressionCodec, tableToIPC } from '@uwdata/flechette';
+import { ZstdCodec } from 'zstd-codec';
+
+// register zstd compression codec
+await new Promise((resolve) => {
+  ZstdCodec.run((zstd) => {
+    const codec = new zstd.Simple();
+    setCompressionCodec(CompressionType.ZSTD, {
+      encode: (data) => codec.compress(data),
+      decode: (data) => codec.decompress(data)
+    });
+    resolve();
+  });
+});
+
+// generate bytes in IPC stream format, with zstd compression of buffers
+const bytes = tableToIPC(table, { codec: CompressionType.ZSTD });
+```
+
+```js
+import { CompressionType, setCompressionCodec, tableToIPC } from '@uwdata/flechette';
+import * as lz4 from 'lz4js';
+
+// register lz4_frame compression codec
+setCompressionCodec(CompressionType.LZ4_FRAME, {
+  encode: (data) => lz4.compress(data),
+  decode: (data) => lz4.decompress(data)
+});
+
+// generate bytes in IPC stream format, with lz4_frame compression of buffers
+const bytes = tableToIPC(table, { codec: CompressionType.LZ4_FRAME });
 ```
 
 <hr/><a id="tableFromArrays" href="#tableFromArrays">#</a>
@@ -166,3 +204,39 @@ const table = tableFromColumns({
   floats: columnFromArray([1.1, 2.2, 3.3, 4.4, 5.5])
 });
 ```
+
+<hr/><a id="setCompressionCodec" href="#setCompressionCodec">#</a>
+<b>setCompressionCodec</b>(<i>type</i>, <i>codec</i>)
+
+Register a compression codec for compressing or decompressing Arrow bufferdata. Flechette does not include compression codecs by default, but can be extended via plugins to handle compressed data. If an appropriate codec implementation is not registered, an error is thrown when attempting to compress or decompress data.
+
+* *type* (`number`): The codec type id, one of the values of the `CompressionType` object.
+* *codec* (`object`): The codec implementation as an object with `encode` (to compress) and `decode` (to decompress) functions, each of which take a `Uint8Array` as input and return a `Uint8Array` as output.
+
+```js
+import { CompressionType, setCompressionCodec } from '@uwdata/flechette';
+import { ZstdCodec } from 'zstd-codec';
+
+// register zstd compression codec
+await new Promise((resolve) => {
+  ZstdCodec.run((zstd) => {
+    const codec = new zstd.Simple();
+    setCompressionCodec(CompressionType.ZSTD, {
+      encode: (data) => codec.compress(data),
+      decode: (data) => codec.decompress(data)
+    });
+    resolve();
+  });
+});
+```
+
+```js
+import { CompressionType, setCompressionCodec } from '@uwdata/flechette';
+import * as lz4 from 'lz4js';
+
+// register lz4_frame compression codec
+setCompressionCodec(CompressionType.LZ4_FRAME, {
+  encode: (data) => lz4.compress(data),
+  decode: (data) => lz4.decompress(data)
+});
+```

package-lock.json

@@ -37,10 +37,12 @@
     "@rollup/plugin-terser": "^0.4.4",
     "apache-arrow": "^21.1.0",
     "eslint": "^9.39.2",
+    "lz4js": "0.2.0",
     "rimraf": "^6.1.2",
     "rollup": "^4.56.0",
     "rollup-plugin-bundle-size": "^1.0.3",
     "typescript": "^5.9.3",
-    "vitest": "^4.0.18"
+    "vitest": "^4.0.18",
+    "zstd-codec": "0.1.5"
   }
 }

package.json

@@ -0,0 +1,71 @@
+/** @import { Codec, CompressionType_ } from './types.js' */
+import { CompressionType } from './constants.js';
+import { writeInt64 } from './encode/builder.js';
+import { keyFor } from './util/objects.js';
+import { readInt64 } from './util/read.js';
+
+const LENGTH_NO_COMPRESSED_DATA = -1;
+const COMPRESS_LENGTH_PREFIX = 8;
+
+/**
+ * Return an error message for a missing codec.
+ * @param {CompressionType_} type The codec type.
+ */
+export function missingCodec(type) {
+  return `Missing compression codec "${keyFor(CompressionType, type)}" (id ${type})`;
+}
+
+/** @type {Map<CompressionType_, Codec>} */
+const codecs = new Map;
+
+/**
+ * Register a codec to use for compressing or decompressing Arrow buffers.
+ * @param {CompressionType_} type The compression type.
+ * @param {Codec} codec The codec implementation.
+ */
+export function setCompressionCodec(type, codec) {
+  codecs.set(type, codec);
+}
+
+/**
+ * Returns a compression codec for the provided type, or null if not found.
+ * Compression codecs must first be registered using *setCompressionCodec*.
+ * @param {CompressionType_ | null} [type] The compression type.
+ * @returns {Codec | null} The compression codec, or null if not registered.
+ */
+export function getCompressionCodec(type) {
+  return (type != null && codecs.get(type)) || null;
+}
+
+/**
+ * Decompress an Arrow buffer, return decompressed bytes and region metadata.
+ * @param {Uint8Array} body The message body.
+ * @param {{ offset: number, length: number }} region Buffer region metadata.
+ * @param {Codec} codec A compression codec.
+ * @returns {{ bytes: Uint8Array, offset: number, length: number }}
+ */
+export function decompressBuffer(body, { offset, length }, codec) {
+  if (length === 0) {
+    return { bytes: new Uint8Array(0), offset: 0, length: 0 };
+  }
+  const ulen = readInt64(body, offset); // uncompressed length
+  const buf = body.subarray(offset + COMPRESS_LENGTH_PREFIX, offset + length);
+  const bytes = (ulen === LENGTH_NO_COMPRESSED_DATA) ? buf : codec.decode(buf);
+  return { bytes, offset: 0, length: bytes.length };
+}
+
+/**
+ * Compress an Arrow buffer, return encoded bytes. If the compression does
+ * not decrease the overall length, retains uncompressed bytes.
+ * @param {Uint8Array} bytes The byte buffer to compress.
+ * @param {Codec} codec A compression codec.
+ */
+export function compressBuffer(bytes, codec) {
+  const compressed = codec.encode(bytes);
+  const keep = compressed.length < bytes.length;
+  const data = keep ? compressed : bytes;
+  const buf = new Uint8Array(COMPRESS_LENGTH_PREFIX + data.length);
+  writeInt64(buf, 0, keep ? bytes.length : LENGTH_NO_COMPRESSED_DATA);
+  buf.set(data, COMPRESS_LENGTH_PREFIX);
+  return buf;
+}

src/compression.js

@@ -374,3 +374,35 @@ export const UnionMode = /** @type {const} */ ({
   /** Dense union layout with offsets into value arrays. */
   Dense: 1
 });
+
+/**
+ * Compression types.
+ */
+export const CompressionType = /** @type {const} */ ({
+  /**
+   * LZ4 frame compression.
+   * Not to be confused with "raw" (also called "block") format.
+   */
+  LZ4_FRAME: 0,
+  /** Zstandard compression. */
+  ZSTD: 1
+});
+
+/**
+ * Body compression methods.
+ * Provided for forward compatibility in case Arrow needs to support
+ * different strategies for compressing the IPC message body (like
+ * whole-body compression rather than buffer-level) in the future.
+ */
+export const BodyCompressionMethod = /** @type {const} */ ({
+  /**
+   * Each constituent buffer is first compressed with the indicated
+   * compressor, and then written with the uncompressed length in the first 8
+   * bytes as a 64-bit little-endian signed integer followed by the compressed
+   * buffer bytes (and then padding as required by the protocol). The
+   * uncompressed length may be set to -1 to indicate that the data that
+   * follows is not compressed, which can be useful for cases where
+   * compression does not yield appreciable savings.
+   */
+  BUFFER: 0
+});

src/constants.js

@@ -0,0 +1,23 @@
+/**
+ * @import { BodyCompression, BodyCompressionMethod_, CompressionType_ } from '../types.js'
+ */
+import { BodyCompressionMethod, CompressionType } from '../constants.js';
+import { readInt8, readObject } from '../util/read.js';
+
+/**
+ * Decode record batch body compression metadata.
+ * @param {Uint8Array} buf A byte buffer of binary Arrow IPC data
+ * @param {number} index The starting index in the byte buffer
+ * @returns {BodyCompression | undefined} The body compression metadata
+ */
+export function decodeBodyCompression(buf, index) {
+  //  4: codec
+  //  6: method
+  const get = readObject(buf, index);
+  return {
+    codec: /** @type {CompressionType_} */(
+      get(4, readInt8, CompressionType.LZ4_FRAME)),
+    method: /** @type {BodyCompressionMethod_} */(
+      get(6, readInt8, BodyCompressionMethod.BUFFER))
+  };
+}

src/decode/body-compression.js

@@ -3,6 +3,7 @@
  */
 import { Version } from '../constants.js';
 import { readInt64, readObject, readOffset, readVector } from '../util/read.js';
+import { decodeBodyCompression } from './body-compression.js';
 
 /**
  * Decode a record batch.
@@ -15,12 +16,9 @@ export function decodeRecordBatch(buf, index, version) {
   //  4: length
   //  6: nodes
   //  8: buffers
-  // 10: compression (not supported)
+  // 10: compression (requires codec plug-in)
   // 12: variadicBuffers (buffer counts for view-typed fields)
   const get = readObject(buf, index);
-  if (get(10, readOffset, 0)) {
-    throw new Error('Record batch compression not implemented');
-  }
 
   // If an Arrow buffer was written before version 4,
   // advance 8 bytes to skip the now-removed page_id field
@@ -36,6 +34,7 @@ export function decodeRecordBatch(buf, index, version) {
       offset: readInt64(buf, pos + offset),
       length: readInt64(buf, pos + offset + 8)
     })),
+    compression: get(10, decodeBodyCompression),
     variadic: readVector(buf, get(12, readOffset), 8, readInt64)
   };
 }