From 3c1a64983b3bad00b6e8feda0b3f7fb8e114c863 Mon Sep 17 00:00:00 2001 From: Kent Wu Date: Wed, 2 Sep 2026 16:36:22 -0400 Subject: [PATCH 1/2] fix(ipc): body prefix must be uncompressed length Per the Arrow columnar format spec, the eight-byte prefix on each compressed IPC body buffer must hold the buffer's uncompressed length so consumers can size the decompression destination buffer. The writer was emitting the compressed length instead, which caused PyArrow (and any other implementation that relies on the prefix for allocation) to reject Arrow JS-produced ZSTD and LZ4_FRAME streams with a "destination buffer too small" error. JS-to-JS round trips did not surface the bug because the reader sizes buffers from the codec frame rather than the prefix. Add a byte-level regression test that inspects each emitted prefix against the decompressed payload length for both allowed codecs (LZ4_FRAME and ZSTD), and factor the shared codec-registration helper out of stream-writer-tests.ts so the new test file does not have to import from another test file. --- src/ipc/writer.ts | 6 +- test/unit/ipc/writer/compression-codecs.ts | 87 +++++++++++++++++++++ test/unit/ipc/writer/file-writer-tests.ts | 64 +++++++-------- test/unit/ipc/writer/stream-writer-tests.ts | 57 ++++++-------- 4 files changed, 145 insertions(+), 69 deletions(-) create mode 100644 test/unit/ipc/writer/compression-codecs.ts diff --git a/src/ipc/writer.ts b/src/ipc/writer.ts index 7d783eb9..6554498c 100644 --- a/src/ipc/writer.ts +++ b/src/ipc/writer.ts @@ -318,7 +318,11 @@ export class RecordBatchWriter extends ReadableInterop< const isCompressionEffective = compressed.length < byteBuf.length; const finalBuffer = isCompressionEffective ? compressed : byteBuf; - const byteLength = isCompressionEffective ? finalBuffer.length : LENGTH_NO_COMPRESSED_DATA; + // Per the Arrow columnar format spec, the 8-byte prefix on a + // compressed body buffer holds the *uncompressed* length so that + // readers can size the decompression destination buffer. When the + // buffer was left uncompressed, the prefix is LENGTH_NO_COMPRESSED_DATA (-1). + const byteLength = isCompressionEffective ? byteBuf.length : LENGTH_NO_COMPRESSED_DATA; const lengthPrefix = new flatbuffers.ByteBuffer(new Uint8Array(COMPRESS_LENGTH_PREFIX)); lengthPrefix.writeInt64(0, BigInt(byteLength)); diff --git a/test/unit/ipc/writer/compression-codecs.ts b/test/unit/ipc/writer/compression-codecs.ts new file mode 100644 index 00000000..8784cbf1 --- /dev/null +++ b/test/unit/ipc/writer/compression-codecs.ts @@ -0,0 +1,87 @@ +// 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 { ByteBuffer } from 'flatbuffers'; + +import { + Codec, + compressionRegistry, + CompressionType, + MessageReader, +} from 'apache-arrow'; +import * as lz4js from 'lz4js'; + +const LENGTH_NO_COMPRESSED_DATA = -1; +const COMPRESS_LENGTH_PREFIX = 8; +// RecordBatchFileWriter prefixes its output with 6 bytes "ARROW1" + 2 bytes padding +// before the stream messages. Skip past those before handing bytes to MessageReader. +export const FILE_FORMAT_HEADER_LENGTH = 8; + +export async function registerCompressionCodecs(): Promise { + if (compressionRegistry.get(CompressionType.LZ4_FRAME) === null) { + 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); + } + + if (compressionRegistry.get(CompressionType.ZSTD) === null) { + const { ZstdCodec } = await import('zstd-codec'); + await new Promise((resolve) => { + ZstdCodec.run((zstd: any) => { + const simple = new zstd.Simple(); + const zstdCodec: Codec = { + encode(data: Uint8Array): Uint8Array { return simple.compress(data); }, + decode(data: Uint8Array): Uint8Array { return simple.decompress(data); } + }; + compressionRegistry.set(CompressionType.ZSTD, zstdCodec); + resolve(); + }); + }); + } +} + +/** + * Walks the IPC messages in `bytes` and returns the (prefix, decompressedLength) + * pair for every compressed body buffer. Per the Arrow columnar format spec, callers + * should assert `prefix === decompressedLength` — the eight-byte prefix must hold + * the uncompressed length so a reader can size the decompression destination buffer. + */ +export function extractCompressedPrefixes( + bytes: Uint8Array, + codec: Codec, +): { prefix: number; decompressedLength: number }[] { + const reader = new MessageReader(bytes); + const results: { prefix: number; decompressedLength: number }[] = []; + + for (const message of reader) { + const body = reader.readMessageBody(message.bodyLength); + if (!message.isRecordBatch()) continue; + + for (const region of message.header().buffers) { + if (region.length === 0) continue; + const buffer = body.subarray(region.offset, region.offset + region.length); + const prefix = Number(new ByteBuffer(buffer).readInt64(0)); + if (prefix === LENGTH_NO_COMPRESSED_DATA) continue; + + const decompressed = codec.decode!(buffer.subarray(COMPRESS_LENGTH_PREFIX)); + results.push({ prefix, decompressedLength: decompressed.length }); + } + } + return results; +} diff --git a/test/unit/ipc/writer/file-writer-tests.ts b/test/unit/ipc/writer/file-writer-tests.ts index f6632d84..3cefa7a8 100644 --- a/test/unit/ipc/writer/file-writer-tests.ts +++ b/test/unit/ipc/writer/file-writer-tests.ts @@ -24,7 +24,6 @@ import { validateRecordBatchIterator } from '../validate.js'; import { builderThroughIterable, - Codec, compressionRegistry, CompressionType, Dictionary, @@ -33,43 +32,16 @@ import { RecordBatchFileWriter, RecordBatchReader, Table, + tableFromArrays, Uint32, Vector } from 'apache-arrow'; -import * as lz4js from 'lz4js'; - -export async function registerCompressionCodecs(): Promise { - if (compressionRegistry.get(CompressionType.LZ4_FRAME) === null) { - 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); - } - if (compressionRegistry.get(CompressionType.ZSTD) === null) { - const { ZstdCodec } = await import('zstd-codec'); - await new Promise((resolve) => { - ZstdCodec.run((zstd: any) => { - const simple = new zstd.Simple(); - const zstdCodec: Codec = { - encode(data: Uint8Array): Uint8Array { - return simple.compress(data); - }, - decode(data: Uint8Array): Uint8Array { - return simple.decompress(data); - } - }; - compressionRegistry.set(CompressionType.ZSTD, zstdCodec); - resolve(); - }); - }); - } -} +import { + extractCompressedPrefixes, + FILE_FORMAT_HEADER_LENGTH, + registerCompressionCodecs, +} from './compression-codecs.js'; describe('RecordBatchFileWriter', () => { for (const table of generateRandomTables([10, 20, 30])) { @@ -90,6 +62,30 @@ describe('RecordBatchFileWriter', () => { testFileWriter(table, testName, { compressionType }); } + describe('compressed body buffer length prefix', () => { + for (const compressionType of compressionTypes) { + it(`writes the uncompressed length for ${CompressionType[compressionType]}`, async () => { + // Highly compressible data so most buffers take the compressed branch. + const fixture = tableFromArrays({ + id: Int32Array.from({ length: 1000 }, (_, i) => i), + label: Array.from({ length: 1000 }, () => 'a highly compressible value'), + }); + + const bytes = await RecordBatchFileWriter.writeAll(fixture, { compressionType }).toUint8Array(); + const prefixes = extractCompressedPrefixes( + bytes.subarray(FILE_FORMAT_HEADER_LENGTH), + compressionRegistry.get(compressionType)!, + ); + + // Guard against a vacuous pass — the fixture must exercise the branch. + expect(prefixes.length).toBeGreaterThan(0); + for (const { prefix, decompressedLength } of prefixes) { + expect(prefix).toBe(decompressedLength); + } + }); + } + }); + it('should throw if attempting to write replacement dictionary batches', async () => { const type = new Dictionary(new Uint32, new Int32, 0); const writer = new RecordBatchFileWriter(); diff --git a/test/unit/ipc/writer/stream-writer-tests.ts b/test/unit/ipc/writer/stream-writer-tests.ts index 2c2e3d3d..8abb58cb 100644 --- a/test/unit/ipc/writer/stream-writer-tests.ts +++ b/test/unit/ipc/writer/stream-writer-tests.ts @@ -25,7 +25,6 @@ import { validateRecordBatchIterator } from '../validate.js'; import type { RecordBatchStreamWriterOptions } from 'apache-arrow/ipc/writer'; import { builderThroughIterable, - Codec, compressionRegistry, CompressionType, Data, @@ -37,43 +36,12 @@ import { RecordBatchStreamWriter, Schema, Table, + tableFromArrays, Uint32, Vector } from 'apache-arrow'; -import * as lz4js from 'lz4js'; - -export async function registerCompressionCodecs(): Promise { - if (compressionRegistry.get(CompressionType.LZ4_FRAME) === null) { - 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); - } - if (compressionRegistry.get(CompressionType.ZSTD) === null) { - const { ZstdCodec } = await import('zstd-codec'); - await new Promise((resolve) => { - ZstdCodec.run((zstd: any) => { - const simple = new zstd.Simple(); - const zstdCodec: Codec = { - encode(data: Uint8Array): Uint8Array { - return simple.compress(data); - }, - decode(data: Uint8Array): Uint8Array { - return simple.decompress(data); - } - }; - compressionRegistry.set(CompressionType.ZSTD, zstdCodec); - resolve(); - }); - }); - } -} +import { extractCompressedPrefixes, registerCompressionCodecs } from './compression-codecs.js'; describe('RecordBatchStreamWriter', () => { @@ -94,6 +62,27 @@ describe('RecordBatchStreamWriter', () => { testStreamWriter(table, testName, { compressionType }); } + describe('compressed body buffer length prefix', () => { + for (const compressionType of compressionTypes) { + it(`writes the uncompressed length for ${CompressionType[compressionType]}`, async () => { + // Highly compressible data so most buffers take the compressed branch. + const table = tableFromArrays({ + id: Int32Array.from({ length: 1000 }, (_, i) => i), + label: Array.from({ length: 1000 }, () => 'a highly compressible value'), + }); + + const bytes = await RecordBatchStreamWriter.writeAll(table, { compressionType }).toUint8Array(); + const prefixes = extractCompressedPrefixes(bytes, compressionRegistry.get(compressionType)!); + + // Guard against a vacuous pass — the fixture must exercise the branch. + expect(prefixes.length).toBeGreaterThan(0); + for (const { prefix, decompressedLength } of prefixes) { + expect(prefix).toBe(decompressedLength); + } + }); + } + }); + for (const table of generateRandomTables([10, 20, 30])) { const testName = `[${table.schema.fields.join(', ')}]`; testStreamWriter(table, testName, { writeLegacyIpcFormat: true }); From 59fde71f9cf75a29a52d5fa7724f851ca213eb7d Mon Sep 17 00:00:00 2001 From: Kent Wu Date: Thu, 3 Sep 2026 10:16:17 -0400 Subject: [PATCH 2/2] fix(build): preserve compressionType in UMD externs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Closure externs generator crawls exported classes for their static and prototype property names, so option-object properties (referenced only inside constructor bodies) never get reserved. Closure ADVANCED renames them, and external UMD callers passing `{ compressionType: X }` to a writer receive silently-uncompressed output — the writer's lookup targets the renamed property, gets nothing, and sets `_compression = null`. Declare `compressionType` as reserved so Closure preserves it. `autoDestroy` is already preserved coincidentally (it exists as a prototype property on `RecordBatchReader`); the identically-shaped bug for `writeLegacyIpcFormat` is left for follow-up. --- gulp/closure-task.js | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/gulp/closure-task.js b/gulp/closure-task.js index 6916b59a..e6049868 100644 --- a/gulp/closure-task.js +++ b/gulp/closure-task.js @@ -215,6 +215,10 @@ Encoding[2] = function() {}; Encoding.UTF8_BYTES = function() {}; /** @type {?} */ Encoding.UTF16_STRING = function() {}; + +var RecordBatchWriterOptions = function() {}; +/** @type {?} */ +RecordBatchWriterOptions.prototype.compressionType; `); }