fix(ipc): body prefix must be uncompressed length

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;
 `);
 }
 

src/ipc/writer.ts

@@ -318,7 +318,11 @@ export class RecordBatchWriter<T extends TypeMap = any> 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));

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<void> {
+    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<void>((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;
+}

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<void> {
-    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<void>((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<Uint32, Int32>(new Uint32, new Int32, 0);
         const writer = new RecordBatchFileWriter();

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<void> {
-    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<void>((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 });