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