From c1efe68d7601d2404893ab9fdce9293aed105694 Mon Sep 17 00:00:00 2001 From: Pedro Holanda Date: Mon, 7 Sep 2026 14:09:37 +0200 Subject: [PATCH 01/13] feat(ipc): add LZ4 and ZSTD compression functions and ArrowIpcCompressor --- src/nanoarrow/ipc/codecs.c | 167 ++++++++++++++++++++++++++++--- src/nanoarrow/ipc/codecs_test.cc | 146 +++++++++++++++++++++++++++ src/nanoarrow/nanoarrow_ipc.h | 67 +++++++++++++ src/nanoarrow/nanoarrow_ipc.hpp | 22 ++++ 4 files changed, 386 insertions(+), 16 deletions(-) diff --git a/src/nanoarrow/ipc/codecs.c b/src/nanoarrow/ipc/codecs.c index ab2979893..cf8d552af 100644 --- a/src/nanoarrow/ipc/codecs.c +++ b/src/nanoarrow/ipc/codecs.c @@ -44,6 +44,32 @@ static ArrowErrorCode ArrowIpcDecompressZstd(struct ArrowBufferView src, uint8_t return NANOARROW_OK; } + +static ArrowErrorCode ArrowIpcCompressZstd(struct ArrowBufferView src, + struct ArrowBuffer* dst, + struct ArrowError* error) { + size_t dst_capacity = ZSTD_compressBound((size_t)src.size_bytes); + if (ZSTD_isError(dst_capacity)) { + ArrowErrorSet(error, "ZSTD_compressBound(%" PRId64 ") failed with error '%s'", + src.size_bytes, ZSTD_getErrorName(dst_capacity)); + return EINVAL; + } + + NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowBufferReserve(dst, (int64_t)dst_capacity), + error); + + size_t code = ZSTD_compress((void*)(dst->data + dst->size_bytes), dst_capacity, + src.data.data, (size_t)src.size_bytes, ZSTD_CLEVEL_DEFAULT); + if (ZSTD_isError(code)) { + ArrowErrorSet(error, + "ZSTD_compress([buffer with %" PRId64 " bytes]) failed with error '%s'", + src.size_bytes, ZSTD_getErrorName(code)); + return EIO; + } + + dst->size_bytes += (int64_t)code; + return NANOARROW_OK; +} #endif ArrowIpcDecompressFunction ArrowIpcGetZstdDecompressionFunction(void) { @@ -54,6 +80,14 @@ ArrowIpcDecompressFunction ArrowIpcGetZstdDecompressionFunction(void) { #endif } +ArrowIpcCompressFunction ArrowIpcGetZstdCompressionFunction(void) { +#if defined(NANOARROW_IPC_WITH_ZSTD) + return &ArrowIpcCompressZstd; +#else + return NULL; +#endif +} + #if defined(NANOARROW_IPC_WITH_LZ4) #include #include @@ -103,6 +137,33 @@ static ArrowErrorCode ArrowIpcDecompressLZ4(struct ArrowBufferView src, uint8_t* NANOARROW_UNUSED(LZ4F_freeDecompressionContext(ctx)); return NANOARROW_OK; } + +static ArrowErrorCode ArrowIpcCompressLZ4(struct ArrowBufferView src, + struct ArrowBuffer* dst, + struct ArrowError* error) { + // Default preferences (default compression level, no content size, no checksums). + // This produces a single complete frame, which is what ArrowIpcDecompressLZ4() + // and Arrow C++ expect. + LZ4F_preferences_t prefs; + memset(&prefs, 0, sizeof(prefs)); + + size_t dst_capacity = LZ4F_compressFrameBound((size_t)src.size_bytes, &prefs); + NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowBufferReserve(dst, (int64_t)dst_capacity), + error); + + size_t code = LZ4F_compressFrame((void*)(dst->data + dst->size_bytes), dst_capacity, + src.data.data, (size_t)src.size_bytes, &prefs); + if (LZ4F_isError(code)) { + ArrowErrorSet(error, + "LZ4F_compressFrame([buffer with %" PRId64 + " bytes]) failed with error '%s'", + src.size_bytes, LZ4F_getErrorName(code)); + return EIO; + } + + dst->size_bytes += (int64_t)code; + return NANOARROW_OK; +} #endif ArrowIpcDecompressFunction ArrowIpcGetLZ4DecompressionFunction(void) { @@ -113,6 +174,26 @@ ArrowIpcDecompressFunction ArrowIpcGetLZ4DecompressionFunction(void) { #endif } +ArrowIpcCompressFunction ArrowIpcGetLZ4CompressionFunction(void) { +#if defined(NANOARROW_IPC_WITH_LZ4) + return &ArrowIpcCompressLZ4; +#else + return NULL; +#endif +} + +// The serial decompressor and compressor keep one function per codec, indexed by +// enum ArrowIpcCompressionType (NONE is never a codec) +static int ArrowIpcCompressionTypeIsCodec(enum ArrowIpcCompressionType compression_type) { + switch (compression_type) { + case NANOARROW_IPC_COMPRESSION_TYPE_ZSTD: + case NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME: + return 1; + default: + return 0; + } +} + struct ArrowIpcSerialDecompressorPrivate { ArrowIpcDecompressFunction decompress_functions[3]; }; @@ -124,18 +205,13 @@ static ArrowErrorCode ArrowIpcSerialDecompressorAdd( struct ArrowIpcSerialDecompressorPrivate* private_data = (struct ArrowIpcSerialDecompressorPrivate*)decompressor->private_data; - ArrowIpcDecompressFunction fn = NULL; - switch (compression_type) { - case NANOARROW_IPC_COMPRESSION_TYPE_ZSTD: - case NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME: - fn = private_data->decompress_functions[compression_type]; - break; - default: - ArrowErrorSet(error, "Unknown decompression type with value %d", - (int)compression_type); - return EINVAL; + if (!ArrowIpcCompressionTypeIsCodec(compression_type)) { + ArrowErrorSet(error, "Unknown decompression type with value %d", + (int)compression_type); + return EINVAL; } + ArrowIpcDecompressFunction fn = private_data->decompress_functions[compression_type]; if (fn == NULL) { ArrowErrorSet( error, "Compression type with value %d not supported by this build of nanoarrow", @@ -187,14 +263,73 @@ ArrowErrorCode ArrowIpcSerialDecompressorSetFunction( struct ArrowIpcSerialDecompressorPrivate* private_data = (struct ArrowIpcSerialDecompressorPrivate*)decompressor->private_data; - switch (compression_type) { - case NANOARROW_IPC_COMPRESSION_TYPE_ZSTD: - case NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME: - break; - default: - return EINVAL; + if (!ArrowIpcCompressionTypeIsCodec(compression_type)) { + return EINVAL; } private_data->decompress_functions[compression_type] = decompress_function; return NANOARROW_OK; } + +struct ArrowIpcSerialCompressorPrivate { + ArrowIpcCompressFunction compress_functions[3]; +}; + +static ArrowErrorCode ArrowIpcSerialCompressorCompress( + struct ArrowIpcCompressor* compressor, enum ArrowIpcCompressionType compression_type, + struct ArrowBufferView src, struct ArrowBuffer* dst, struct ArrowError* error) { + struct ArrowIpcSerialCompressorPrivate* private_data = + (struct ArrowIpcSerialCompressorPrivate*)compressor->private_data; + + if (!ArrowIpcCompressionTypeIsCodec(compression_type)) { + ArrowErrorSet(error, "Unknown compression type with value %d", (int)compression_type); + return EINVAL; + } + + ArrowIpcCompressFunction fn = private_data->compress_functions[compression_type]; + if (fn == NULL) { + ArrowErrorSet( + error, "Compression type with value %d not supported by this build of nanoarrow", + (int)compression_type); + return ENOTSUP; + } + + NANOARROW_RETURN_NOT_OK(fn(src, dst, error)); + return NANOARROW_OK; +} + +static void ArrowIpcSerialCompressorRelease(struct ArrowIpcCompressor* compressor) { + ArrowFree(compressor->private_data); + compressor->release = NULL; +} + +ArrowErrorCode ArrowIpcSerialCompressor(struct ArrowIpcCompressor* compressor) { + compressor->compress = &ArrowIpcSerialCompressorCompress; + compressor->release = &ArrowIpcSerialCompressorRelease; + compressor->private_data = ArrowMalloc(sizeof(struct ArrowIpcSerialCompressorPrivate)); + if (compressor->private_data == NULL) { + return ENOMEM; + } + + memset(compressor->private_data, 0, sizeof(struct ArrowIpcSerialCompressorPrivate)); + ArrowIpcSerialCompressorSetFunction(compressor, NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, + ArrowIpcGetZstdCompressionFunction()); + ArrowIpcSerialCompressorSetFunction(compressor, + NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, + ArrowIpcGetLZ4CompressionFunction()); + return NANOARROW_OK; +} + +ArrowErrorCode ArrowIpcSerialCompressorSetFunction( + struct ArrowIpcCompressor* compressor, enum ArrowIpcCompressionType compression_type, + ArrowIpcCompressFunction compress_function) { + struct ArrowIpcSerialCompressorPrivate* private_data = + (struct ArrowIpcSerialCompressorPrivate*)compressor->private_data; + + if (!ArrowIpcCompressionTypeIsCodec(compression_type)) { + return EINVAL; + } + + private_data->compress_functions[compression_type] = compress_function; + return NANOARROW_OK; +} diff --git a/src/nanoarrow/ipc/codecs_test.cc b/src/nanoarrow/ipc/codecs_test.cc index 76278174a..f298e4149 100644 --- a/src/nanoarrow/ipc/codecs_test.cc +++ b/src/nanoarrow/ipc/codecs_test.cc @@ -16,6 +16,7 @@ // under the License. #include +#include #include #include @@ -185,3 +186,148 @@ TEST(NanoarrowIpcTest, SerialDecompressor) { EXPECT_STREQ(error.message, "Compression type with value 2 not supported by this build of nanoarrow"); } + +// Compress input (appending to a buffer that already has content), decompress the +// appended bytes, and check that the result matches the input. Returns the number of +// compressed bytes that were appended. +static int64_t TestCompressRoundtrip(ArrowIpcCompressFunction compress, + ArrowIpcDecompressFunction decompress, + const std::vector& input) { + struct ArrowError error {}; + nanoarrow::UniqueBuffer compressed; + + // Content already in dst must be preserved (compress functions only append) + const char* existing = "existing"; + const int64_t existing_size = 8; + EXPECT_EQ(ArrowBufferAppend(compressed.get(), existing, existing_size), NANOARROW_OK); + + EXPECT_EQ(compress({{input.data()}, static_cast(input.size())}, + compressed.get(), &error), + NANOARROW_OK) + << error.message; + EXPECT_GT(compressed->size_bytes, existing_size); + EXPECT_EQ(std::memcmp(compressed->data, existing, existing_size), 0); + + std::vector output(input.size()); + struct ArrowBufferView compressed_view = {{compressed->data + existing_size}, + compressed->size_bytes - existing_size}; + EXPECT_EQ(decompress(compressed_view, output.data(), + static_cast(output.size()), &error), + NANOARROW_OK) + << error.message; + EXPECT_EQ(output, input); + + return compressed->size_bytes - existing_size; +} + +static std::vector CompressibleInput(size_t n) { + std::vector out(n); + for (size_t i = 0; i < n; i++) { + out[i] = static_cast(i % 7); + } + return out; +} + +// Check compress/decompress on empty, small, and multi-block inputs +static void TestCompressionFunctions(ArrowIpcCompressFunction compress, + ArrowIpcDecompressFunction decompress) { + ASSERT_NE(compress, nullptr); + ASSERT_NE(decompress, nullptr); + + TestCompressRoundtrip(compress, decompress, {}); + TestCompressRoundtrip( + compress, decompress, + std::vector(kUncompressed012, + kUncompressed012 + sizeof(kUncompressed012))); + + // Large enough to span several blocks; a repetitive input must actually shrink + auto input = CompressibleInput(1 << 20); + int64_t compressed_size = TestCompressRoundtrip(compress, decompress, input); + EXPECT_LT(compressed_size, static_cast(input.size() / 10)); +} + +TEST(NanoarrowIpcTest, NanoarrowIpcZstdCompressBuildMatchesRuntime) { +#if defined(NANOARROW_IPC_WITH_ZSTD) + ASSERT_NE(ArrowIpcGetZstdCompressionFunction(), nullptr); +#else + ASSERT_EQ(ArrowIpcGetZstdCompressionFunction(), nullptr); +#endif +} + +TEST(NanoarrowIpcTest, ZstdCompressRoundtrip) { + if (ArrowIpcGetZstdCompressionFunction() == nullptr) { + GTEST_SKIP() << "nanoarrow_ipc not built with NANOARROW_IPC_WITH_ZSTD"; + } + TestCompressionFunctions(ArrowIpcGetZstdCompressionFunction(), + ArrowIpcGetZstdDecompressionFunction()); +} + +TEST(NanoarrowIpcTest, NanoarrowIpcLZ4CompressBuildMatchesRuntime) { +#if defined(NANOARROW_IPC_WITH_LZ4) + ASSERT_NE(ArrowIpcGetLZ4CompressionFunction(), nullptr); +#else + ASSERT_EQ(ArrowIpcGetLZ4CompressionFunction(), nullptr); +#endif +} + +TEST(NanoarrowIpcTest, LZ4CompressRoundtrip) { + if (ArrowIpcGetLZ4CompressionFunction() == nullptr) { + GTEST_SKIP() << "nanoarrow_ipc not built with NANOARROW_IPC_WITH_LZ4"; + } + TestCompressionFunctions(ArrowIpcGetLZ4CompressionFunction(), + ArrowIpcGetLZ4DecompressionFunction()); +} + +TEST(NanoarrowIpcTest, SerialCompressor) { + struct ArrowError error {}; + nanoarrow::ipc::UniqueCompressor compressor; + + ASSERT_EQ(ArrowIpcSerialCompressor(compressor.get()), NANOARROW_OK); + + // Check the function setter error + ASSERT_EQ(ArrowIpcSerialCompressorSetFunction( + compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_NONE, nullptr), + EINVAL); + + // NONE is not a codec that can be used to compress + nanoarrow::UniqueBuffer dst; + EXPECT_EQ(compressor->compress(compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_NONE, + {{nullptr}, 0}, dst.get(), &error), + EINVAL); + EXPECT_STREQ(error.message, "Unknown compression type with value 0"); + + // Check a compress for a supported codec if we have one (or for an error if we don't) + if (ArrowIpcGetZstdCompressionFunction() != nullptr) { + ASSERT_EQ(compressor->compress(compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, + {{kUncompressed012}, sizeof(kUncompressed012)}, + dst.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_GT(dst->size_bytes, 0); + + uint8_t out[sizeof(kUncompressed012)]; + std::memset(out, 0, sizeof(out)); + ASSERT_EQ(ArrowIpcGetZstdDecompressionFunction()({{dst->data}, dst->size_bytes}, out, + sizeof(out), &error), + NANOARROW_OK) + << error.message; + EXPECT_TRUE(std::memcmp(out, kUncompressed012, sizeof(kUncompressed012)) == 0); + } else { + EXPECT_EQ(compressor->compress(compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, + {{nullptr}, 0}, dst.get(), &error), + ENOTSUP); + EXPECT_STREQ( + error.message, + "Compression type with value 2 not supported by this build of nanoarrow"); + } + + // Either way, if we explicitly remove support for a codec, we should get an error + ASSERT_EQ(ArrowIpcSerialCompressorSetFunction( + compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, nullptr), + NANOARROW_OK); + EXPECT_EQ(compressor->compress(compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, + {{nullptr}, 0}, dst.get(), &error), + ENOTSUP); + EXPECT_STREQ(error.message, + "Compression type with value 2 not supported by this build of nanoarrow"); +} diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index fb38750c5..797019120 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -33,6 +33,14 @@ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcSerialDecompressor) #define ArrowIpcSerialDecompressorSetFunction \ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcSerialDecompressorSetFunction) +#define ArrowIpcGetZstdCompressionFunction \ + NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcGetZstdCompressionFunction) +#define ArrowIpcGetLZ4CompressionFunction \ + NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcGetLZ4CompressionFunction) +#define ArrowIpcSerialCompressor \ + NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcSerialCompressor) +#define ArrowIpcSerialCompressorSetFunction \ + NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcSerialCompressorSetFunction) #define ArrowIpcDecoderInit NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcDecoderInit) #define ArrowIpcDecoderReset NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcDecoderReset) #define ArrowIpcDecoderSetDecompressor \ @@ -412,6 +420,65 @@ ArrowIpcSerialDecompressorSetFunction(struct ArrowIpcDecompressor* decompressor, enum ArrowIpcCompressionType compression_type, ArrowIpcDecompressFunction decompress_function); +/// \brief A user-extensible compressor +/// +/// The ArrowIpcCompressor is the underlying object that enables buffer compression +/// in the ArrowIpcEncoder. An implementation of a compressor may support more than one +/// ArrowIpcCompressionType. +struct ArrowIpcCompressor { + /// \brief Compress a buffer + /// + /// Compresses src using compression_type and appends the compressed bytes to dst. + /// Any content already in dst must be preserved (i.e., implementations may only + /// append to dst). + ArrowErrorCode (*compress)(struct ArrowIpcCompressor* compressor, + enum ArrowIpcCompressionType compression_type, + struct ArrowBufferView src, struct ArrowBuffer* dst, + struct ArrowError* error); + + /// \brief Release the compressor and any resources it may be holding + /// + /// Release callback implementations must set the release member to NULL. + /// Callers must check that the release callback is not NULL before calling + /// compress() or release(). + void (*release)(struct ArrowIpcCompressor* compressor); + + /// \brief Implementation-specific opaque data + void* private_data; +}; + +/// \brief A self-contained compression function +/// +/// Compresses src and appends the compressed bytes to dst. Because the compressed +/// size is not known in advance, implementations are responsible for reserving +/// sufficient space in dst (e.g., using the compression library's bound function) +/// and must only append to dst. +typedef ArrowErrorCode (*ArrowIpcCompressFunction)(struct ArrowBufferView src, + struct ArrowBuffer* dst, + struct ArrowError* error); + +/// \brief Get the compression function for ZSTD +/// +/// The result will be NULL if nanoarrow was not built with NANOARROW_IPC_WITH_ZSTD. +NANOARROW_DLL ArrowIpcCompressFunction ArrowIpcGetZstdCompressionFunction(void); + +/// \brief Get the compression function for LZ4 +/// +/// The result will be NULL if nanoarrow was not built with NANOARROW_IPC_WITH_LZ4. +NANOARROW_DLL ArrowIpcCompressFunction ArrowIpcGetLZ4CompressionFunction(void); + +/// \brief An ArrowIpcCompressor implementation that performs compression in serial +NANOARROW_DLL ArrowErrorCode +ArrowIpcSerialCompressor(struct ArrowIpcCompressor* compressor); + +/// \brief Override the ArrowIpcCompressFunction used for a specific compression type +/// +/// This may be used to inject support for a particular type of compression if used +/// with a version of nanoarrow with unknown or minimal capabilities. +NANOARROW_DLL ArrowErrorCode ArrowIpcSerialCompressorSetFunction( + struct ArrowIpcCompressor* compressor, enum ArrowIpcCompressionType compression_type, + ArrowIpcCompressFunction compress_function); + /// \brief Decoder for Arrow IPC messages /// /// This structure is intended to be allocated by the caller, diff --git a/src/nanoarrow/nanoarrow_ipc.hpp b/src/nanoarrow/nanoarrow_ipc.hpp index e6b23924f..8afdac2bd 100644 --- a/src/nanoarrow/nanoarrow_ipc.hpp +++ b/src/nanoarrow/nanoarrow_ipc.hpp @@ -129,6 +129,25 @@ inline void release_pointer(struct ArrowIpcDecompressor* data) { } } +template <> +inline void init_pointer(struct ArrowIpcCompressor* data) { + data->private_data = nullptr; + data->release = nullptr; +} + +template <> +inline void move_pointer(struct ArrowIpcCompressor* src, struct ArrowIpcCompressor* dst) { + memcpy(dst, src, sizeof(struct ArrowIpcCompressor)); + src->release = nullptr; +} + +template <> +inline void release_pointer(struct ArrowIpcCompressor* data) { + if (data->release != nullptr) { + data->release(data); + } +} + template <> inline void init_pointer(struct ArrowIpcInputStream* data) { data->release = nullptr; @@ -215,6 +234,9 @@ using UniqueEncoder = internal::Unique; /// \brief Class wrapping a unique struct ArrowIpcDecompressor using UniqueDecompressor = internal::Unique; +/// \brief Class wrapping a unique struct ArrowIpcCompressor +using UniqueCompressor = internal::Unique; + /// \brief Class wrapping a unique struct ArrowIpcInputStream using UniqueInputStream = internal::Unique; From 6c69a82f686fb5996685efcabadd03f446e52900 Mon Sep 17 00:00:00 2001 From: Pedro Holanda Date: Mon, 7 Sep 2026 14:15:09 +0200 Subject: [PATCH 02/13] feat(ipc): encode RecordBatch bodies with LZ4 or ZSTD buffer compression --- src/nanoarrow/ipc/encoder.c | 171 +++++++++++++++-- src/nanoarrow/ipc/encoder_test.cc | 299 ++++++++++++++++++++++++++++++ src/nanoarrow/nanoarrow_ipc.h | 30 +++ 3 files changed, 481 insertions(+), 19 deletions(-) diff --git a/src/nanoarrow/ipc/encoder.c b/src/nanoarrow/ipc/encoder.c index 2c12a1b57..2022b9c44 100644 --- a/src/nanoarrow/ipc/encoder.c +++ b/src/nanoarrow/ipc/encoder.c @@ -49,6 +49,12 @@ struct ArrowIpcEncoderPrivate { // Metadata to attach to the next encoded Message (in nanoarrow's packed // representation), or an empty buffer if the next Message has no metadata. struct ArrowBuffer message_metadata; + // Compression applied to the body buffers of subsequently encoded RecordBatches + enum ArrowIpcCompressionType codec; + // Compressor used when codec != NONE (release is NULL until one is needed) + struct ArrowIpcCompressor compressor; + // Whether compressor was provided by ArrowIpcEncoderSetCompressor() + int custom_compressor; }; ArrowErrorCode ArrowIpcEncoderInit(struct ArrowIpcEncoder* encoder) { @@ -69,6 +75,9 @@ ArrowErrorCode ArrowIpcEncoderInit(struct ArrowIpcEncoder* encoder) { ArrowBufferInit(&private->nodes); ArrowIpcDictionaryEncodingsInit(&private->dictionary_encodings); ArrowBufferInit(&private->message_metadata); + private->codec = NANOARROW_IPC_COMPRESSION_TYPE_NONE; + private->compressor.release = NULL; + private->custom_compressor = 0; return NANOARROW_OK; } @@ -82,6 +91,9 @@ void ArrowIpcEncoderReset(struct ArrowIpcEncoder* encoder) { ArrowBufferReset(&private->buffers); ArrowIpcDictionaryEncodingsReset(&private->dictionary_encodings); ArrowBufferReset(&private->message_metadata); + if (private->compressor.release != NULL) { + private->compressor.release(&private->compressor); + } ArrowFree(private); } memset(encoder, 0, sizeof(struct ArrowIpcEncoder)); @@ -119,6 +131,68 @@ ArrowErrorCode ArrowIpcEncoderSetMessageMetadata(struct ArrowIpcEncoder* encoder return NANOARROW_OK; } +ArrowErrorCode ArrowIpcEncoderSetCompressor(struct ArrowIpcEncoder* encoder, + struct ArrowIpcCompressor* compressor) { + NANOARROW_DCHECK(encoder != NULL && encoder->private_data != NULL && + compressor != NULL && compressor->release != NULL); + struct ArrowIpcEncoderPrivate* private = + (struct ArrowIpcEncoderPrivate*)encoder->private_data; + + if (private->compressor.release != NULL) { + private->compressor.release(&private->compressor); + } + + memcpy(&private->compressor, compressor, sizeof(struct ArrowIpcCompressor)); + compressor->release = NULL; + private->custom_compressor = 1; + return NANOARROW_OK; +} + +ArrowErrorCode ArrowIpcEncoderSetCompression( + struct ArrowIpcEncoder* encoder, enum ArrowIpcCompressionType compression_type, + struct ArrowError* error) { + NANOARROW_DCHECK(encoder != NULL && encoder->private_data != NULL); + struct ArrowIpcEncoderPrivate* private = + (struct ArrowIpcEncoderPrivate*)encoder->private_data; + + ArrowIpcCompressFunction built_in = NULL; + switch (compression_type) { + case NANOARROW_IPC_COMPRESSION_TYPE_NONE: + break; + case NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME: + built_in = ArrowIpcGetLZ4CompressionFunction(); + break; + case NANOARROW_IPC_COMPRESSION_TYPE_ZSTD: + built_in = ArrowIpcGetZstdCompressionFunction(); + break; + default: + ArrowErrorSet(error, "Unknown compression type with value %d", + (int)compression_type); + return EINVAL; + } + + if (compression_type != NANOARROW_IPC_COMPRESSION_TYPE_NONE) { + // With the default compressor, fail now rather than when the first RecordBatch is + // encoded if this build does not support the codec. A custom compressor may support + // codecs that were not built in, so it is only checked when a RecordBatch is encoded. + if (!private->custom_compressor && built_in == NULL) { + ArrowErrorSet( + error, + "Compression type with value %d not supported by this build of nanoarrow", + (int)compression_type); + return ENOTSUP; + } + + if (private->compressor.release == NULL) { + NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowIpcSerialCompressor(&private->compressor), + error); + } + } + + private->codec = compression_type; + return NANOARROW_OK; +} + static ArrowErrorCode ArrowIpcEncoderWriteContinuationAndSize(struct ArrowBuffer* out, size_t size) { _NANOARROW_CHECK_UPPER_LIMIT(size, INT32_MAX); @@ -643,37 +717,74 @@ struct ArrowIpcBufferEncoder { int64_t body_length; }; +// Append buffer_view to body_buffer as a compressed IPC buffer: the uncompressed length +// as a little-endian int64 followed by the compressed bytes. If compression does not +// reduce the size, the buffer is stored uncompressed with a length prefix of -1 instead. +static ArrowErrorCode ArrowIpcEncoderAppendCompressedBuffer( + struct ArrowIpcEncoderPrivate* private, struct ArrowBufferView buffer_view, + struct ArrowBuffer* body_buffer, struct ArrowError* error) { + NANOARROW_DCHECK(private->compressor.release != NULL); + + // placeholder for the prefix, then compress directly into the body + int64_t prefix_offset = body_buffer->size_bytes; + int64_t payload_offset = prefix_offset + (int64_t)sizeof(int64_t); + NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowBufferAppendInt64(body_buffer, 0), error); + NANOARROW_RETURN_NOT_OK(private->compressor.compress( + &private->compressor, private->codec, buffer_view, body_buffer, error)); + + int64_t prefix = buffer_view.size_bytes; + if (body_buffer->size_bytes - payload_offset >= buffer_view.size_bytes) { + body_buffer->size_bytes = payload_offset; + NANOARROW_RETURN_NOT_OK_WITH_ERROR( + ArrowBufferAppend(body_buffer, buffer_view.data.data, buffer_view.size_bytes), + error); + prefix = -1; + } + + // the prefix is always little endian + if (ArrowIpcSystemEndianness() == NANOARROW_IPC_ENDIANNESS_BIG) { + prefix = (int64_t)bswap64((uint64_t)prefix); + } + memcpy(body_buffer->data + prefix_offset, &prefix, sizeof(int64_t)); + return NANOARROW_OK; +} + static ArrowErrorCode ArrowIpcEncoderBuildContiguousBodyBufferCallback( struct ArrowBufferView buffer_view, struct ArrowIpcEncoder* encoder, struct ArrowIpcBufferEncoder* buffer_encoder, int64_t* offset, int64_t* length, struct ArrowError* error) { - NANOARROW_UNUSED(encoder); - + struct ArrowIpcEncoderPrivate* private = + (struct ArrowIpcEncoderPrivate*)encoder->private_data; struct ArrowBuffer* body_buffer = (struct ArrowBuffer*)buffer_encoder->encode_buffer_state; - int64_t old_size = body_buffer->size_bytes; - int64_t buffer_begin = _ArrowRoundUpToMultipleOf8(old_size); - int64_t buffer_end = buffer_begin + buffer_view.size_bytes; - int64_t new_size = _ArrowRoundUpToMultipleOf8(buffer_end); - - // reserve all the memory we'll need now - NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowBufferReserve(body_buffer, new_size - old_size), - error); - // zero padding up to the start of the buffer - NANOARROW_ASSERT_OK(ArrowBufferAppendFill(body_buffer, 0, buffer_begin - old_size)); + int64_t buffer_begin = _ArrowRoundUpToMultipleOf8(body_buffer->size_bytes); + NANOARROW_RETURN_NOT_OK_WITH_ERROR( + ArrowBufferAppendFill(body_buffer, 0, buffer_begin - body_buffer->size_bytes), + error); - // store offset and length of the buffer - *offset = buffer_begin; - *length = buffer_view.size_bytes; + // empty buffers are never compressed (nor length-prefixed), matching Arrow C++ + if (private->codec != NANOARROW_IPC_COMPRESSION_TYPE_NONE && + buffer_view.size_bytes > 0) { + NANOARROW_RETURN_NOT_OK( + ArrowIpcEncoderAppendCompressedBuffer(private, buffer_view, body_buffer, error)); + } else { + NANOARROW_RETURN_NOT_OK_WITH_ERROR( + ArrowBufferAppend(body_buffer, buffer_view.data.data, buffer_view.size_bytes), + error); + } - NANOARROW_ASSERT_OK( - ArrowBufferAppend(body_buffer, buffer_view.data.data, buffer_view.size_bytes)); + // store offset and length (including any prefix) of the buffer + *offset = buffer_begin; + *length = body_buffer->size_bytes - buffer_begin; // zero padding after writing the buffer - NANOARROW_DCHECK(body_buffer->size_bytes == buffer_end); - NANOARROW_ASSERT_OK(ArrowBufferAppendFill(body_buffer, 0, new_size - buffer_end)); + int64_t buffer_end = body_buffer->size_bytes; + NANOARROW_RETURN_NOT_OK_WITH_ERROR( + ArrowBufferAppendFill(body_buffer, 0, + _ArrowRoundUpToMultipleOf8(buffer_end) - buffer_end), + error); buffer_encoder->body_length = body_buffer->size_bytes; return NANOARROW_OK; @@ -744,6 +855,28 @@ static ArrowErrorCode ArrowIpcEncoderEncodeRecordBatch( FLATCC_RETURN_UNLESS_0(Message_header_RecordBatch_start(builder), error); FLATCC_RETURN_UNLESS_0(RecordBatch_length_add(builder, array_view->length), error); + if (private->codec != NANOARROW_IPC_COMPRESSION_TYPE_NONE) { + ns(CompressionType_enum_t) codec; + switch (private->codec) { + case NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME: + codec = ns(CompressionType_LZ4_FRAME); + break; + case NANOARROW_IPC_COMPRESSION_TYPE_ZSTD: + codec = ns(CompressionType_ZSTD); + break; + default: + ArrowErrorSet(error, "Unknown compression type with value %d", + (int)private->codec); + return EINVAL; + } + + FLATCC_RETURN_UNLESS_0(RecordBatch_compression_start(builder), error); + FLATCC_RETURN_UNLESS_0(BodyCompression_codec_add(builder, codec), error); + FLATCC_RETURN_UNLESS_0( + BodyCompression_method_add(builder, ns(BodyCompressionMethod_BUFFER)), error); + FLATCC_RETURN_UNLESS_0(RecordBatch_compression_end(builder), error); + } + NANOARROW_ASSERT_OK(ArrowBufferResize(&private->buffers, 0, 0)); NANOARROW_ASSERT_OK(ArrowBufferResize(&private->nodes, 0, 0)); NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderEncodeRecordBatchImpl( diff --git a/src/nanoarrow/ipc/encoder_test.cc b/src/nanoarrow/ipc/encoder_test.cc index 3eeff1ca5..0512705be 100644 --- a/src/nanoarrow/ipc/encoder_test.cc +++ b/src/nanoarrow/ipc/encoder_test.cc @@ -17,6 +17,7 @@ #include +#include #include #include #include @@ -430,3 +431,301 @@ TEST(NanoarrowIpcTest, NanoarrowIpcVisitMessageMetadataError) { EXPECT_EQ(visited, (KeyValues{{"key1", "value1"}})); EXPECT_STREQ(error.message, "visitor stopped at key1"); } + +// A record batch whose columns exercise each path of the compressed body builder: +// - "compressible": int32s with a repeating pattern +// - "with_nulls": int32s with a validity buffer +// - "incompressible": pseudo-random bytes, which are stored uncompressed (prefix -1) +// Columns without nulls have a zero-length validity buffer, which is never compressed. +class CompressibleRecordBatch { + public: + static constexpr int64_t kLength = 4096; + static constexpr int64_t kBytesPerValue = 16; + + CompressibleRecordBatch() { + NANOARROW_THROW_NOT_OK(ArrowSchemaInitFromType(schema_.get(), NANOARROW_TYPE_STRUCT)); + NANOARROW_THROW_NOT_OK(ArrowSchemaAllocateChildren(schema_.get(), 3)); + NANOARROW_THROW_NOT_OK( + ArrowSchemaInitFromType(schema_->children[0], NANOARROW_TYPE_INT32)); + NANOARROW_THROW_NOT_OK(ArrowSchemaSetName(schema_->children[0], "compressible")); + NANOARROW_THROW_NOT_OK( + ArrowSchemaInitFromType(schema_->children[1], NANOARROW_TYPE_INT32)); + NANOARROW_THROW_NOT_OK(ArrowSchemaSetName(schema_->children[1], "with_nulls")); + NANOARROW_THROW_NOT_OK( + ArrowSchemaInitFromType(schema_->children[2], NANOARROW_TYPE_BINARY)); + NANOARROW_THROW_NOT_OK(ArrowSchemaSetName(schema_->children[2], "incompressible")); + + NANOARROW_THROW_NOT_OK( + ArrowArrayInitFromSchema(array_.get(), schema_.get(), nullptr)); + NANOARROW_THROW_NOT_OK(ArrowArrayStartAppending(array_.get())); + + uint32_t state = 2463534242u; + uint8_t random_bytes[kBytesPerValue]; + for (int64_t i = 0; i < kLength; i++) { + NANOARROW_THROW_NOT_OK(ArrowArrayAppendInt(array_->children[0], i % 4)); + + if (i % 3 == 0) { + NANOARROW_THROW_NOT_OK(ArrowArrayAppendNull(array_->children[1], 1)); + } else { + NANOARROW_THROW_NOT_OK(ArrowArrayAppendInt(array_->children[1], i)); + } + + // xorshift32 so that the bytes are deterministic but not compressible + for (int64_t j = 0; j < kBytesPerValue; j += 4) { + state ^= state << 13; + state ^= state >> 17; + state ^= state << 5; + std::memcpy(random_bytes + j, &state, sizeof(state)); + } + struct ArrowBufferView bytes = {{random_bytes}, kBytesPerValue}; + NANOARROW_THROW_NOT_OK(ArrowArrayAppendBytes(array_->children[2], bytes)); + + NANOARROW_THROW_NOT_OK(ArrowArrayFinishElement(array_.get())); + } + + NANOARROW_THROW_NOT_OK(ArrowArrayFinishBuildingDefault(array_.get(), nullptr)); + NANOARROW_THROW_NOT_OK( + ArrowArrayViewInitFromSchema(array_view_.get(), schema_.get(), nullptr)); + NANOARROW_THROW_NOT_OK( + ArrowArrayViewSetArray(array_view_.get(), array_.get(), nullptr)); + } + + struct ArrowSchema* schema() { return schema_.get(); } + const struct ArrowArrayView* array_view() { return array_view_.get(); } + + private: + nanoarrow::UniqueSchema schema_; + nanoarrow::UniqueArray array_; + nanoarrow::UniqueArrayView array_view_; +}; + +static void AssertArrayViewsEqual(const struct ArrowArrayView* expected, + const struct ArrowArrayView* actual) { + ASSERT_EQ(actual->length, expected->length); + ASSERT_EQ(actual->null_count, expected->null_count); + ASSERT_EQ(actual->n_children, expected->n_children); + + for (int i = 0; i < NANOARROW_MAX_FIXED_BUFFERS; i++) { + SCOPED_TRACE("buffer " + std::to_string(i)); + ASSERT_EQ(actual->buffer_views[i].size_bytes, expected->buffer_views[i].size_bytes); + if (expected->buffer_views[i].size_bytes > 0) { + EXPECT_EQ(std::memcmp(actual->buffer_views[i].data.data, + expected->buffer_views[i].data.data, + expected->buffer_views[i].size_bytes), + 0); + } + } + + for (int64_t i = 0; i < expected->n_children; i++) { + SCOPED_TRACE("child " + std::to_string(i)); + AssertArrayViewsEqual(expected->children[i], actual->children[i]); + } +} + +static int64_t ReadLittleEndianInt64(const uint8_t* data) { + int64_t value; + std::memcpy(&value, data, sizeof(value)); + if (ArrowIpcSystemEndianness() == NANOARROW_IPC_ENDIANNESS_BIG) { + value = static_cast(bswap64(static_cast(value))); + } + return value; +} + +static void TestCompressedRecordBatchRoundtrip(enum ArrowIpcCompressionType codec) { + nanoarrow::ipc::UniqueEncoder encoder; + ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); + nanoarrow::ipc::UniqueDecoder decoder; + ASSERT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK); + + CompressibleRecordBatch batch; + struct ArrowError error; + ASSERT_EQ(ArrowIpcDecoderSetSchema(decoder.get(), batch.schema(), &error), NANOARROW_OK) + << error.message; + + // Encode without compression for reference + nanoarrow::UniqueBuffer uncompressed_message, uncompressed_body; + ASSERT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), batch.array_view(), + uncompressed_body.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/true, + uncompressed_message.get()), + NANOARROW_OK); + + ASSERT_EQ(ArrowIpcEncoderSetCompression(encoder.get(), codec, &error), NANOARROW_OK) + << error.message; + + nanoarrow::UniqueBuffer message, body; + ASSERT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), batch.array_view(), + body.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ( + ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/true, message.get()), + NANOARROW_OK); + + // The compressible column should have made the body smaller, and the body must + // still be padded to a multiple of 8 bytes + EXPECT_LT(body->size_bytes, uncompressed_body->size_bytes); + EXPECT_EQ(body->size_bytes % 8, 0); + + // The first buffer in the body is the data buffer of "compressible" (its validity + // buffer is empty and takes no space). It should be prefixed with its uncompressed + // length. + const int64_t int32_data_size = CompressibleRecordBatch::kLength * sizeof(int32_t); + EXPECT_EQ(ReadLittleEndianInt64(body->data), int32_data_size); + + // The last buffer in the body is the data buffer of "incompressible", which should + // have been stored uncompressed with a prefix of -1 (and is a multiple of 8 bytes, + // so ends exactly at the end of the body). + const int64_t binary_data_size = + CompressibleRecordBatch::kLength * CompressibleRecordBatch::kBytesPerValue; + const uint8_t* last_buffer = body->data + body->size_bytes - binary_data_size - 8; + EXPECT_EQ(ReadLittleEndianInt64(last_buffer), -1); + EXPECT_EQ(std::memcmp(last_buffer + 8, + batch.array_view()->children[2]->buffer_views[2].data.data, + binary_data_size), + 0); + + // Decode the header: the codec is recorded and the body length is correct + struct ArrowBufferView message_view = {{message->data}, message->size_bytes}; + ASSERT_EQ(ArrowIpcDecoderVerifyHeader(decoder.get(), message_view, &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcDecoderDecodeHeader(decoder.get(), message_view, &error), + NANOARROW_OK) + << error.message; + EXPECT_EQ(decoder->message_type, NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH); + EXPECT_EQ(decoder->codec, codec); + EXPECT_EQ(decoder->body_size_bytes, body->size_bytes); + + // Decode the body and compare with the original + nanoarrow::UniqueArray decoded; + struct ArrowBufferView body_view = {{body->data}, body->size_bytes}; + ASSERT_EQ(ArrowIpcDecoderDecodeArray(decoder.get(), body_view, -1, decoded.get(), + NANOARROW_VALIDATION_LEVEL_FULL, &error), + NANOARROW_OK) + << error.message; + + nanoarrow::UniqueArrayView decoded_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(decoded_view.get(), batch.schema(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayViewSetArray(decoded_view.get(), decoded.get(), &error), + NANOARROW_OK) + << error.message; + AssertArrayViewsEqual(batch.array_view(), decoded_view.get()); + + // Compression can be turned off again + ASSERT_EQ(ArrowIpcEncoderSetCompression(encoder.get(), + NANOARROW_IPC_COMPRESSION_TYPE_NONE, &error), + NANOARROW_OK) + << error.message; + message->size_bytes = 0; + body->size_bytes = 0; + ASSERT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), batch.array_view(), + body.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ( + ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/true, message.get()), + NANOARROW_OK); + EXPECT_EQ(body->size_bytes, uncompressed_body->size_bytes); + EXPECT_EQ(std::memcmp(body->data, uncompressed_body->data, body->size_bytes), 0); + + message_view = {{message->data}, message->size_bytes}; + ASSERT_EQ(ArrowIpcDecoderDecodeHeader(decoder.get(), message_view, &error), + NANOARROW_OK) + << error.message; + EXPECT_EQ(decoder->codec, NANOARROW_IPC_COMPRESSION_TYPE_NONE); +} + +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderCompressedRecordBatchLZ4) { + if (ArrowIpcGetLZ4CompressionFunction() == nullptr) { + GTEST_SKIP() << "nanoarrow_ipc not built with NANOARROW_IPC_WITH_LZ4"; + } + TestCompressedRecordBatchRoundtrip(NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME); +} + +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderCompressedRecordBatchZstd) { + if (ArrowIpcGetZstdCompressionFunction() == nullptr) { + GTEST_SKIP() << "nanoarrow_ipc not built with NANOARROW_IPC_WITH_ZSTD"; + } + TestCompressedRecordBatchRoundtrip(NANOARROW_IPC_COMPRESSION_TYPE_ZSTD); +} + +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSetCompressionErrors) { + nanoarrow::ipc::UniqueEncoder encoder; + ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); + struct ArrowError error; + + EXPECT_EQ(ArrowIpcEncoderSetCompression( + encoder.get(), static_cast(99), &error), + EINVAL); + EXPECT_STREQ(error.message, "Unknown compression type with value 99"); + + // NONE is always supported + EXPECT_EQ(ArrowIpcEncoderSetCompression(encoder.get(), + NANOARROW_IPC_COMPRESSION_TYPE_NONE, &error), + NANOARROW_OK) + << error.message; + + // Codecs that were not built in are rejected when they are set rather than when + // the first batch is encoded +#if defined(NANOARROW_IPC_WITH_LZ4) + EXPECT_EQ(ArrowIpcEncoderSetCompression( + encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, &error), + NANOARROW_OK) + << error.message; +#else + EXPECT_EQ(ArrowIpcEncoderSetCompression( + encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, &error), + ENOTSUP); + EXPECT_STREQ(error.message, + "Compression type with value 1 not supported by this build of nanoarrow"); +#endif + +#if defined(NANOARROW_IPC_WITH_ZSTD) + EXPECT_EQ(ArrowIpcEncoderSetCompression(encoder.get(), + NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, &error), + NANOARROW_OK) + << error.message; +#else + EXPECT_EQ(ArrowIpcEncoderSetCompression(encoder.get(), + NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, &error), + ENOTSUP); + EXPECT_STREQ(error.message, + "Compression type with value 2 not supported by this build of nanoarrow"); +#endif +} + +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSetCompressor) { + nanoarrow::ipc::UniqueEncoder encoder; + ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); + struct ArrowError error; + + // A custom compressor that explicitly does not support LZ4 + nanoarrow::ipc::UniqueCompressor compressor; + ASSERT_EQ(ArrowIpcSerialCompressor(compressor.get()), NANOARROW_OK); + ASSERT_EQ(ArrowIpcSerialCompressorSetFunction( + compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, nullptr), + NANOARROW_OK); + + ASSERT_EQ(ArrowIpcEncoderSetCompressor(encoder.get(), compressor.get()), NANOARROW_OK); + // The encoder took ownership of the compressor + EXPECT_EQ(compressor->release, nullptr); + + // With a custom compressor, support is not checked until a batch is encoded + ASSERT_EQ(ArrowIpcEncoderSetCompression( + encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, &error), + NANOARROW_OK) + << error.message; + + CompressibleRecordBatch batch; + nanoarrow::UniqueBuffer body; + EXPECT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), batch.array_view(), + body.get(), &error), + ENOTSUP); + EXPECT_STREQ(error.message, + "Compression type with value 1 not supported by this build of nanoarrow"); +} diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index 797019120..1d6ea0b7d 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -100,6 +100,10 @@ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcEncoderFinalizeBuffer) #define ArrowIpcEncoderSetMessageMetadata \ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcEncoderSetMessageMetadata) +#define ArrowIpcEncoderSetCompression \ + NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcEncoderSetCompression) +#define ArrowIpcEncoderSetCompressor \ + NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcEncoderSetCompressor) #define ArrowIpcEncoderEncodeSchema \ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcEncoderEncodeSchema) #define ArrowIpcEncoderEncodeSimpleRecordBatch \ @@ -931,6 +935,32 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcEncoderSetMessageMetadata(struct ArrowIpcEncoder* encoder, struct ArrowBuffer* metadata, struct ArrowError* error); +/// \brief Set the buffer compression used by subsequently encoded RecordBatch messages +/// +/// When compression_type is not NANOARROW_IPC_COMPRESSION_TYPE_NONE, the body buffers +/// of every RecordBatch encoded after this call are compressed with the given codec as +/// described by the Arrow IPC format: each non-empty buffer is written as its +/// uncompressed length (a little-endian int64) followed by the compressed bytes. +/// Buffers that do not shrink when compressed are written uncompressed with a length +/// prefix of -1, and empty buffers are written as-is. The setting persists until it is +/// changed and does not affect Schema messages. +/// +/// Returns EINVAL for an unknown compression type and ENOTSUP if the compression type +/// is not supported by this build of nanoarrow (i.e., nanoarrow was not built with +/// NANOARROW_IPC_WITH_LZ4 or NANOARROW_IPC_WITH_ZSTD). If a custom compressor was +/// set with ArrowIpcEncoderSetCompressor(), support is not checked until a RecordBatch +/// is encoded. +NANOARROW_DLL ArrowErrorCode ArrowIpcEncoderSetCompression( + struct ArrowIpcEncoder* encoder, enum ArrowIpcCompressionType compression_type, + struct ArrowError* error); + +/// \brief Set the compressor implementation used by this encoder +/// +/// The encoder takes ownership of compressor. If this is not called, an +/// ArrowIpcSerialCompressor() is used when compression is first required. +NANOARROW_DLL ArrowErrorCode ArrowIpcEncoderSetCompressor( + struct ArrowIpcEncoder* encoder, struct ArrowIpcCompressor* compressor); + /// \brief Encode an ArrowSchema /// /// Returns ENOMEM if allocation fails, NANOARROW_OK otherwise. From 3356ee81278ca3e7309c2a9df245cf912edafb3b Mon Sep 17 00:00:00 2001 From: Pedro Holanda Date: Mon, 7 Sep 2026 14:18:51 +0200 Subject: [PATCH 03/13] feat(ipc): add ArrowIpcWriterSetCompression() for LZ4/ZSTD compressed streams and files --- src/nanoarrow/ipc/files_test.cc | 77 ++++++++++------ src/nanoarrow/ipc/writer.c | 9 ++ src/nanoarrow/ipc/writer_test.cc | 150 +++++++++++++++++++++++++++++++ src/nanoarrow/nanoarrow_ipc.h | 13 +++ 4 files changed, 222 insertions(+), 27 deletions(-) diff --git a/src/nanoarrow/ipc/files_test.cc b/src/nanoarrow/ipc/files_test.cc index 46c7203c1..f5f77ddca 100644 --- a/src/nanoarrow/ipc/files_test.cc +++ b/src/nanoarrow/ipc/files_test.cc @@ -18,6 +18,9 @@ #include #include #include +#include +#include +#include #include @@ -212,6 +215,7 @@ class TestFile { ArrowErrorCode WriteNanoarrowStream(const nanoarrow::UniqueSchema& schema, const std::vector& arrays, + enum ArrowIpcCompressionType codec, struct ArrowBuffer* buffer, struct ArrowError* error) { nanoarrow::ipc::UniqueOutputStream output_stream; @@ -219,6 +223,7 @@ class TestFile { nanoarrow::ipc::UniqueWriter writer; NANOARROW_RETURN_NOT_OK(ArrowIpcWriterInit(writer.get(), output_stream.get())); + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterSetCompression(writer.get(), codec, error)); nanoarrow::UniqueArrayView array_view; NANOARROW_RETURN_NOT_OK( @@ -259,14 +264,27 @@ class TestFile { GTEST_FAIL() << MakeError(NANOARROW_OK, ""); } - // Write back to a buffer using nanoarrow if supported. We do this here - // because we need to move the arrays into the comparison for the Arrow C++ - // read. - nanoarrow::UniqueBuffer roundtripped; + // Write back to a buffer using nanoarrow if supported: once uncompressed and once + // with each compression codec available in this build. We do this here because we + // need to move the arrays into the comparison for the Arrow C++ read. + std::vector> codecs = { + {"uncompressed", NANOARROW_IPC_COMPRESSION_TYPE_NONE}}; + if (ArrowIpcGetLZ4CompressionFunction() != nullptr) { + codecs.emplace_back("lz4", NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME); + } + if (ArrowIpcGetZstdCompressionFunction() != nullptr) { + codecs.emplace_back("zstd", NANOARROW_IPC_COMPRESSION_TYPE_ZSTD); + } + + std::vector roundtripped(codecs.size()); if (write_supported_) { - ASSERT_EQ(WriteNanoarrowStream(schema, arrays, roundtripped.get(), &error), - NANOARROW_OK) - << error.message; + for (size_t i = 0; i < codecs.size(); i++) { + SCOPED_TRACE("Write the " + codecs[i].first + " stream using nanoarrow"); + ASSERT_EQ(WriteNanoarrowStream(schema, arrays, codecs[i].second, + roundtripped[i].get(), &error), + NANOARROW_OK) + << error.message; + } } // Read the same file with Arrow C++ @@ -283,28 +301,33 @@ class TestFile { return; } - auto maybe_table_roundtripped = ReadTable(BufferInputStream(roundtripped.get())); - { - SCOPED_TRACE("Read the roundtripped buffer using Arrow C++"); - FAIL_RESULT_NOT_OK(maybe_table_roundtripped); + for (size_t i = 0; i < codecs.size(); i++) { + SCOPED_TRACE("Roundtrip of the " + codecs[i].first + " stream"); - AssertEqualsTable(maybe_table_roundtripped.ValueUnsafe(), - maybe_table_arrow.ValueUnsafe()); - } + auto maybe_table_roundtripped = ReadTable(BufferInputStream(roundtripped[i].get())); + { + SCOPED_TRACE("Read the roundtripped buffer using Arrow C++"); + FAIL_RESULT_NOT_OK(maybe_table_roundtripped); - nanoarrow::UniqueSchema roundtripped_schema; - std::vector roundtripped_arrays; - { - SCOPED_TRACE("Read the roundtripped buffer using nanoarrow"); - nanoarrow::UniqueArrayStream array_stream; - ASSERT_EQ(GetArrowArrayStreamIPC(roundtripped.get(), array_stream.get(), &error), - NANOARROW_OK); - ASSERT_EQ(ReadArrowArrayStreamIPC(array_stream.get(), roundtripped_schema.get(), - &roundtripped_arrays, &error), - NANOARROW_OK); - - AssertEqualsTable(std::move(roundtripped_schema), std::move(roundtripped_arrays), - maybe_table_arrow.ValueUnsafe()); + AssertEqualsTable(maybe_table_roundtripped.ValueUnsafe(), + maybe_table_arrow.ValueUnsafe()); + } + + nanoarrow::UniqueSchema roundtripped_schema; + std::vector roundtripped_arrays; + { + SCOPED_TRACE("Read the roundtripped buffer using nanoarrow"); + nanoarrow::UniqueArrayStream array_stream; + ASSERT_EQ( + GetArrowArrayStreamIPC(roundtripped[i].get(), array_stream.get(), &error), + NANOARROW_OK); + ASSERT_EQ(ReadArrowArrayStreamIPC(array_stream.get(), roundtripped_schema.get(), + &roundtripped_arrays, &error), + NANOARROW_OK); + + AssertEqualsTable(std::move(roundtripped_schema), std::move(roundtripped_arrays), + maybe_table_arrow.ValueUnsafe()); + } } } diff --git a/src/nanoarrow/ipc/writer.c b/src/nanoarrow/ipc/writer.c index 1f100865b..2aab91d37 100644 --- a/src/nanoarrow/ipc/writer.c +++ b/src/nanoarrow/ipc/writer.c @@ -226,6 +226,15 @@ void ArrowIpcWriterReset(struct ArrowIpcWriter* writer) { memset(writer, 0, sizeof(struct ArrowIpcWriter)); } +ArrowErrorCode ArrowIpcWriterSetCompression(struct ArrowIpcWriter* writer, + enum ArrowIpcCompressionType compression_type, + struct ArrowError* error) { + NANOARROW_DCHECK(writer != NULL && writer->private_data != NULL); + struct ArrowIpcWriterPrivate* private = + (struct ArrowIpcWriterPrivate*)writer->private_data; + return ArrowIpcEncoderSetCompression(&private->encoder, compression_type, error); +} + static struct ArrowBufferView ArrowBufferToBufferView(const struct ArrowBuffer* buffer) { struct ArrowBufferView buffer_view = { .data.as_uint8 = buffer->data, diff --git a/src/nanoarrow/ipc/writer_test.cc b/src/nanoarrow/ipc/writer_test.cc index a07ae516c..5a74b3ce9 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -204,3 +204,153 @@ TEST(NanoarrowIpcWriter, FileWriting) { auto after_footer = p->bytes_written; EXPECT_GT(after_footer, after_eos); } + +// A struct array with a single int32 column of repeating values (i.e., compressible) +static constexpr int64_t kCompressibleBatchLength = 1024; + +static void InitCompressibleBatch(struct ArrowSchema* schema, struct ArrowArray* array) { + ASSERT_EQ(ArrowSchemaInitFromType(schema, NANOARROW_TYPE_STRUCT), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateChildren(schema, 1), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaInitFromType(schema->children[0], NANOARROW_TYPE_INT32), + NANOARROW_OK); + ASSERT_EQ(ArrowSchemaSetName(schema->children[0], "col"), NANOARROW_OK); + + ASSERT_EQ(ArrowArrayInitFromSchema(array, schema, nullptr), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(array), NANOARROW_OK); + for (int64_t i = 0; i < kCompressibleBatchLength; i++) { + ASSERT_EQ(ArrowArrayAppendInt(array->children[0], i % 8), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishElement(array), NANOARROW_OK); + } + ASSERT_EQ(ArrowArrayFinishBuildingDefault(array, nullptr), NANOARROW_OK); +} + +// Write schema + batch + EOS (optionally as an IPC file) to output using codec +static void WriteCompressibleBatch(enum ArrowIpcCompressionType codec, bool as_file, + struct ArrowBuffer* output) { + struct ArrowError error; + + nanoarrow::UniqueSchema schema; + nanoarrow::UniqueArray array; + ASSERT_NO_FATAL_FAILURE(InitCompressibleBatch(schema.get(), array.get())); + nanoarrow::UniqueArrayView array_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(array_view.get(), schema.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), array.get(), &error), NANOARROW_OK) + << error.message; + + nanoarrow::ipc::UniqueOutputStream stream; + ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream.get(), output), NANOARROW_OK); + nanoarrow::ipc::UniqueWriter writer; + ASSERT_EQ(ArrowIpcWriterInit(writer.get(), stream.get()), NANOARROW_OK); + ASSERT_EQ(ArrowIpcWriterSetCompression(writer.get(), codec, &error), NANOARROW_OK) + << error.message; + + if (as_file) { + ASSERT_EQ(ArrowIpcWriterStartFile(writer.get(), &error), NANOARROW_OK) + << error.message; + } + ASSERT_EQ(ArrowIpcWriterWriteSchema(writer.get(), schema.get(), &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcWriterWriteArrayView(writer.get(), array_view.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcWriterWriteArrayView(writer.get(), nullptr, &error), NANOARROW_OK) + << error.message; + + if (as_file) { + // The block for the record batch records the (compressed) body length + auto* p = static_cast(writer->private_data); + ASSERT_EQ(p->footer.record_batch_blocks.size_bytes, sizeof(struct ArrowIpcFileBlock)); + auto* block = + reinterpret_cast(p->footer.record_batch_blocks.data); + EXPECT_EQ(block->body_length, p->body_buffer.size_bytes); + ASSERT_EQ(ArrowIpcWriterFinalizeFile(writer.get(), &error), NANOARROW_OK) + << error.message; + } +} + +// Read the stream starting at offset back with the array stream reader and check that +// the values match what InitCompressibleBatch() produced +static void CheckCompressibleBatch(const struct ArrowBuffer* output, int64_t offset) { + struct ArrowError error; + + nanoarrow::UniqueBuffer input_buffer; + ASSERT_EQ(ArrowBufferAppend(input_buffer.get(), output->data + offset, + output->size_bytes - offset), + NANOARROW_OK); + nanoarrow::ipc::UniqueInputStream input; + ASSERT_EQ(ArrowIpcInputStreamInitBuffer(input.get(), input_buffer.get()), NANOARROW_OK); + nanoarrow::UniqueArrayStream stream; + ASSERT_EQ(ArrowIpcArrayStreamReaderInit(stream.get(), input.get(), nullptr), + NANOARROW_OK); + + nanoarrow::UniqueSchema schema; + ASSERT_EQ(ArrowArrayStreamGetSchema(stream.get(), schema.get(), &error), NANOARROW_OK) + << error.message; + EXPECT_STREQ(schema->format, "+s"); + + nanoarrow::UniqueArray array; + ASSERT_EQ(ArrowArrayStreamGetNext(stream.get(), array.get(), &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(array->length, kCompressibleBatchLength); + + nanoarrow::UniqueArrayView array_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(array_view.get(), schema.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), array.get(), &error), NANOARROW_OK) + << error.message; + for (int64_t i = 0; i < kCompressibleBatchLength; i++) { + ASSERT_EQ(ArrowArrayViewGetIntUnsafe(array_view->children[0], i), i % 8); + } + + nanoarrow::UniqueArray eos; + ASSERT_EQ(ArrowArrayStreamGetNext(stream.get(), eos.get(), &error), NANOARROW_OK) + << error.message; + EXPECT_EQ(eos->release, nullptr); +} + +static void TestCompressedWriting(enum ArrowIpcCompressionType codec) { + for (bool as_file : {false, true}) { + SCOPED_TRACE(as_file ? "file" : "stream"); + + nanoarrow::UniqueBuffer uncompressed, compressed; + ASSERT_NO_FATAL_FAILURE(WriteCompressibleBatch(NANOARROW_IPC_COMPRESSION_TYPE_NONE, + as_file, uncompressed.get())); + ASSERT_NO_FATAL_FAILURE(WriteCompressibleBatch(codec, as_file, compressed.get())); + EXPECT_LT(compressed->size_bytes, uncompressed->size_bytes); + + // The stream portion of a file follows the padded magic + int64_t offset = as_file ? sizeof(NANOARROW_IPC_FILE_PADDED_MAGIC) : 0; + ASSERT_NO_FATAL_FAILURE(CheckCompressibleBatch(compressed.get(), offset)); + } +} + +TEST(NanoarrowIpcWriter, CompressedWritingLZ4) { + if (ArrowIpcGetLZ4CompressionFunction() == nullptr) { + GTEST_SKIP() << "nanoarrow_ipc not built with NANOARROW_IPC_WITH_LZ4"; + } + TestCompressedWriting(NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME); +} + +TEST(NanoarrowIpcWriter, CompressedWritingZstd) { + if (ArrowIpcGetZstdCompressionFunction() == nullptr) { + GTEST_SKIP() << "nanoarrow_ipc not built with NANOARROW_IPC_WITH_ZSTD"; + } + TestCompressedWriting(NANOARROW_IPC_COMPRESSION_TYPE_ZSTD); +} + +TEST(NanoarrowIpcWriter, SetCompressionErrors) { + nanoarrow::UniqueBuffer output; + nanoarrow::ipc::UniqueOutputStream stream; + ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream.get(), output.get()), NANOARROW_OK); + nanoarrow::ipc::UniqueWriter writer; + ASSERT_EQ(ArrowIpcWriterInit(writer.get(), stream.get()), NANOARROW_OK); + + struct ArrowError error; + EXPECT_EQ(ArrowIpcWriterSetCompression( + writer.get(), static_cast(99), &error), + EINVAL); + EXPECT_STREQ(error.message, "Unknown compression type with value 99"); +} diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index 1d6ea0b7d..25d26ca5e 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -118,6 +118,8 @@ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcOutputStreamMove) #define ArrowIpcWriterInit NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcWriterInit) #define ArrowIpcWriterReset NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcWriterReset) +#define ArrowIpcWriterSetCompression \ + NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcWriterSetCompression) #define ArrowIpcWriterWriteSchema \ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcWriterWriteSchema) #define ArrowIpcWriterWriteArrayView \ @@ -1044,6 +1046,17 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcWriterInit( /// \brief Release all resources attached to a writer NANOARROW_DLL void ArrowIpcWriterReset(struct ArrowIpcWriter* writer); +/// \brief Set the buffer compression used for subsequently written record batches +/// +/// See ArrowIpcEncoderSetCompression(). Compression applies to record batches written +/// after this call (in both stream and file mode) and may be changed between batches. +/// +/// Returns EINVAL for an unknown compression type and ENOTSUP if the compression type +/// is not supported by this build of nanoarrow. +NANOARROW_DLL ArrowErrorCode ArrowIpcWriterSetCompression( + struct ArrowIpcWriter* writer, enum ArrowIpcCompressionType compression_type, + struct ArrowError* error); + /// \brief Write a schema to the output byte stream /// /// Errors are propagated from the underlying encoder and output byte stream. From 75be6c6e332b97416112413d0a413f00a3af176e Mon Sep 17 00:00:00 2001 From: Pedro Holanda Date: Tue, 8 Sep 2026 14:14:44 +0200 Subject: [PATCH 04/13] fix(ipc): handle compressor allocation failures and reserve buffer padding --- src/nanoarrow/ipc/codecs.c | 5 +- src/nanoarrow/ipc/encoder.c | 16 +++++-- src/nanoarrow/ipc/encoder_test.cc | 77 +++++++++++++++++++++---------- 3 files changed, 68 insertions(+), 30 deletions(-) diff --git a/src/nanoarrow/ipc/codecs.c b/src/nanoarrow/ipc/codecs.c index cf8d552af..0f2aa36cf 100644 --- a/src/nanoarrow/ipc/codecs.c +++ b/src/nanoarrow/ipc/codecs.c @@ -304,8 +304,7 @@ static void ArrowIpcSerialCompressorRelease(struct ArrowIpcCompressor* compresso } ArrowErrorCode ArrowIpcSerialCompressor(struct ArrowIpcCompressor* compressor) { - compressor->compress = &ArrowIpcSerialCompressorCompress; - compressor->release = &ArrowIpcSerialCompressorRelease; + compressor->release = NULL; compressor->private_data = ArrowMalloc(sizeof(struct ArrowIpcSerialCompressorPrivate)); if (compressor->private_data == NULL) { return ENOMEM; @@ -317,6 +316,8 @@ ArrowErrorCode ArrowIpcSerialCompressor(struct ArrowIpcCompressor* compressor) { ArrowIpcSerialCompressorSetFunction(compressor, NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, ArrowIpcGetLZ4CompressionFunction()); + compressor->compress = &ArrowIpcSerialCompressorCompress; + compressor->release = &ArrowIpcSerialCompressorRelease; return NANOARROW_OK; } diff --git a/src/nanoarrow/ipc/encoder.c b/src/nanoarrow/ipc/encoder.c index 2022b9c44..454141e0c 100644 --- a/src/nanoarrow/ipc/encoder.c +++ b/src/nanoarrow/ipc/encoder.c @@ -758,15 +758,23 @@ static ArrowErrorCode ArrowIpcEncoderBuildContiguousBodyBufferCallback( struct ArrowBuffer* body_buffer = (struct ArrowBuffer*)buffer_encoder->encode_buffer_state; - // zero padding up to the start of the buffer int64_t buffer_begin = _ArrowRoundUpToMultipleOf8(body_buffer->size_bytes); + // Empty buffers are never compressed (nor length-prefixed), matching Arrow C++. + int needs_compression = + private->codec != NANOARROW_IPC_COMPRESSION_TYPE_NONE && buffer_view.size_bytes > 0; + if (!needs_compression) { + // Reserve the data and padding together to avoid growing the buffer twice. + int64_t new_size = _ArrowRoundUpToMultipleOf8(buffer_begin + buffer_view.size_bytes); + NANOARROW_RETURN_NOT_OK_WITH_ERROR( + ArrowBufferReserve(body_buffer, new_size - body_buffer->size_bytes), error); + } + + // zero padding up to the start of the buffer NANOARROW_RETURN_NOT_OK_WITH_ERROR( ArrowBufferAppendFill(body_buffer, 0, buffer_begin - body_buffer->size_bytes), error); - // empty buffers are never compressed (nor length-prefixed), matching Arrow C++ - if (private->codec != NANOARROW_IPC_COMPRESSION_TYPE_NONE && - buffer_view.size_bytes > 0) { + if (needs_compression) { NANOARROW_RETURN_NOT_OK( ArrowIpcEncoderAppendCompressedBuffer(private, buffer_view, body_buffer, error)); } else { diff --git a/src/nanoarrow/ipc/encoder_test.cc b/src/nanoarrow/ipc/encoder_test.cc index 0512705be..4c640a7a8 100644 --- a/src/nanoarrow/ipc/encoder_test.cc +++ b/src/nanoarrow/ipc/encoder_test.cc @@ -499,29 +499,6 @@ class CompressibleRecordBatch { nanoarrow::UniqueArrayView array_view_; }; -static void AssertArrayViewsEqual(const struct ArrowArrayView* expected, - const struct ArrowArrayView* actual) { - ASSERT_EQ(actual->length, expected->length); - ASSERT_EQ(actual->null_count, expected->null_count); - ASSERT_EQ(actual->n_children, expected->n_children); - - for (int i = 0; i < NANOARROW_MAX_FIXED_BUFFERS; i++) { - SCOPED_TRACE("buffer " + std::to_string(i)); - ASSERT_EQ(actual->buffer_views[i].size_bytes, expected->buffer_views[i].size_bytes); - if (expected->buffer_views[i].size_bytes > 0) { - EXPECT_EQ(std::memcmp(actual->buffer_views[i].data.data, - expected->buffer_views[i].data.data, - expected->buffer_views[i].size_bytes), - 0); - } - } - - for (int64_t i = 0; i < expected->n_children; i++) { - SCOPED_TRACE("child " + std::to_string(i)); - AssertArrayViewsEqual(expected->children[i], actual->children[i]); - } -} - static int64_t ReadLittleEndianInt64(const uint8_t* data) { int64_t value; std::memcpy(&value, data, sizeof(value)); @@ -614,7 +591,11 @@ static void TestCompressedRecordBatchRoundtrip(enum ArrowIpcCompressionType code ASSERT_EQ(ArrowArrayViewSetArray(decoded_view.get(), decoded.get(), &error), NANOARROW_OK) << error.message; - AssertArrayViewsEqual(batch.array_view(), decoded_view.get()); + int is_equal = 0; + ASSERT_EQ(ArrowArrayViewCompare(decoded_view.get(), batch.array_view(), + NANOARROW_COMPARE_IDENTICAL, &is_equal, &error), + NANOARROW_OK); + EXPECT_EQ(is_equal, 1) << error.message; // Compression can be turned off again ASSERT_EQ(ArrowIpcEncoderSetCompression(encoder.get(), @@ -640,6 +621,54 @@ static void TestCompressedRecordBatchRoundtrip(enum ArrowIpcCompressionType code EXPECT_EQ(decoder->codec, NANOARROW_IPC_COMPRESSION_TYPE_NONE); } +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderUncompressedRecordBatchAllocation) { + nanoarrow::ipc::UniqueEncoder encoder; + ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); + + // An odd number of int32 values requires four bytes of trailing padding. + std::vector values(1025, 42); + struct ArrowError error; + nanoarrow::UniqueSchema schema; + ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_STRUCT), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateChildren(schema.get(), 1), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaInitFromType(schema->children[0], NANOARROW_TYPE_INT32), + NANOARROW_OK); + nanoarrow::UniqueArray array; + ASSERT_EQ(ArrowArrayInitFromSchema(array.get(), schema.get(), &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(array.get()), NANOARROW_OK); + for (int32_t value : values) { + ASSERT_EQ(ArrowArrayAppendInt(array->children[0], value), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishElement(array.get()), NANOARROW_OK); + } + ASSERT_EQ(ArrowArrayFinishBuildingDefault(array.get(), &error), NANOARROW_OK); + nanoarrow::UniqueArrayView array_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(array_view.get(), schema.get(), &error), + NANOARROW_OK); + ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), array.get(), &error), NANOARROW_OK); + + int allocations = 0; + auto allocator = ArrowBufferAllocatorDefault(); + allocator.private_data = &allocations; + allocator.reallocate = [](struct ArrowBufferAllocator* allocator, uint8_t* ptr, + int64_t old_size, int64_t new_size) { + ++*static_cast(allocator->private_data); + auto default_allocator = ArrowBufferAllocatorDefault(); + return default_allocator.reallocate(&default_allocator, ptr, old_size, new_size); + }; + nanoarrow::UniqueBuffer body; + ASSERT_EQ(ArrowBufferSetAllocator(body.get(), allocator), NANOARROW_OK); + + ASSERT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), array_view.get(), + body.get(), &error), + NANOARROW_OK) + << error.message; + EXPECT_EQ(allocations, 1); + ASSERT_EQ(body->size_bytes, 4104); + EXPECT_EQ(body->capacity_bytes, 4104); + EXPECT_EQ(std::memcmp(body->data, values.data(), 4100), 0); + EXPECT_EQ(std::memcmp(body->data + 4100, "\0\0\0\0", 4), 0); +} + TEST(NanoarrowIpcTest, NanoarrowIpcEncoderCompressedRecordBatchLZ4) { if (ArrowIpcGetLZ4CompressionFunction() == nullptr) { GTEST_SKIP() << "nanoarrow_ipc not built with NANOARROW_IPC_WITH_LZ4"; From 6ee135bf5c258e9463a5486c1a0ed3e79c6a5a0b Mon Sep 17 00:00:00 2001 From: Pedro Holanda Date: Tue, 8 Sep 2026 14:51:18 +0200 Subject: [PATCH 05/13] feat(ipc): add a compression level parameter to the IPC compression API --- src/nanoarrow/ipc/codecs.c | 15 ++-- src/nanoarrow/ipc/codecs_test.cc | 80 ++++++++++++++----- src/nanoarrow/ipc/encoder.c | 9 ++- src/nanoarrow/ipc/encoder_test.cc | 123 +++++++++++++++++++++++++----- src/nanoarrow/ipc/files_test.cc | 3 +- src/nanoarrow/ipc/writer.c | 4 +- src/nanoarrow/ipc/writer_test.cc | 9 ++- src/nanoarrow/nanoarrow_ipc.h | 39 +++++++--- 8 files changed, 222 insertions(+), 60 deletions(-) diff --git a/src/nanoarrow/ipc/codecs.c b/src/nanoarrow/ipc/codecs.c index 0f2aa36cf..5f3a03144 100644 --- a/src/nanoarrow/ipc/codecs.c +++ b/src/nanoarrow/ipc/codecs.c @@ -46,7 +46,7 @@ static ArrowErrorCode ArrowIpcDecompressZstd(struct ArrowBufferView src, uint8_t } static ArrowErrorCode ArrowIpcCompressZstd(struct ArrowBufferView src, - struct ArrowBuffer* dst, + int compression_level, struct ArrowBuffer* dst, struct ArrowError* error) { size_t dst_capacity = ZSTD_compressBound((size_t)src.size_bytes); if (ZSTD_isError(dst_capacity)) { @@ -59,7 +59,7 @@ static ArrowErrorCode ArrowIpcCompressZstd(struct ArrowBufferView src, error); size_t code = ZSTD_compress((void*)(dst->data + dst->size_bytes), dst_capacity, - src.data.data, (size_t)src.size_bytes, ZSTD_CLEVEL_DEFAULT); + src.data.data, (size_t)src.size_bytes, compression_level); if (ZSTD_isError(code)) { ArrowErrorSet(error, "ZSTD_compress([buffer with %" PRId64 " bytes]) failed with error '%s'", @@ -139,13 +139,15 @@ static ArrowErrorCode ArrowIpcDecompressLZ4(struct ArrowBufferView src, uint8_t* } static ArrowErrorCode ArrowIpcCompressLZ4(struct ArrowBufferView src, - struct ArrowBuffer* dst, + int compression_level, struct ArrowBuffer* dst, struct ArrowError* error) { - // Default preferences (default compression level, no content size, no checksums). + // Default preferences except for the compression level (no content size, no + // checksums). // This produces a single complete frame, which is what ArrowIpcDecompressLZ4() // and Arrow C++ expect. LZ4F_preferences_t prefs; memset(&prefs, 0, sizeof(prefs)); + prefs.compressionLevel = compression_level; size_t dst_capacity = LZ4F_compressFrameBound((size_t)src.size_bytes, &prefs); NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowBufferReserve(dst, (int64_t)dst_capacity), @@ -277,7 +279,8 @@ struct ArrowIpcSerialCompressorPrivate { static ArrowErrorCode ArrowIpcSerialCompressorCompress( struct ArrowIpcCompressor* compressor, enum ArrowIpcCompressionType compression_type, - struct ArrowBufferView src, struct ArrowBuffer* dst, struct ArrowError* error) { + int compression_level, struct ArrowBufferView src, struct ArrowBuffer* dst, + struct ArrowError* error) { struct ArrowIpcSerialCompressorPrivate* private_data = (struct ArrowIpcSerialCompressorPrivate*)compressor->private_data; @@ -294,7 +297,7 @@ static ArrowErrorCode ArrowIpcSerialCompressorCompress( return ENOTSUP; } - NANOARROW_RETURN_NOT_OK(fn(src, dst, error)); + NANOARROW_RETURN_NOT_OK(fn(src, compression_level, dst, error)); return NANOARROW_OK; } diff --git a/src/nanoarrow/ipc/codecs_test.cc b/src/nanoarrow/ipc/codecs_test.cc index f298e4149..c58e4aa16 100644 --- a/src/nanoarrow/ipc/codecs_test.cc +++ b/src/nanoarrow/ipc/codecs_test.cc @@ -16,6 +16,7 @@ // under the License. #include +#include #include #include @@ -187,11 +188,12 @@ TEST(NanoarrowIpcTest, SerialDecompressor) { "Compression type with value 2 not supported by this build of nanoarrow"); } -// Compress input (appending to a buffer that already has content), decompress the -// appended bytes, and check that the result matches the input. Returns the number of -// compressed bytes that were appended. +// Compress input at compression_level (appending to a buffer that already has content), +// decompress the appended bytes, and check that the result matches the input. Returns +// the number of compressed bytes that were appended. static int64_t TestCompressRoundtrip(ArrowIpcCompressFunction compress, ArrowIpcDecompressFunction decompress, + int compression_level, const std::vector& input) { struct ArrowError error {}; nanoarrow::UniqueBuffer compressed; @@ -202,7 +204,7 @@ static int64_t TestCompressRoundtrip(ArrowIpcCompressFunction compress, EXPECT_EQ(ArrowBufferAppend(compressed.get(), existing, existing_size), NANOARROW_OK); EXPECT_EQ(compress({{input.data()}, static_cast(input.size())}, - compressed.get(), &error), + compression_level, compressed.get(), &error), NANOARROW_OK) << error.message; EXPECT_GT(compressed->size_bytes, existing_size); @@ -228,22 +230,26 @@ static std::vector CompressibleInput(size_t n) { return out; } -// Check compress/decompress on empty, small, and multi-block inputs +// Check compress/decompress on empty, small, and multi-block inputs at each level static void TestCompressionFunctions(ArrowIpcCompressFunction compress, - ArrowIpcDecompressFunction decompress) { + ArrowIpcDecompressFunction decompress, + const std::vector& compression_levels) { ASSERT_NE(compress, nullptr); ASSERT_NE(decompress, nullptr); - TestCompressRoundtrip(compress, decompress, {}); - TestCompressRoundtrip( - compress, decompress, - std::vector(kUncompressed012, - kUncompressed012 + sizeof(kUncompressed012))); - - // Large enough to span several blocks; a repetitive input must actually shrink auto input = CompressibleInput(1 << 20); - int64_t compressed_size = TestCompressRoundtrip(compress, decompress, input); - EXPECT_LT(compressed_size, static_cast(input.size() / 10)); + for (int level : compression_levels) { + SCOPED_TRACE("compression level " + std::to_string(level)); + TestCompressRoundtrip(compress, decompress, level, {}); + TestCompressRoundtrip( + compress, decompress, level, + std::vector(kUncompressed012, + kUncompressed012 + sizeof(kUncompressed012))); + + // Large enough to span several blocks; a repetitive input must actually shrink + int64_t compressed_size = TestCompressRoundtrip(compress, decompress, level, input); + EXPECT_LT(compressed_size, static_cast(input.size() / 10)); + } } TEST(NanoarrowIpcTest, NanoarrowIpcZstdCompressBuildMatchesRuntime) { @@ -258,8 +264,10 @@ TEST(NanoarrowIpcTest, ZstdCompressRoundtrip) { if (ArrowIpcGetZstdCompressionFunction() == nullptr) { GTEST_SKIP() << "nanoarrow_ipc not built with NANOARROW_IPC_WITH_ZSTD"; } + // Default, a negative (fast) level, the lowest regular level, and a high level TestCompressionFunctions(ArrowIpcGetZstdCompressionFunction(), - ArrowIpcGetZstdDecompressionFunction()); + ArrowIpcGetZstdDecompressionFunction(), + {NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, -5, 1, 19}); } TEST(NanoarrowIpcTest, NanoarrowIpcLZ4CompressBuildMatchesRuntime) { @@ -274,8 +282,22 @@ TEST(NanoarrowIpcTest, LZ4CompressRoundtrip) { if (ArrowIpcGetLZ4CompressionFunction() == nullptr) { GTEST_SKIP() << "nanoarrow_ipc not built with NANOARROW_IPC_WITH_LZ4"; } + // Default (fast), acceleration, the last fast level, and LZ4HC levels TestCompressionFunctions(ArrowIpcGetLZ4CompressionFunction(), - ArrowIpcGetLZ4DecompressionFunction()); + ArrowIpcGetLZ4DecompressionFunction(), + {NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, -1, 2, 9, 12}); +} + +// A stand-in compression function that records the level it was called with and +// "compresses" by copying +static int last_compression_level = 0; + +static ArrowErrorCode RecordLevelAndCopy(struct ArrowBufferView src, + int compression_level, struct ArrowBuffer* dst, + struct ArrowError* error) { + NANOARROW_UNUSED(error); + last_compression_level = compression_level; + return ArrowBufferAppend(dst, src.data.data, src.size_bytes); } TEST(NanoarrowIpcTest, SerialCompressor) { @@ -292,13 +314,15 @@ TEST(NanoarrowIpcTest, SerialCompressor) { // NONE is not a codec that can be used to compress nanoarrow::UniqueBuffer dst; EXPECT_EQ(compressor->compress(compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_NONE, - {{nullptr}, 0}, dst.get(), &error), + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, {{nullptr}, 0}, + dst.get(), &error), EINVAL); EXPECT_STREQ(error.message, "Unknown compression type with value 0"); // Check a compress for a supported codec if we have one (or for an error if we don't) if (ArrowIpcGetZstdCompressionFunction() != nullptr) { ASSERT_EQ(compressor->compress(compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, {{kUncompressed012}, sizeof(kUncompressed012)}, dst.get(), &error), NANOARROW_OK) @@ -314,6 +338,7 @@ TEST(NanoarrowIpcTest, SerialCompressor) { EXPECT_TRUE(std::memcmp(out, kUncompressed012, sizeof(kUncompressed012)) == 0); } else { EXPECT_EQ(compressor->compress(compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, {{nullptr}, 0}, dst.get(), &error), ENOTSUP); EXPECT_STREQ( @@ -326,8 +351,25 @@ TEST(NanoarrowIpcTest, SerialCompressor) { compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, nullptr), NANOARROW_OK); EXPECT_EQ(compressor->compress(compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, - {{nullptr}, 0}, dst.get(), &error), + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, {{nullptr}, 0}, + dst.get(), &error), ENOTSUP); EXPECT_STREQ(error.message, "Compression type with value 2 not supported by this build of nanoarrow"); + + // The compression level is passed through to the function for the codec + ASSERT_EQ(ArrowIpcSerialCompressorSetFunction(compressor.get(), + NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, + &RecordLevelAndCopy), + NANOARROW_OK); + dst->size_bytes = 0; + last_compression_level = 0; + ASSERT_EQ(compressor->compress( + compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, 7, + {{kUncompressed012}, sizeof(kUncompressed012)}, dst.get(), &error), + NANOARROW_OK) + << error.message; + EXPECT_EQ(last_compression_level, 7); + ASSERT_EQ(dst->size_bytes, static_cast(sizeof(kUncompressed012))); + EXPECT_EQ(std::memcmp(dst->data, kUncompressed012, sizeof(kUncompressed012)), 0); } diff --git a/src/nanoarrow/ipc/encoder.c b/src/nanoarrow/ipc/encoder.c index 454141e0c..ee36a8127 100644 --- a/src/nanoarrow/ipc/encoder.c +++ b/src/nanoarrow/ipc/encoder.c @@ -55,6 +55,8 @@ struct ArrowIpcEncoderPrivate { struct ArrowIpcCompressor compressor; // Whether compressor was provided by ArrowIpcEncoderSetCompressor() int custom_compressor; + // Compression level passed to the compressor when codec != NONE + int compression_level; }; ArrowErrorCode ArrowIpcEncoderInit(struct ArrowIpcEncoder* encoder) { @@ -78,6 +80,7 @@ ArrowErrorCode ArrowIpcEncoderInit(struct ArrowIpcEncoder* encoder) { private->codec = NANOARROW_IPC_COMPRESSION_TYPE_NONE; private->compressor.release = NULL; private->custom_compressor = 0; + private->compression_level = NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT; return NANOARROW_OK; } @@ -150,7 +153,7 @@ ArrowErrorCode ArrowIpcEncoderSetCompressor(struct ArrowIpcEncoder* encoder, ArrowErrorCode ArrowIpcEncoderSetCompression( struct ArrowIpcEncoder* encoder, enum ArrowIpcCompressionType compression_type, - struct ArrowError* error) { + int compression_level, struct ArrowError* error) { NANOARROW_DCHECK(encoder != NULL && encoder->private_data != NULL); struct ArrowIpcEncoderPrivate* private = (struct ArrowIpcEncoderPrivate*)encoder->private_data; @@ -190,6 +193,7 @@ ArrowErrorCode ArrowIpcEncoderSetCompression( } private->codec = compression_type; + private->compression_level = compression_level; return NANOARROW_OK; } @@ -730,7 +734,8 @@ static ArrowErrorCode ArrowIpcEncoderAppendCompressedBuffer( int64_t payload_offset = prefix_offset + (int64_t)sizeof(int64_t); NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowBufferAppendInt64(body_buffer, 0), error); NANOARROW_RETURN_NOT_OK(private->compressor.compress( - &private->compressor, private->codec, buffer_view, body_buffer, error)); + &private->compressor, private->codec, private->compression_level, buffer_view, + body_buffer, error)); int64_t prefix = buffer_view.size_bytes; if (body_buffer->size_bytes - payload_offset >= buffer_view.size_bytes) { diff --git a/src/nanoarrow/ipc/encoder_test.cc b/src/nanoarrow/ipc/encoder_test.cc index 4c640a7a8..2d5843fae 100644 --- a/src/nanoarrow/ipc/encoder_test.cc +++ b/src/nanoarrow/ipc/encoder_test.cc @@ -529,7 +529,9 @@ static void TestCompressedRecordBatchRoundtrip(enum ArrowIpcCompressionType code uncompressed_message.get()), NANOARROW_OK); - ASSERT_EQ(ArrowIpcEncoderSetCompression(encoder.get(), codec, &error), NANOARROW_OK) + ASSERT_EQ(ArrowIpcEncoderSetCompression( + encoder.get(), codec, NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), + NANOARROW_OK) << error.message; nanoarrow::UniqueBuffer message, body; @@ -598,9 +600,10 @@ static void TestCompressedRecordBatchRoundtrip(enum ArrowIpcCompressionType code EXPECT_EQ(is_equal, 1) << error.message; // Compression can be turned off again - ASSERT_EQ(ArrowIpcEncoderSetCompression(encoder.get(), - NANOARROW_IPC_COMPRESSION_TYPE_NONE, &error), - NANOARROW_OK) + ASSERT_EQ( + ArrowIpcEncoderSetCompression(encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_NONE, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), + NANOARROW_OK) << error.message; message->size_bytes = 0; body->size_bytes = 0; @@ -689,40 +692,46 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSetCompressionErrors) { struct ArrowError error; EXPECT_EQ(ArrowIpcEncoderSetCompression( - encoder.get(), static_cast(99), &error), + encoder.get(), static_cast(99), + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), EINVAL); EXPECT_STREQ(error.message, "Unknown compression type with value 99"); // NONE is always supported - EXPECT_EQ(ArrowIpcEncoderSetCompression(encoder.get(), - NANOARROW_IPC_COMPRESSION_TYPE_NONE, &error), - NANOARROW_OK) + EXPECT_EQ( + ArrowIpcEncoderSetCompression(encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_NONE, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), + NANOARROW_OK) << error.message; // Codecs that were not built in are rejected when they are set rather than when // the first batch is encoded #if defined(NANOARROW_IPC_WITH_LZ4) EXPECT_EQ(ArrowIpcEncoderSetCompression( - encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, &error), + encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), NANOARROW_OK) << error.message; #else EXPECT_EQ(ArrowIpcEncoderSetCompression( - encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, &error), + encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), ENOTSUP); EXPECT_STREQ(error.message, "Compression type with value 1 not supported by this build of nanoarrow"); #endif #if defined(NANOARROW_IPC_WITH_ZSTD) - EXPECT_EQ(ArrowIpcEncoderSetCompression(encoder.get(), - NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, &error), - NANOARROW_OK) + EXPECT_EQ( + ArrowIpcEncoderSetCompression(encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), + NANOARROW_OK) << error.message; #else - EXPECT_EQ(ArrowIpcEncoderSetCompression(encoder.get(), - NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, &error), - ENOTSUP); + EXPECT_EQ( + ArrowIpcEncoderSetCompression(encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), + ENOTSUP); EXPECT_STREQ(error.message, "Compression type with value 2 not supported by this build of nanoarrow"); #endif @@ -746,7 +755,8 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSetCompressor) { // With a custom compressor, support is not checked until a batch is encoded ASSERT_EQ(ArrowIpcEncoderSetCompression( - encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, &error), + encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), NANOARROW_OK) << error.message; @@ -758,3 +768,82 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSetCompressor) { EXPECT_STREQ(error.message, "Compression type with value 1 not supported by this build of nanoarrow"); } + +// A stand-in compression function that records the level it was called with and +// "compresses" by copying (so that every buffer takes the uncompressed fallback path) +static int last_compression_level = 0; + +static ArrowErrorCode RecordLevelAndCopy(struct ArrowBufferView src, + int compression_level, struct ArrowBuffer* dst, + struct ArrowError* error) { + NANOARROW_UNUSED(error); + last_compression_level = compression_level; + return ArrowBufferAppend(dst, src.data.data, src.size_bytes); +} + +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderCompressionLevel) { + nanoarrow::ipc::UniqueEncoder encoder; + ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); + nanoarrow::ipc::UniqueDecoder decoder; + ASSERT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK); + struct ArrowError error; + + nanoarrow::ipc::UniqueCompressor compressor; + ASSERT_EQ(ArrowIpcSerialCompressor(compressor.get()), NANOARROW_OK); + ASSERT_EQ( + ArrowIpcSerialCompressorSetFunction( + compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, &RecordLevelAndCopy), + NANOARROW_OK); + ASSERT_EQ(ArrowIpcEncoderSetCompressor(encoder.get(), compressor.get()), NANOARROW_OK); + ASSERT_EQ(ArrowIpcEncoderSetCompression( + encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, 11, &error), + NANOARROW_OK) + << error.message; + + CompressibleRecordBatch batch; + ASSERT_EQ(ArrowIpcDecoderSetSchema(decoder.get(), batch.schema(), &error), NANOARROW_OK) + << error.message; + + // The level reaches the codec function + last_compression_level = 0; + nanoarrow::UniqueBuffer message, body; + ASSERT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), batch.array_view(), + body.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ( + ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/true, message.get()), + NANOARROW_OK); + EXPECT_EQ(last_compression_level, 11); + + // Copying never shrinks a buffer, so every buffer took the uncompressed (-1) path; + // the message still declares the codec and must decode + struct ArrowBufferView message_view = {{message->data}, message->size_bytes}; + ASSERT_EQ(ArrowIpcDecoderVerifyHeader(decoder.get(), message_view, &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcDecoderDecodeHeader(decoder.get(), message_view, &error), + NANOARROW_OK) + << error.message; + EXPECT_EQ(decoder->codec, NANOARROW_IPC_COMPRESSION_TYPE_ZSTD); + + nanoarrow::UniqueArray decoded; + struct ArrowBufferView body_view = {{body->data}, body->size_bytes}; + ASSERT_EQ(ArrowIpcDecoderDecodeArray(decoder.get(), body_view, -1, decoded.get(), + NANOARROW_VALIDATION_LEVEL_FULL, &error), + NANOARROW_OK) + << error.message; + + nanoarrow::UniqueArrayView decoded_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(decoded_view.get(), batch.schema(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayViewSetArray(decoded_view.get(), decoded.get(), &error), + NANOARROW_OK) + << error.message; + int is_equal = 0; + ASSERT_EQ(ArrowArrayViewCompare(decoded_view.get(), batch.array_view(), + NANOARROW_COMPARE_IDENTICAL, &is_equal, &error), + NANOARROW_OK); + EXPECT_EQ(is_equal, 1) << error.message; +} diff --git a/src/nanoarrow/ipc/files_test.cc b/src/nanoarrow/ipc/files_test.cc index f5f77ddca..73fc9724f 100644 --- a/src/nanoarrow/ipc/files_test.cc +++ b/src/nanoarrow/ipc/files_test.cc @@ -223,7 +223,8 @@ class TestFile { nanoarrow::ipc::UniqueWriter writer; NANOARROW_RETURN_NOT_OK(ArrowIpcWriterInit(writer.get(), output_stream.get())); - NANOARROW_RETURN_NOT_OK(ArrowIpcWriterSetCompression(writer.get(), codec, error)); + NANOARROW_RETURN_NOT_OK(ArrowIpcWriterSetCompression( + writer.get(), codec, NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, error)); nanoarrow::UniqueArrayView array_view; NANOARROW_RETURN_NOT_OK( diff --git a/src/nanoarrow/ipc/writer.c b/src/nanoarrow/ipc/writer.c index 2aab91d37..28672f242 100644 --- a/src/nanoarrow/ipc/writer.c +++ b/src/nanoarrow/ipc/writer.c @@ -228,11 +228,13 @@ void ArrowIpcWriterReset(struct ArrowIpcWriter* writer) { ArrowErrorCode ArrowIpcWriterSetCompression(struct ArrowIpcWriter* writer, enum ArrowIpcCompressionType compression_type, + int compression_level, struct ArrowError* error) { NANOARROW_DCHECK(writer != NULL && writer->private_data != NULL); struct ArrowIpcWriterPrivate* private = (struct ArrowIpcWriterPrivate*)writer->private_data; - return ArrowIpcEncoderSetCompression(&private->encoder, compression_type, error); + return ArrowIpcEncoderSetCompression(&private->encoder, compression_type, + compression_level, error); } static struct ArrowBufferView ArrowBufferToBufferView(const struct ArrowBuffer* buffer) { diff --git a/src/nanoarrow/ipc/writer_test.cc b/src/nanoarrow/ipc/writer_test.cc index 5a74b3ce9..8ae5b7918 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -243,7 +243,9 @@ static void WriteCompressibleBatch(enum ArrowIpcCompressionType codec, bool as_f ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream.get(), output), NANOARROW_OK); nanoarrow::ipc::UniqueWriter writer; ASSERT_EQ(ArrowIpcWriterInit(writer.get(), stream.get()), NANOARROW_OK); - ASSERT_EQ(ArrowIpcWriterSetCompression(writer.get(), codec, &error), NANOARROW_OK) + ASSERT_EQ(ArrowIpcWriterSetCompression(writer.get(), codec, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), + NANOARROW_OK) << error.message; if (as_file) { @@ -349,8 +351,9 @@ TEST(NanoarrowIpcWriter, SetCompressionErrors) { ASSERT_EQ(ArrowIpcWriterInit(writer.get(), stream.get()), NANOARROW_OK); struct ArrowError error; - EXPECT_EQ(ArrowIpcWriterSetCompression( - writer.get(), static_cast(99), &error), + EXPECT_EQ(ArrowIpcWriterSetCompression(writer.get(), + static_cast(99), + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), EINVAL); EXPECT_STREQ(error.message, "Unknown compression type with value 99"); } diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index 25d26ca5e..df427963f 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -426,6 +426,9 @@ ArrowIpcSerialDecompressorSetFunction(struct ArrowIpcDecompressor* decompressor, enum ArrowIpcCompressionType compression_type, ArrowIpcDecompressFunction decompress_function); +/// \brief Compression level that selects the codec's default level +#define NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT 0 + /// \brief A user-extensible compressor /// /// The ArrowIpcCompressor is the underlying object that enables buffer compression @@ -434,13 +437,14 @@ ArrowIpcSerialDecompressorSetFunction(struct ArrowIpcDecompressor* decompressor, struct ArrowIpcCompressor { /// \brief Compress a buffer /// - /// Compresses src using compression_type and appends the compressed bytes to dst. - /// Any content already in dst must be preserved (i.e., implementations may only - /// append to dst). + /// Compresses src using compression_type at compression_level and appends the + /// compressed bytes to dst. Any content already in dst must be preserved (i.e., + /// implementations may only append to dst). See ArrowIpcCompressFunction for the + /// interpretation of compression_level. ArrowErrorCode (*compress)(struct ArrowIpcCompressor* compressor, enum ArrowIpcCompressionType compression_type, - struct ArrowBufferView src, struct ArrowBuffer* dst, - struct ArrowError* error); + int compression_level, struct ArrowBufferView src, + struct ArrowBuffer* dst, struct ArrowError* error); /// \brief Release the compressor and any resources it may be holding /// @@ -455,11 +459,19 @@ struct ArrowIpcCompressor { /// \brief A self-contained compression function /// -/// Compresses src and appends the compressed bytes to dst. Because the compressed -/// size is not known in advance, implementations are responsible for reserving -/// sufficient space in dst (e.g., using the compression library's bound function) -/// and must only append to dst. +/// Compresses src at compression_level and appends the compressed bytes to dst. Because +/// the compressed size is not known in advance, implementations are responsible for +/// reserving sufficient space in dst (e.g., using the compression library's bound +/// function) and must only append to dst. +/// +/// The interpretation of compression_level is codec-specific. +/// NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT selects the codec's default level; other +/// values are passed to the underlying library as-is (for ZSTD, ZSTD_minCLevel() to +/// ZSTD_maxCLevel() where negative levels favour speed; for LZ4, up to +/// LZ4F_compressionLevel_max() where levels >= 3 use LZ4HC and negative levels select +/// acceleration). Both libraries clamp out-of-range levels. typedef ArrowErrorCode (*ArrowIpcCompressFunction)(struct ArrowBufferView src, + int compression_level, struct ArrowBuffer* dst, struct ArrowError* error); @@ -947,6 +959,10 @@ ArrowIpcEncoderSetMessageMetadata(struct ArrowIpcEncoder* encoder, /// prefix of -1, and empty buffers are written as-is. The setting persists until it is /// changed and does not affect Schema messages. /// +/// compression_level is passed to the compressor unchanged; see +/// ArrowIpcCompressFunction for its interpretation. Use +/// NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT for the codec's default level. +/// /// Returns EINVAL for an unknown compression type and ENOTSUP if the compression type /// is not supported by this build of nanoarrow (i.e., nanoarrow was not built with /// NANOARROW_IPC_WITH_LZ4 or NANOARROW_IPC_WITH_ZSTD). If a custom compressor was @@ -954,7 +970,7 @@ ArrowIpcEncoderSetMessageMetadata(struct ArrowIpcEncoder* encoder, /// is encoded. NANOARROW_DLL ArrowErrorCode ArrowIpcEncoderSetCompression( struct ArrowIpcEncoder* encoder, enum ArrowIpcCompressionType compression_type, - struct ArrowError* error); + int compression_level, struct ArrowError* error); /// \brief Set the compressor implementation used by this encoder /// @@ -1050,12 +1066,13 @@ NANOARROW_DLL void ArrowIpcWriterReset(struct ArrowIpcWriter* writer); /// /// See ArrowIpcEncoderSetCompression(). Compression applies to record batches written /// after this call (in both stream and file mode) and may be changed between batches. +/// Use NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT for the codec's default level. /// /// Returns EINVAL for an unknown compression type and ENOTSUP if the compression type /// is not supported by this build of nanoarrow. NANOARROW_DLL ArrowErrorCode ArrowIpcWriterSetCompression( struct ArrowIpcWriter* writer, enum ArrowIpcCompressionType compression_type, - struct ArrowError* error); + int compression_level, struct ArrowError* error); /// \brief Write a schema to the output byte stream /// From 43d39bb6a6b8ec57537742aa9d3947cba3cecd06 Mon Sep 17 00:00:00 2001 From: Pedro Holanda Date: Tue, 8 Sep 2026 15:32:39 +0200 Subject: [PATCH 06/13] fix(ipc): clamp extreme LZ4 levels and verify compression level behavior --- src/nanoarrow/ipc/codecs.c | 7 +++++++ src/nanoarrow/ipc/codecs_test.cc | 31 +++++++++++++++++++++++++++++++ src/nanoarrow/ipc/writer_test.cc | 24 ++++++++++++++++-------- src/nanoarrow/nanoarrow_ipc.h | 4 ++-- 4 files changed, 56 insertions(+), 10 deletions(-) diff --git a/src/nanoarrow/ipc/codecs.c b/src/nanoarrow/ipc/codecs.c index 5f3a03144..e0155823d 100644 --- a/src/nanoarrow/ipc/codecs.c +++ b/src/nanoarrow/ipc/codecs.c @@ -16,6 +16,7 @@ // under the License. #include +#include #include "nanoarrow/nanoarrow_ipc.h" @@ -141,6 +142,12 @@ static ArrowErrorCode ArrowIpcDecompressLZ4(struct ArrowBufferView src, uint8_t* static ArrowErrorCode ArrowIpcCompressLZ4(struct ArrowBufferView src, int compression_level, struct ArrowBuffer* dst, struct ArrowError* error) { + // LZ4 computes acceleration as -level + 1 before clamping it. Keep that + // calculation representable even for the most negative int values. + if (compression_level < 1 - INT_MAX) { + compression_level = 1 - INT_MAX; + } + // Default preferences except for the compression level (no content size, no // checksums). // This produces a single complete frame, which is what ArrowIpcDecompressLZ4() diff --git a/src/nanoarrow/ipc/codecs_test.cc b/src/nanoarrow/ipc/codecs_test.cc index c58e4aa16..4ed8d6264 100644 --- a/src/nanoarrow/ipc/codecs_test.cc +++ b/src/nanoarrow/ipc/codecs_test.cc @@ -16,6 +16,7 @@ // under the License. #include +#include #include #include @@ -238,6 +239,7 @@ static void TestCompressionFunctions(ArrowIpcCompressFunction compress, ASSERT_NE(decompress, nullptr); auto input = CompressibleInput(1 << 20); + int64_t default_size = 0; for (int level : compression_levels) { SCOPED_TRACE("compression level " + std::to_string(level)); TestCompressRoundtrip(compress, decompress, level, {}); @@ -249,7 +251,16 @@ static void TestCompressionFunctions(ArrowIpcCompressFunction compress, // Large enough to span several blocks; a repetitive input must actually shrink int64_t compressed_size = TestCompressRoundtrip(compress, decompress, level, input); EXPECT_LT(compressed_size, static_cast(input.size() / 10)); + if (level == NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT) { + default_size = compressed_size; + } } + + // High acceleration must trade compression ratio for speed on this input. Merely + // roundtripping at several levels would also pass if the level were ignored. + ASSERT_GT(default_size, 0); + int64_t accelerated_size = TestCompressRoundtrip(compress, decompress, -65536, input); + EXPECT_GT(accelerated_size, default_size * 2); } TEST(NanoarrowIpcTest, NanoarrowIpcZstdCompressBuildMatchesRuntime) { @@ -288,6 +299,26 @@ TEST(NanoarrowIpcTest, LZ4CompressRoundtrip) { {NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, -1, 2, 9, 12}); } +TEST(NanoarrowIpcTest, LZ4CompressMinimumLevels) { + auto compress = ArrowIpcGetLZ4CompressionFunction(); + if (compress == nullptr) { + GTEST_SKIP() << "nanoarrow_ipc not built with NANOARROW_IPC_WITH_LZ4"; + } + auto decompress = ArrowIpcGetLZ4DecompressionFunction(); + ASSERT_NE(decompress, nullptr); + + auto input = CompressibleInput(1 << 20); + int64_t accelerated_size = TestCompressRoundtrip(compress, decompress, -65536, input); + for (int level : {std::numeric_limits::min(), std::numeric_limits::min() + 1, + std::numeric_limits::min() + 2}) { + SCOPED_TRACE("compression level " + std::to_string(level)); + // The most negative levels must saturate at maximum acceleration rather than + // overflow and fall back to the default compression level. + EXPECT_EQ(TestCompressRoundtrip(compress, decompress, level, input), + accelerated_size); + } +} + // A stand-in compression function that records the level it was called with and // "compresses" by copying static int last_compression_level = 0; diff --git a/src/nanoarrow/ipc/writer_test.cc b/src/nanoarrow/ipc/writer_test.cc index 8ae5b7918..3ab02e676 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -224,8 +224,10 @@ static void InitCompressibleBatch(struct ArrowSchema* schema, struct ArrowArray* ASSERT_EQ(ArrowArrayFinishBuildingDefault(array, nullptr), NANOARROW_OK); } -// Write schema + batch + EOS (optionally as an IPC file) to output using codec -static void WriteCompressibleBatch(enum ArrowIpcCompressionType codec, bool as_file, +// Write schema + batch + EOS (optionally as an IPC file) using codec and +// compression_level +static void WriteCompressibleBatch(enum ArrowIpcCompressionType codec, + int compression_level, bool as_file, struct ArrowBuffer* output) { struct ArrowError error; @@ -243,8 +245,7 @@ static void WriteCompressibleBatch(enum ArrowIpcCompressionType codec, bool as_f ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream.get(), output), NANOARROW_OK); nanoarrow::ipc::UniqueWriter writer; ASSERT_EQ(ArrowIpcWriterInit(writer.get(), stream.get()), NANOARROW_OK); - ASSERT_EQ(ArrowIpcWriterSetCompression(writer.get(), codec, - NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), + ASSERT_EQ(ArrowIpcWriterSetCompression(writer.get(), codec, compression_level, &error), NANOARROW_OK) << error.message; @@ -317,15 +318,22 @@ static void TestCompressedWriting(enum ArrowIpcCompressionType codec) { for (bool as_file : {false, true}) { SCOPED_TRACE(as_file ? "file" : "stream"); - nanoarrow::UniqueBuffer uncompressed, compressed; - ASSERT_NO_FATAL_FAILURE(WriteCompressibleBatch(NANOARROW_IPC_COMPRESSION_TYPE_NONE, - as_file, uncompressed.get())); - ASSERT_NO_FATAL_FAILURE(WriteCompressibleBatch(codec, as_file, compressed.get())); + nanoarrow::UniqueBuffer uncompressed, compressed, accelerated; + ASSERT_NO_FATAL_FAILURE(WriteCompressibleBatch( + NANOARROW_IPC_COMPRESSION_TYPE_NONE, NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, + as_file, uncompressed.get())); + ASSERT_NO_FATAL_FAILURE(WriteCompressibleBatch( + codec, NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, as_file, compressed.get())); + ASSERT_NO_FATAL_FAILURE( + WriteCompressibleBatch(codec, -65536, as_file, accelerated.get())); EXPECT_LT(compressed->size_bytes, uncompressed->size_bytes); + // A nondefault level must reach the codec through the writer and encoder. + EXPECT_GT(accelerated->size_bytes, compressed->size_bytes); // The stream portion of a file follows the padded magic int64_t offset = as_file ? sizeof(NANOARROW_IPC_FILE_PADDED_MAGIC) : 0; ASSERT_NO_FATAL_FAILURE(CheckCompressibleBatch(compressed.get(), offset)); + ASSERT_NO_FATAL_FAILURE(CheckCompressibleBatch(accelerated.get(), offset)); } } diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index df427963f..9f2bab675 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -466,10 +466,10 @@ struct ArrowIpcCompressor { /// /// The interpretation of compression_level is codec-specific. /// NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT selects the codec's default level; other -/// values are passed to the underlying library as-is (for ZSTD, ZSTD_minCLevel() to +/// values follow the underlying library's conventions (for ZSTD, ZSTD_minCLevel() to /// ZSTD_maxCLevel() where negative levels favour speed; for LZ4, up to /// LZ4F_compressionLevel_max() where levels >= 3 use LZ4HC and negative levels select -/// acceleration). Both libraries clamp out-of-range levels. +/// acceleration). The built-in implementations clamp out-of-range levels. typedef ArrowErrorCode (*ArrowIpcCompressFunction)(struct ArrowBufferView src, int compression_level, struct ArrowBuffer* dst, From d3e42ee55c3eb1e86823ba7548b019f65c52289e Mon Sep 17 00:00:00 2001 From: Pedro Holanda Date: Tue, 8 Sep 2026 16:04:55 +0200 Subject: [PATCH 07/13] chore(ipc): suppress clang-tidy EnumCastOutOfRange in compression error tests --- src/nanoarrow/ipc/encoder_test.cc | 11 +++++++---- src/nanoarrow/ipc/writer_test.cc | 6 ++++-- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/src/nanoarrow/ipc/encoder_test.cc b/src/nanoarrow/ipc/encoder_test.cc index 2d5843fae..c8338882b 100644 --- a/src/nanoarrow/ipc/encoder_test.cc +++ b/src/nanoarrow/ipc/encoder_test.cc @@ -691,10 +691,13 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSetCompressionErrors) { ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); struct ArrowError error; - EXPECT_EQ(ArrowIpcEncoderSetCompression( - encoder.get(), static_cast(99), - NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), - EINVAL); + // 99 is not an enumerator; it exercises the EINVAL path + // NOLINTNEXTLINE(clang-analyzer-optin.core.EnumCastOutOfRange) + auto unknown_type = static_cast(99); + EXPECT_EQ( + ArrowIpcEncoderSetCompression(encoder.get(), unknown_type, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), + EINVAL); EXPECT_STREQ(error.message, "Unknown compression type with value 99"); // NONE is always supported diff --git a/src/nanoarrow/ipc/writer_test.cc b/src/nanoarrow/ipc/writer_test.cc index 3ab02e676..d753dea83 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -359,8 +359,10 @@ TEST(NanoarrowIpcWriter, SetCompressionErrors) { ASSERT_EQ(ArrowIpcWriterInit(writer.get(), stream.get()), NANOARROW_OK); struct ArrowError error; - EXPECT_EQ(ArrowIpcWriterSetCompression(writer.get(), - static_cast(99), + // 99 is not an enumerator; it exercises the EINVAL path + // NOLINTNEXTLINE(clang-analyzer-optin.core.EnumCastOutOfRange) + auto unknown_type = static_cast(99); + EXPECT_EQ(ArrowIpcWriterSetCompression(writer.get(), unknown_type, NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), EINVAL); EXPECT_STREQ(error.message, "Unknown compression type with value 99"); From 20c3d0ba3ad87b6b6cd41e241e6309301b98c2b9 Mon Sep 17 00:00:00 2001 From: Pedro Holanda Date: Tue, 8 Sep 2026 16:24:17 +0200 Subject: [PATCH 08/13] feat(ipc): add ArrowIpcCompressionTypeToString() and ArrowIpcCompressionTypeFromString() --- src/nanoarrow/ipc/codecs.c | 37 ++++++++++++++++++++++++++++++++ src/nanoarrow/ipc/codecs_test.cc | 36 +++++++++++++++++++++++++++++++ src/nanoarrow/nanoarrow_ipc.h | 18 ++++++++++++++++ 3 files changed, 91 insertions(+) diff --git a/src/nanoarrow/ipc/codecs.c b/src/nanoarrow/ipc/codecs.c index e0155823d..6357ae542 100644 --- a/src/nanoarrow/ipc/codecs.c +++ b/src/nanoarrow/ipc/codecs.c @@ -203,6 +203,43 @@ static int ArrowIpcCompressionTypeIsCodec(enum ArrowIpcCompressionType compressi } } +const char* ArrowIpcCompressionTypeToString( + enum ArrowIpcCompressionType compression_type) { + switch (compression_type) { + case NANOARROW_IPC_COMPRESSION_TYPE_NONE: + return "none"; + case NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME: + return "lz4"; + case NANOARROW_IPC_COMPRESSION_TYPE_ZSTD: + return "zstd"; + default: + return NULL; + } +} + +ArrowErrorCode ArrowIpcCompressionTypeFromString( + const char* name, enum ArrowIpcCompressionType* compression_type_out, + struct ArrowError* error) { + NANOARROW_DCHECK(compression_type_out != NULL); + static const enum ArrowIpcCompressionType types[] = { + NANOARROW_IPC_COMPRESSION_TYPE_NONE, NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, + NANOARROW_IPC_COMPRESSION_TYPE_ZSTD}; + + if (name != NULL) { + for (size_t i = 0; i < sizeof(types) / sizeof(types[0]); i++) { + if (strcmp(name, ArrowIpcCompressionTypeToString(types[i])) == 0) { + *compression_type_out = types[i]; + return NANOARROW_OK; + } + } + } + + ArrowErrorSet(error, + "Unknown compression type name '%s' (expected 'none', 'lz4', or 'zstd')", + name == NULL ? "" : name); + return EINVAL; +} + struct ArrowIpcSerialDecompressorPrivate { ArrowIpcDecompressFunction decompress_functions[3]; }; diff --git a/src/nanoarrow/ipc/codecs_test.cc b/src/nanoarrow/ipc/codecs_test.cc index 4ed8d6264..826b920b6 100644 --- a/src/nanoarrow/ipc/codecs_test.cc +++ b/src/nanoarrow/ipc/codecs_test.cc @@ -189,6 +189,42 @@ TEST(NanoarrowIpcTest, SerialDecompressor) { "Compression type with value 2 not supported by this build of nanoarrow"); } +TEST(NanoarrowIpcTest, CompressionTypeStrings) { + EXPECT_STREQ(ArrowIpcCompressionTypeToString(NANOARROW_IPC_COMPRESSION_TYPE_NONE), + "none"); + EXPECT_STREQ(ArrowIpcCompressionTypeToString(NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME), + "lz4"); + EXPECT_STREQ(ArrowIpcCompressionTypeToString(NANOARROW_IPC_COMPRESSION_TYPE_ZSTD), + "zstd"); + // 99 is not an enumerator + // NOLINTNEXTLINE(clang-analyzer-optin.core.EnumCastOutOfRange) + auto unknown_type = static_cast(99); + EXPECT_EQ(ArrowIpcCompressionTypeToString(unknown_type), nullptr); + + struct ArrowError error {}; + for (auto type : + {NANOARROW_IPC_COMPRESSION_TYPE_NONE, NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, + NANOARROW_IPC_COMPRESSION_TYPE_ZSTD}) { + enum ArrowIpcCompressionType parsed = unknown_type; + ASSERT_EQ(ArrowIpcCompressionTypeFromString(ArrowIpcCompressionTypeToString(type), + &parsed, &error), + NANOARROW_OK) + << error.message; + EXPECT_EQ(parsed, type); + } + + enum ArrowIpcCompressionType parsed = NANOARROW_IPC_COMPRESSION_TYPE_NONE; + EXPECT_EQ(ArrowIpcCompressionTypeFromString("LZ4", &parsed, &error), EINVAL); + EXPECT_STREQ(error.message, + "Unknown compression type name 'LZ4' (expected 'none', 'lz4', or 'zstd')"); + EXPECT_EQ(ArrowIpcCompressionTypeFromString("", &parsed, &error), EINVAL); + EXPECT_EQ(ArrowIpcCompressionTypeFromString(nullptr, &parsed, &error), EINVAL); + EXPECT_STREQ(error.message, + "Unknown compression type name '' (expected 'none', 'lz4', or 'zstd')"); + // A failed lookup leaves the output untouched + EXPECT_EQ(parsed, NANOARROW_IPC_COMPRESSION_TYPE_NONE); +} + // Compress input at compression_level (appending to a buffer that already has content), // decompress the appended bytes, and check that the result matches the input. Returns // the number of compressed bytes that were appended. diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index 9f2bab675..55fcfb7c4 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -41,6 +41,10 @@ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcSerialCompressor) #define ArrowIpcSerialCompressorSetFunction \ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcSerialCompressorSetFunction) +#define ArrowIpcCompressionTypeToString \ + NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcCompressionTypeToString) +#define ArrowIpcCompressionTypeFromString \ + NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcCompressionTypeFromString) #define ArrowIpcDecoderInit NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcDecoderInit) #define ArrowIpcDecoderReset NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcDecoderReset) #define ArrowIpcDecoderSetDecompressor \ @@ -497,6 +501,20 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcSerialCompressorSetFunction( struct ArrowIpcCompressor* compressor, enum ArrowIpcCompressionType compression_type, ArrowIpcCompressFunction compress_function); +/// \brief Get the name of a compression type +/// +/// Returns "none", "lz4", or "zstd", or NULL for an unknown compression type. +NANOARROW_DLL const char* ArrowIpcCompressionTypeToString( + enum ArrowIpcCompressionType compression_type); + +/// \brief Look up a compression type by name +/// +/// Accepts exactly the names returned by ArrowIpcCompressionTypeToString(). Returns +/// EINVAL (with an error message listing the accepted names) if name is not one of them. +NANOARROW_DLL ArrowErrorCode ArrowIpcCompressionTypeFromString( + const char* name, enum ArrowIpcCompressionType* compression_type_out, + struct ArrowError* error); + /// \brief Decoder for Arrow IPC messages /// /// This structure is intended to be allocated by the caller, From cf9d0604b7f15c6f455a712155c50d8b4a129292 Mon Sep 17 00:00:00 2001 From: Pedro Holanda Date: Tue, 8 Sep 2026 16:26:11 +0200 Subject: [PATCH 09/13] feat(ipc): validate compression levels against ArrowIpcGetCompressionLevelRange() --- src/nanoarrow/ipc/codecs.c | 36 ++++++++++++++++++++++++++++ src/nanoarrow/ipc/codecs_test.cc | 38 +++++++++++++++++++++++++++++ src/nanoarrow/ipc/encoder.c | 35 +++++++++++++++++---------- src/nanoarrow/ipc/encoder_test.cc | 35 ++++++++++++++++++++++++--- src/nanoarrow/ipc/writer_test.cc | 9 +++++++ src/nanoarrow/nanoarrow_ipc.h | 40 +++++++++++++++++++++---------- 6 files changed, 164 insertions(+), 29 deletions(-) diff --git a/src/nanoarrow/ipc/codecs.c b/src/nanoarrow/ipc/codecs.c index 6357ae542..3bdb3d968 100644 --- a/src/nanoarrow/ipc/codecs.c +++ b/src/nanoarrow/ipc/codecs.c @@ -191,6 +191,42 @@ ArrowIpcCompressFunction ArrowIpcGetLZ4CompressionFunction(void) { #endif } +ArrowErrorCode ArrowIpcGetCompressionLevelRange( + enum ArrowIpcCompressionType compression_type, int* min_level_out, + int* max_level_out) { + NANOARROW_DCHECK(min_level_out != NULL && max_level_out != NULL); + NANOARROW_UNUSED(min_level_out); + NANOARROW_UNUSED(max_level_out); + + switch (compression_type) { + case NANOARROW_IPC_COMPRESSION_TYPE_ZSTD: +#if defined(NANOARROW_IPC_WITH_ZSTD) +#if ZSTD_VERSION_NUMBER >= 10400 + *min_level_out = ZSTD_minCLevel(); +#else + // Negative (fast) levels can't be queried before zstd 1.4.0 + *min_level_out = NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT; +#endif + *max_level_out = ZSTD_maxCLevel(); + return NANOARROW_OK; +#else + return ENOTSUP; +#endif + case NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME: +#if defined(NANOARROW_IPC_WITH_LZ4) + // A negative level selects an acceleration of 1 - level, which lz4 caps at 65537 + // (LZ4_ACCELERATION_MAX, which is not part of its public headers) + *min_level_out = 1 - 65537; + *max_level_out = LZ4F_compressionLevel_max(); + return NANOARROW_OK; +#else + return ENOTSUP; +#endif + default: + return EINVAL; + } +} + // The serial decompressor and compressor keep one function per codec, indexed by // enum ArrowIpcCompressionType (NONE is never a codec) static int ArrowIpcCompressionTypeIsCodec(enum ArrowIpcCompressionType compression_type) { diff --git a/src/nanoarrow/ipc/codecs_test.cc b/src/nanoarrow/ipc/codecs_test.cc index 826b920b6..d364f2856 100644 --- a/src/nanoarrow/ipc/codecs_test.cc +++ b/src/nanoarrow/ipc/codecs_test.cc @@ -225,6 +225,44 @@ TEST(NanoarrowIpcTest, CompressionTypeStrings) { EXPECT_EQ(parsed, NANOARROW_IPC_COMPRESSION_TYPE_NONE); } +TEST(NanoarrowIpcTest, CompressionLevelRange) { + int min_level = 1; + int max_level = -1; + EXPECT_EQ(ArrowIpcGetCompressionLevelRange(NANOARROW_IPC_COMPRESSION_TYPE_NONE, + &min_level, &max_level), + EINVAL); + // NOLINTNEXTLINE(clang-analyzer-optin.core.EnumCastOutOfRange) + auto unknown_type = static_cast(99); + EXPECT_EQ(ArrowIpcGetCompressionLevelRange(unknown_type, &min_level, &max_level), + EINVAL); + + if (ArrowIpcGetLZ4CompressionFunction() != nullptr) { + ASSERT_EQ(ArrowIpcGetCompressionLevelRange(NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, + &min_level, &max_level), + NANOARROW_OK); + EXPECT_EQ(min_level, -65536); + EXPECT_EQ(max_level, 12); + } else { + EXPECT_EQ(ArrowIpcGetCompressionLevelRange(NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, + &min_level, &max_level), + ENOTSUP); + } + + if (ArrowIpcGetZstdCompressionFunction() != nullptr) { + ASSERT_EQ(ArrowIpcGetCompressionLevelRange(NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, + &min_level, &max_level), + NANOARROW_OK); + // The levels used by the roundtrip tests below must be in range + EXPECT_LE(min_level, -5); + EXPECT_LE(min_level, NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT); + EXPECT_GE(max_level, 19); + } else { + EXPECT_EQ(ArrowIpcGetCompressionLevelRange(NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, + &min_level, &max_level), + ENOTSUP); + } +} + // Compress input at compression_level (appending to a buffer that already has content), // decompress the appended bytes, and check that the result matches the input. Returns // the number of compressed bytes that were appended. diff --git a/src/nanoarrow/ipc/encoder.c b/src/nanoarrow/ipc/encoder.c index ee36a8127..2ac35f21a 100644 --- a/src/nanoarrow/ipc/encoder.c +++ b/src/nanoarrow/ipc/encoder.c @@ -158,15 +158,10 @@ ArrowErrorCode ArrowIpcEncoderSetCompression( struct ArrowIpcEncoderPrivate* private = (struct ArrowIpcEncoderPrivate*)encoder->private_data; - ArrowIpcCompressFunction built_in = NULL; switch (compression_type) { case NANOARROW_IPC_COMPRESSION_TYPE_NONE: - break; case NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME: - built_in = ArrowIpcGetLZ4CompressionFunction(); - break; case NANOARROW_IPC_COMPRESSION_TYPE_ZSTD: - built_in = ArrowIpcGetZstdCompressionFunction(); break; default: ArrowErrorSet(error, "Unknown compression type with value %d", @@ -176,14 +171,28 @@ ArrowErrorCode ArrowIpcEncoderSetCompression( if (compression_type != NANOARROW_IPC_COMPRESSION_TYPE_NONE) { // With the default compressor, fail now rather than when the first RecordBatch is - // encoded if this build does not support the codec. A custom compressor may support - // codecs that were not built in, so it is only checked when a RecordBatch is encoded. - if (!private->custom_compressor && built_in == NULL) { - ArrowErrorSet( - error, - "Compression type with value %d not supported by this build of nanoarrow", - (int)compression_type); - return ENOTSUP; + // encoded if this build does not support the codec or the level is out of range. + // A custom compressor may support other codecs and levels, so it is only checked + // when a RecordBatch is encoded. + if (!private->custom_compressor) { + int min_level; + int max_level; + if (ArrowIpcGetCompressionLevelRange(compression_type, &min_level, &max_level) != + NANOARROW_OK) { + ArrowErrorSet( + error, + "Compression type with value %d not supported by this build of nanoarrow", + (int)compression_type); + return ENOTSUP; + } + + if (compression_level < min_level || compression_level > max_level) { + ArrowErrorSet( + error, "Compression level %d is out of range for %s (expected %d to %d)", + compression_level, ArrowIpcCompressionTypeToString(compression_type), + min_level, max_level); + return EINVAL; + } } if (private->compressor.release == NULL) { diff --git a/src/nanoarrow/ipc/encoder_test.cc b/src/nanoarrow/ipc/encoder_test.cc index c8338882b..8f22d7f10 100644 --- a/src/nanoarrow/ipc/encoder_test.cc +++ b/src/nanoarrow/ipc/encoder_test.cc @@ -715,6 +715,35 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSetCompressionErrors) { NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), NANOARROW_OK) << error.message; + + // Levels outside the codec's range are rejected when set rather than clamped + int min_level; + int max_level; + ASSERT_EQ(ArrowIpcGetCompressionLevelRange(NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, + &min_level, &max_level), + NANOARROW_OK); + EXPECT_EQ( + ArrowIpcEncoderSetCompression( + encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, max_level + 1, &error), + EINVAL); + EXPECT_EQ(std::string(error.message), + "Compression level " + std::to_string(max_level + 1) + + " is out of range for lz4 (expected " + std::to_string(min_level) + + " to " + std::to_string(max_level) + ")"); + EXPECT_EQ( + ArrowIpcEncoderSetCompression( + encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, min_level - 1, &error), + EINVAL); + EXPECT_EQ( + ArrowIpcEncoderSetCompression( + encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, min_level, &error), + NANOARROW_OK) + << error.message; + EXPECT_EQ( + ArrowIpcEncoderSetCompression( + encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, max_level, &error), + NANOARROW_OK) + << error.message; #else EXPECT_EQ(ArrowIpcEncoderSetCompression( encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, @@ -756,10 +785,10 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSetCompressor) { // The encoder took ownership of the compressor EXPECT_EQ(compressor->release, nullptr); - // With a custom compressor, support is not checked until a batch is encoded + // With a custom compressor, neither codec support nor the level is checked until + // a batch is encoded ASSERT_EQ(ArrowIpcEncoderSetCompression( - encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, - NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), + encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, 1000000, &error), NANOARROW_OK) << error.message; diff --git a/src/nanoarrow/ipc/writer_test.cc b/src/nanoarrow/ipc/writer_test.cc index d753dea83..de3d1b30a 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -366,4 +366,13 @@ TEST(NanoarrowIpcWriter, SetCompressionErrors) { NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), EINVAL); EXPECT_STREQ(error.message, "Unknown compression type with value 99"); + +#if defined(NANOARROW_IPC_WITH_LZ4) + EXPECT_EQ(ArrowIpcWriterSetCompression( + writer.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, 1000000, &error), + EINVAL); + EXPECT_STREQ( + error.message, + "Compression level 1000000 is out of range for lz4 (expected -65536 to 12)"); +#endif } diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index 55fcfb7c4..07878e118 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -45,6 +45,8 @@ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcCompressionTypeToString) #define ArrowIpcCompressionTypeFromString \ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcCompressionTypeFromString) +#define ArrowIpcGetCompressionLevelRange \ + NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcGetCompressionLevelRange) #define ArrowIpcDecoderInit NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcDecoderInit) #define ArrowIpcDecoderReset NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcDecoderReset) #define ArrowIpcDecoderSetDecompressor \ @@ -468,12 +470,11 @@ struct ArrowIpcCompressor { /// reserving sufficient space in dst (e.g., using the compression library's bound /// function) and must only append to dst. /// -/// The interpretation of compression_level is codec-specific. -/// NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT selects the codec's default level; other -/// values follow the underlying library's conventions (for ZSTD, ZSTD_minCLevel() to -/// ZSTD_maxCLevel() where negative levels favour speed; for LZ4, up to -/// LZ4F_compressionLevel_max() where levels >= 3 use LZ4HC and negative levels select -/// acceleration). The built-in implementations clamp out-of-range levels. +/// The interpretation of compression_level is codec-specific: +/// NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT selects the codec's default level and other +/// values follow the underlying library's conventions (see +/// ArrowIpcGetCompressionLevelRange()). ArrowIpcEncoderSetCompression() rejects levels +/// outside that range; the built-in functions clamp them if called directly. typedef ArrowErrorCode (*ArrowIpcCompressFunction)(struct ArrowBufferView src, int compression_level, struct ArrowBuffer* dst, @@ -515,6 +516,18 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcCompressionTypeFromString( const char* name, enum ArrowIpcCompressionType* compression_type_out, struct ArrowError* error); +/// \brief Get the range of compression levels accepted by a codec +/// +/// The range is inclusive and always contains NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT. +/// For ZSTD it is ZSTD_minCLevel() to ZSTD_maxCLevel(), where negative levels favour +/// speed; for LZ4 it is -65536 to LZ4F_compressionLevel_max(), where levels >= 3 use +/// LZ4HC and negative levels select an acceleration of 1 - level. Returns ENOTSUP if +/// the codec is not supported by this build of nanoarrow and EINVAL if compression_type +/// is not a codec (including NANOARROW_IPC_COMPRESSION_TYPE_NONE). +NANOARROW_DLL ArrowErrorCode +ArrowIpcGetCompressionLevelRange(enum ArrowIpcCompressionType compression_type, + int* min_level_out, int* max_level_out); + /// \brief Decoder for Arrow IPC messages /// /// This structure is intended to be allocated by the caller, @@ -981,11 +994,12 @@ ArrowIpcEncoderSetMessageMetadata(struct ArrowIpcEncoder* encoder, /// ArrowIpcCompressFunction for its interpretation. Use /// NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT for the codec's default level. /// -/// Returns EINVAL for an unknown compression type and ENOTSUP if the compression type -/// is not supported by this build of nanoarrow (i.e., nanoarrow was not built with -/// NANOARROW_IPC_WITH_LZ4 or NANOARROW_IPC_WITH_ZSTD). If a custom compressor was -/// set with ArrowIpcEncoderSetCompressor(), support is not checked until a RecordBatch -/// is encoded. +/// Returns EINVAL for an unknown compression type or a compression_level outside the +/// range reported by ArrowIpcGetCompressionLevelRange(), and ENOTSUP if the compression +/// type is not supported by this build of nanoarrow (i.e., nanoarrow was not built with +/// NANOARROW_IPC_WITH_LZ4 or NANOARROW_IPC_WITH_ZSTD). If a custom compressor was set +/// with ArrowIpcEncoderSetCompressor(), neither is checked until a RecordBatch is +/// encoded. NANOARROW_DLL ArrowErrorCode ArrowIpcEncoderSetCompression( struct ArrowIpcEncoder* encoder, enum ArrowIpcCompressionType compression_type, int compression_level, struct ArrowError* error); @@ -1086,8 +1100,8 @@ NANOARROW_DLL void ArrowIpcWriterReset(struct ArrowIpcWriter* writer); /// after this call (in both stream and file mode) and may be changed between batches. /// Use NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT for the codec's default level. /// -/// Returns EINVAL for an unknown compression type and ENOTSUP if the compression type -/// is not supported by this build of nanoarrow. +/// Returns EINVAL for an unknown compression type or an out-of-range compression_level, +/// and ENOTSUP if the compression type is not supported by this build of nanoarrow. NANOARROW_DLL ArrowErrorCode ArrowIpcWriterSetCompression( struct ArrowIpcWriter* writer, enum ArrowIpcCompressionType compression_type, int compression_level, struct ArrowError* error); From a0a392d49b9a600ba3c6c7cdd37b4b2c997b223e Mon Sep 17 00:00:00 2001 From: Pedro Holanda Date: Tue, 8 Sep 2026 17:47:58 +0200 Subject: [PATCH 10/13] fix(ipc): link shared libzstd into the shared library and cover compression error paths --- CMakeLists.txt | 18 ++++- src/nanoarrow/ipc/codecs_test.cc | 89 +++++++++++++++++++++++- src/nanoarrow/ipc/encoder.c | 26 +++---- src/nanoarrow/ipc/encoder_test.cc | 109 ++++++++++++++++++++++++++++-- src/nanoarrow/ipc/writer_test.cc | 7 +- 5 files changed, 222 insertions(+), 27 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 6335b9029..848b933e8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -247,8 +247,21 @@ if(NANOARROW_IPC) # This could be configurable if shared zstd is a must if(TARGET zstd::libzstd_static) set(NANOARROW_IPC_EXTRA_LIBS zstd::libzstd_static) - else() + elseif(TARGET zstd::libzstd) set(NANOARROW_IPC_EXTRA_LIBS zstd::libzstd) + else() + set(NANOARROW_IPC_EXTRA_LIBS zstd::libzstd_shared) + endif() + + # A static libzstd from a system package manager is usually not built with -fPIC + # and can't be linked into a shared object, so the shared nanoarrow_ipc library + # links the shared libzstd instead + if(TARGET zstd::libzstd_shared) + set(NANOARROW_IPC_EXTRA_LIBS_SHARED zstd::libzstd_shared) + elseif(TARGET zstd::libzstd) + set(NANOARROW_IPC_EXTRA_LIBS_SHARED zstd::libzstd) + else() + set(NANOARROW_IPC_EXTRA_LIBS_SHARED zstd::libzstd_static) endif() endif() @@ -276,6 +289,7 @@ if(NANOARROW_IPC) set(NANOARROW_IPC_EXTRA_FLAGS ${NANOARROW_IPC_EXTRA_FLAGS} "-DNANOARROW_IPC_WITH_LZ4") set(NANOARROW_IPC_EXTRA_LIBS ${NANOARROW_IPC_EXTRA_LIBS} lz4::lz4) + set(NANOARROW_IPC_EXTRA_LIBS_SHARED ${NANOARROW_IPC_EXTRA_LIBS_SHARED} lz4::lz4) endif() if(NOT NANOARROW_BUNDLE) @@ -294,7 +308,7 @@ if(NANOARROW_IPC) add_library(nanoarrow_ipc_shared SHARED ${NANOARROW_IPC_BUILD_SOURCES}) target_link_libraries(nanoarrow_ipc_shared - PRIVATE flatccrt ${NANOARROW_IPC_EXTRA_LIBS} + PRIVATE flatccrt ${NANOARROW_IPC_EXTRA_LIBS_SHARED} PUBLIC nanoarrow_shared nanoarrow_coverage_config) install(FILES src/nanoarrow/nanoarrow_ipc.h src/nanoarrow/nanoarrow_ipc.hpp diff --git a/src/nanoarrow/ipc/codecs_test.cc b/src/nanoarrow/ipc/codecs_test.cc index d364f2856..3d125f2a0 100644 --- a/src/nanoarrow/ipc/codecs_test.cc +++ b/src/nanoarrow/ipc/codecs_test.cc @@ -155,6 +155,13 @@ TEST(NanoarrowIpcTest, SerialDecompressor) { decompressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_NONE, nullptr), EINVAL); + // NONE is not a codec that can be used to decompress + EXPECT_EQ(decompressor->decompress_add(decompressor.get(), + NANOARROW_IPC_COMPRESSION_TYPE_NONE, + {{nullptr}, 0}, nullptr, 0, &error), + EINVAL); + EXPECT_STREQ(error.message, "Unknown decompression type with value 0"); + // The serial decompressor never waits and always succeeds when requested to EXPECT_EQ(decompressor->decompress_wait(decompressor.get(), 0, &error), NANOARROW_OK); @@ -196,9 +203,10 @@ TEST(NanoarrowIpcTest, CompressionTypeStrings) { "lz4"); EXPECT_STREQ(ArrowIpcCompressionTypeToString(NANOARROW_IPC_COMPRESSION_TYPE_ZSTD), "zstd"); - // 99 is not an enumerator + // 3 is not an enumerator but is within the enum's value range (unlike, e.g., 99, + // which C++ can't represent in this enum) // NOLINTNEXTLINE(clang-analyzer-optin.core.EnumCastOutOfRange) - auto unknown_type = static_cast(99); + auto unknown_type = static_cast(3); EXPECT_EQ(ArrowIpcCompressionTypeToString(unknown_type), nullptr); struct ArrowError error {}; @@ -232,7 +240,7 @@ TEST(NanoarrowIpcTest, CompressionLevelRange) { &min_level, &max_level), EINVAL); // NOLINTNEXTLINE(clang-analyzer-optin.core.EnumCastOutOfRange) - auto unknown_type = static_cast(99); + auto unknown_type = static_cast(3); EXPECT_EQ(ArrowIpcGetCompressionLevelRange(unknown_type, &min_level, &max_level), EINVAL); @@ -393,6 +401,69 @@ TEST(NanoarrowIpcTest, LZ4CompressMinimumLevels) { } } +// An allocator whose reallocate() fails on the fail_on-th call (1-based) and otherwise +// delegates to the default allocator +struct FailingAllocatorState { + int calls; + int fail_on; +}; + +static uint8_t* FailingReallocate(struct ArrowBufferAllocator* allocator, uint8_t* ptr, + int64_t old_size, int64_t new_size) { + auto* state = static_cast(allocator->private_data); + auto default_allocator = ArrowBufferAllocatorDefault(); + if (++state->calls == state->fail_on) { + // nanoarrow discards the buffer on failure, so the old allocation is freed here + default_allocator.free(&default_allocator, ptr, old_size); + return nullptr; + } + return default_allocator.reallocate(&default_allocator, ptr, old_size, new_size); +} + +static void FailingFree(struct ArrowBufferAllocator* allocator, uint8_t* ptr, + int64_t size) { + NANOARROW_UNUSED(allocator); + auto default_allocator = ArrowBufferAllocatorDefault(); + default_allocator.free(&default_allocator, ptr, size); +} + +static struct ArrowBufferAllocator FailingAllocator(FailingAllocatorState* state) { + struct ArrowBufferAllocator allocator = ArrowBufferAllocatorDefault(); + allocator.reallocate = &FailingReallocate; + allocator.free = &FailingFree; + allocator.private_data = state; + return allocator; +} + +TEST(NanoarrowIpcTest, CompressAllocationFailure) { + struct ArrowError error {}; + for (auto compress : + {ArrowIpcGetLZ4CompressionFunction(), ArrowIpcGetZstdCompressionFunction()}) { + if (compress == nullptr) { + continue; + } + + FailingAllocatorState state{0, 1}; + nanoarrow::UniqueBuffer dst; + ASSERT_EQ(ArrowBufferSetAllocator(dst.get(), FailingAllocator(&state)), NANOARROW_OK); + EXPECT_EQ(compress({{kUncompressed012}, sizeof(kUncompressed012)}, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, dst.get(), &error), + ENOMEM); + EXPECT_THAT(error.message, ::testing::HasSubstr("ArrowBufferReserve")); + EXPECT_EQ(state.calls, 1); + } +} + +// A stand-in compression function that always fails +static ArrowErrorCode FailCompress(struct ArrowBufferView src, int compression_level, + struct ArrowBuffer* dst, struct ArrowError* error) { + NANOARROW_UNUSED(src); + NANOARROW_UNUSED(compression_level); + NANOARROW_UNUSED(dst); + ArrowErrorSet(error, "FailCompress() failed"); + return EIO; +} + // A stand-in compression function that records the level it was called with and // "compresses" by copying static int last_compression_level = 0; @@ -477,4 +548,16 @@ TEST(NanoarrowIpcTest, SerialCompressor) { EXPECT_EQ(last_compression_level, 7); ASSERT_EQ(dst->size_bytes, static_cast(sizeof(kUncompressed012))); EXPECT_EQ(std::memcmp(dst->data, kUncompressed012, sizeof(kUncompressed012)), 0); + + // Errors from the function for the codec are propagated + ASSERT_EQ( + ArrowIpcSerialCompressorSetFunction( + compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, &FailCompress), + NANOARROW_OK); + EXPECT_EQ(compressor->compress( + compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, + {{kUncompressed012}, sizeof(kUncompressed012)}, dst.get(), &error), + EIO); + EXPECT_STREQ(error.message, "FailCompress() failed"); } diff --git a/src/nanoarrow/ipc/encoder.c b/src/nanoarrow/ipc/encoder.c index 2ac35f21a..eaf8d240d 100644 --- a/src/nanoarrow/ipc/encoder.c +++ b/src/nanoarrow/ipc/encoder.c @@ -57,6 +57,8 @@ struct ArrowIpcEncoderPrivate { int custom_compressor; // Compression level passed to the compressor when codec != NONE int compression_level; + // The flatbuffer equivalent of codec (only meaningful when codec != NONE) + ns(CompressionType_enum_t) flatbuf_codec; }; ArrowErrorCode ArrowIpcEncoderInit(struct ArrowIpcEncoder* encoder) { @@ -81,6 +83,7 @@ ArrowErrorCode ArrowIpcEncoderInit(struct ArrowIpcEncoder* encoder) { private->compressor.release = NULL; private->custom_compressor = 0; private->compression_level = NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT; + private->flatbuf_codec = ns(CompressionType_LZ4_FRAME); return NANOARROW_OK; } @@ -158,10 +161,15 @@ ArrowErrorCode ArrowIpcEncoderSetCompression( struct ArrowIpcEncoderPrivate* private = (struct ArrowIpcEncoderPrivate*)encoder->private_data; + ns(CompressionType_enum_t) flatbuf_codec = ns(CompressionType_LZ4_FRAME); switch (compression_type) { case NANOARROW_IPC_COMPRESSION_TYPE_NONE: + break; case NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME: + flatbuf_codec = ns(CompressionType_LZ4_FRAME); + break; case NANOARROW_IPC_COMPRESSION_TYPE_ZSTD: + flatbuf_codec = ns(CompressionType_ZSTD); break; default: ArrowErrorSet(error, "Unknown compression type with value %d", @@ -202,6 +210,7 @@ ArrowErrorCode ArrowIpcEncoderSetCompression( } private->codec = compression_type; + private->flatbuf_codec = flatbuf_codec; private->compression_level = compression_level; return NANOARROW_OK; } @@ -878,22 +887,9 @@ static ArrowErrorCode ArrowIpcEncoderEncodeRecordBatch( FLATCC_RETURN_UNLESS_0(RecordBatch_length_add(builder, array_view->length), error); if (private->codec != NANOARROW_IPC_COMPRESSION_TYPE_NONE) { - ns(CompressionType_enum_t) codec; - switch (private->codec) { - case NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME: - codec = ns(CompressionType_LZ4_FRAME); - break; - case NANOARROW_IPC_COMPRESSION_TYPE_ZSTD: - codec = ns(CompressionType_ZSTD); - break; - default: - ArrowErrorSet(error, "Unknown compression type with value %d", - (int)private->codec); - return EINVAL; - } - FLATCC_RETURN_UNLESS_0(RecordBatch_compression_start(builder), error); - FLATCC_RETURN_UNLESS_0(BodyCompression_codec_add(builder, codec), error); + FLATCC_RETURN_UNLESS_0(BodyCompression_codec_add(builder, private->flatbuf_codec), + error); FLATCC_RETURN_UNLESS_0( BodyCompression_method_add(builder, ns(BodyCompressionMethod_BUFFER)), error); FLATCC_RETURN_UNLESS_0(RecordBatch_compression_end(builder), error); diff --git a/src/nanoarrow/ipc/encoder_test.cc b/src/nanoarrow/ipc/encoder_test.cc index 8f22d7f10..21ffa4fbd 100644 --- a/src/nanoarrow/ipc/encoder_test.cc +++ b/src/nanoarrow/ipc/encoder_test.cc @@ -691,14 +691,15 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSetCompressionErrors) { ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); struct ArrowError error; - // 99 is not an enumerator; it exercises the EINVAL path + // 3 is not an enumerator but is within the enum's value range (unlike, e.g., 99, + // which C++ can't represent in this enum); it exercises the EINVAL path // NOLINTNEXTLINE(clang-analyzer-optin.core.EnumCastOutOfRange) - auto unknown_type = static_cast(99); + auto unknown_type = static_cast(3); EXPECT_EQ( ArrowIpcEncoderSetCompression(encoder.get(), unknown_type, NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), EINVAL); - EXPECT_STREQ(error.message, "Unknown compression type with value 99"); + EXPECT_STREQ(error.message, "Unknown compression type with value 3"); // NONE is always supported EXPECT_EQ( @@ -769,11 +770,30 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSetCompressionErrors) { #endif } +static void (*original_compressor_release)(struct ArrowIpcCompressor*) = nullptr; +static int compressor_release_calls = 0; + +static void CountingCompressorRelease(struct ArrowIpcCompressor* compressor) { + compressor_release_calls++; + original_compressor_release(compressor); +} + TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSetCompressor) { nanoarrow::ipc::UniqueEncoder encoder; ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); struct ArrowError error; + // A compressor whose release we can observe + nanoarrow::ipc::UniqueCompressor first_compressor; + ASSERT_EQ(ArrowIpcSerialCompressor(first_compressor.get()), NANOARROW_OK); + original_compressor_release = first_compressor->release; + first_compressor->release = &CountingCompressorRelease; + compressor_release_calls = 0; + ASSERT_EQ(ArrowIpcEncoderSetCompressor(encoder.get(), first_compressor.get()), + NANOARROW_OK); + EXPECT_EQ(first_compressor->release, nullptr); + EXPECT_EQ(compressor_release_calls, 0); + // A custom compressor that explicitly does not support LZ4 nanoarrow::ipc::UniqueCompressor compressor; ASSERT_EQ(ArrowIpcSerialCompressor(compressor.get()), NANOARROW_OK); @@ -782,8 +802,9 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSetCompressor) { NANOARROW_OK); ASSERT_EQ(ArrowIpcEncoderSetCompressor(encoder.get(), compressor.get()), NANOARROW_OK); - // The encoder took ownership of the compressor + // The encoder took ownership of the compressor and released the previous one EXPECT_EQ(compressor->release, nullptr); + EXPECT_EQ(compressor_release_calls, 1); // With a custom compressor, neither codec support nor the level is checked until // a batch is encoded @@ -879,3 +900,83 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderCompressionLevel) { NANOARROW_OK); EXPECT_EQ(is_equal, 1) << error.message; } + +// An allocator whose reallocate() fails on the fail_on-th call (1-based) and otherwise +// delegates to the default allocator +struct FailingAllocatorState { + int calls; + int fail_on; +}; + +static uint8_t* FailingReallocate(struct ArrowBufferAllocator* allocator, uint8_t* ptr, + int64_t old_size, int64_t new_size) { + auto* state = static_cast(allocator->private_data); + auto default_allocator = ArrowBufferAllocatorDefault(); + if (++state->calls == state->fail_on) { + // nanoarrow discards the buffer on failure, so the old allocation is freed here + default_allocator.free(&default_allocator, ptr, old_size); + return nullptr; + } + return default_allocator.reallocate(&default_allocator, ptr, old_size, new_size); +} + +static void FailingFree(struct ArrowBufferAllocator* allocator, uint8_t* ptr, + int64_t size) { + NANOARROW_UNUSED(allocator); + auto default_allocator = ArrowBufferAllocatorDefault(); + default_allocator.free(&default_allocator, ptr, size); +} + +static struct ArrowBufferAllocator FailingAllocator(FailingAllocatorState* state) { + struct ArrowBufferAllocator allocator = ArrowBufferAllocatorDefault(); + allocator.reallocate = &FailingReallocate; + allocator.free = &FailingFree; + allocator.private_data = state; + return allocator; +} + +// Encode a batch with a body allocator that fails on the fail_on-th allocation, for +// every fail_on until encoding succeeds, so that each allocation site reports ENOMEM +static void TestEncodeAllocationFailures(enum ArrowIpcCompressionType codec) { + struct ArrowError error; + CompressibleRecordBatch batch; + + int fail_on = 1; + for (; fail_on < 100; fail_on++) { + SCOPED_TRACE("fail_on " + std::to_string(fail_on)); + // A fresh encoder each time so that a failed encode can't affect the next one + nanoarrow::ipc::UniqueEncoder encoder; + ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); + ASSERT_EQ(ArrowIpcEncoderSetCompression( + encoder.get(), codec, NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), + NANOARROW_OK) + << error.message; + + FailingAllocatorState state{0, fail_on}; + nanoarrow::UniqueBuffer body; + ASSERT_EQ(ArrowBufferSetAllocator(body.get(), FailingAllocator(&state)), + NANOARROW_OK); + int result = ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), batch.array_view(), + body.get(), &error); + if (state.calls < fail_on) { + // No allocation failed, so this is one more than the number of allocations + EXPECT_EQ(result, NANOARROW_OK) << error.message; + break; + } + EXPECT_EQ(result, ENOMEM); + } + + EXPECT_GT(fail_on, 1); + EXPECT_LT(fail_on, 100); +} + +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderUncompressedAllocationFailures) { + TestEncodeAllocationFailures(NANOARROW_IPC_COMPRESSION_TYPE_NONE); +} + +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderCompressedAllocationFailures) { + if (ArrowIpcGetLZ4CompressionFunction() == nullptr) { + GTEST_SKIP() << "nanoarrow_ipc not built with NANOARROW_IPC_WITH_LZ4"; + } + TestEncodeAllocationFailures(NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME); +} diff --git a/src/nanoarrow/ipc/writer_test.cc b/src/nanoarrow/ipc/writer_test.cc index de3d1b30a..9a9421f5e 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -359,13 +359,14 @@ TEST(NanoarrowIpcWriter, SetCompressionErrors) { ASSERT_EQ(ArrowIpcWriterInit(writer.get(), stream.get()), NANOARROW_OK); struct ArrowError error; - // 99 is not an enumerator; it exercises the EINVAL path + // 3 is not an enumerator but is within the enum's value range (unlike, e.g., 99, + // which C++ can't represent in this enum); it exercises the EINVAL path // NOLINTNEXTLINE(clang-analyzer-optin.core.EnumCastOutOfRange) - auto unknown_type = static_cast(99); + auto unknown_type = static_cast(3); EXPECT_EQ(ArrowIpcWriterSetCompression(writer.get(), unknown_type, NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), EINVAL); - EXPECT_STREQ(error.message, "Unknown compression type with value 99"); + EXPECT_STREQ(error.message, "Unknown compression type with value 3"); #if defined(NANOARROW_IPC_WITH_LZ4) EXPECT_EQ(ArrowIpcWriterSetCompression( From 7e7e96ab78b0d0b176a238754325cec4eab5f06f Mon Sep 17 00:00:00 2001 From: Pedro Holanda Date: Wed, 9 Sep 2026 12:21:32 +0200 Subject: [PATCH 11/13] refactor(ipc): carry codec and level in ArrowIpcCompressor and never return NULL from ToString --- src/nanoarrow/ipc/codecs.c | 12 +-- src/nanoarrow/ipc/codecs_test.cc | 39 +++++----- src/nanoarrow/ipc/encoder.c | 118 +++++++++++++++--------------- src/nanoarrow/ipc/encoder_test.cc | 29 +++++--- src/nanoarrow/nanoarrow_ipc.h | 75 +++++++++++-------- src/nanoarrow/nanoarrow_ipc.hpp | 2 + 6 files changed, 149 insertions(+), 126 deletions(-) diff --git a/src/nanoarrow/ipc/codecs.c b/src/nanoarrow/ipc/codecs.c index 3bdb3d968..a2c0d21e9 100644 --- a/src/nanoarrow/ipc/codecs.c +++ b/src/nanoarrow/ipc/codecs.c @@ -249,7 +249,7 @@ const char* ArrowIpcCompressionTypeToString( case NANOARROW_IPC_COMPRESSION_TYPE_ZSTD: return "zstd"; default: - return NULL; + return ""; } } @@ -358,11 +358,11 @@ struct ArrowIpcSerialCompressorPrivate { }; static ArrowErrorCode ArrowIpcSerialCompressorCompress( - struct ArrowIpcCompressor* compressor, enum ArrowIpcCompressionType compression_type, - int compression_level, struct ArrowBufferView src, struct ArrowBuffer* dst, - struct ArrowError* error) { + struct ArrowIpcCompressor* compressor, struct ArrowBufferView src, + struct ArrowBuffer* dst, struct ArrowError* error) { struct ArrowIpcSerialCompressorPrivate* private_data = (struct ArrowIpcSerialCompressorPrivate*)compressor->private_data; + enum ArrowIpcCompressionType compression_type = compressor->compression_type; if (!ArrowIpcCompressionTypeIsCodec(compression_type)) { ArrowErrorSet(error, "Unknown compression type with value %d", (int)compression_type); @@ -377,7 +377,7 @@ static ArrowErrorCode ArrowIpcSerialCompressorCompress( return ENOTSUP; } - NANOARROW_RETURN_NOT_OK(fn(src, compression_level, dst, error)); + NANOARROW_RETURN_NOT_OK(fn(src, compressor->compression_level, dst, error)); return NANOARROW_OK; } @@ -387,6 +387,8 @@ static void ArrowIpcSerialCompressorRelease(struct ArrowIpcCompressor* compresso } ArrowErrorCode ArrowIpcSerialCompressor(struct ArrowIpcCompressor* compressor) { + compressor->compression_type = NANOARROW_IPC_COMPRESSION_TYPE_NONE; + compressor->compression_level = NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT; compressor->release = NULL; compressor->private_data = ArrowMalloc(sizeof(struct ArrowIpcSerialCompressorPrivate)); if (compressor->private_data == NULL) { diff --git a/src/nanoarrow/ipc/codecs_test.cc b/src/nanoarrow/ipc/codecs_test.cc index 3d125f2a0..f8f19f625 100644 --- a/src/nanoarrow/ipc/codecs_test.cc +++ b/src/nanoarrow/ipc/codecs_test.cc @@ -207,7 +207,8 @@ TEST(NanoarrowIpcTest, CompressionTypeStrings) { // which C++ can't represent in this enum) // NOLINTNEXTLINE(clang-analyzer-optin.core.EnumCastOutOfRange) auto unknown_type = static_cast(3); - EXPECT_EQ(ArrowIpcCompressionTypeToString(unknown_type), nullptr); + EXPECT_STREQ(ArrowIpcCompressionTypeToString(unknown_type), + ""); struct ArrowError error {}; for (auto type : @@ -226,6 +227,9 @@ TEST(NanoarrowIpcTest, CompressionTypeStrings) { EXPECT_STREQ(error.message, "Unknown compression type name 'LZ4' (expected 'none', 'lz4', or 'zstd')"); EXPECT_EQ(ArrowIpcCompressionTypeFromString("", &parsed, &error), EINVAL); + EXPECT_EQ( + ArrowIpcCompressionTypeFromString("", &parsed, &error), + EINVAL); EXPECT_EQ(ArrowIpcCompressionTypeFromString(nullptr, &parsed, &error), EINVAL); EXPECT_STREQ(error.message, "Unknown compression type name '' (expected 'none', 'lz4', or 'zstd')"); @@ -481,6 +485,8 @@ TEST(NanoarrowIpcTest, SerialCompressor) { nanoarrow::ipc::UniqueCompressor compressor; ASSERT_EQ(ArrowIpcSerialCompressor(compressor.get()), NANOARROW_OK); + EXPECT_EQ(compressor->compression_type, NANOARROW_IPC_COMPRESSION_TYPE_NONE); + EXPECT_EQ(compressor->compression_level, NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT); // Check the function setter error ASSERT_EQ(ArrowIpcSerialCompressorSetFunction( @@ -489,16 +495,14 @@ TEST(NanoarrowIpcTest, SerialCompressor) { // NONE is not a codec that can be used to compress nanoarrow::UniqueBuffer dst; - EXPECT_EQ(compressor->compress(compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_NONE, - NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, {{nullptr}, 0}, - dst.get(), &error), + EXPECT_EQ(compressor->compress(compressor.get(), {{nullptr}, 0}, dst.get(), &error), EINVAL); EXPECT_STREQ(error.message, "Unknown compression type with value 0"); // Check a compress for a supported codec if we have one (or for an error if we don't) + compressor->compression_type = NANOARROW_IPC_COMPRESSION_TYPE_ZSTD; if (ArrowIpcGetZstdCompressionFunction() != nullptr) { - ASSERT_EQ(compressor->compress(compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, - NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, + ASSERT_EQ(compressor->compress(compressor.get(), {{kUncompressed012}, sizeof(kUncompressed012)}, dst.get(), &error), NANOARROW_OK) @@ -513,9 +517,7 @@ TEST(NanoarrowIpcTest, SerialCompressor) { << error.message; EXPECT_TRUE(std::memcmp(out, kUncompressed012, sizeof(kUncompressed012)) == 0); } else { - EXPECT_EQ(compressor->compress(compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, - NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, - {{nullptr}, 0}, dst.get(), &error), + EXPECT_EQ(compressor->compress(compressor.get(), {{nullptr}, 0}, dst.get(), &error), ENOTSUP); EXPECT_STREQ( error.message, @@ -526,9 +528,7 @@ TEST(NanoarrowIpcTest, SerialCompressor) { ASSERT_EQ(ArrowIpcSerialCompressorSetFunction( compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, nullptr), NANOARROW_OK); - EXPECT_EQ(compressor->compress(compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, - NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, {{nullptr}, 0}, - dst.get(), &error), + EXPECT_EQ(compressor->compress(compressor.get(), {{nullptr}, 0}, dst.get(), &error), ENOTSUP); EXPECT_STREQ(error.message, "Compression type with value 2 not supported by this build of nanoarrow"); @@ -538,11 +538,13 @@ TEST(NanoarrowIpcTest, SerialCompressor) { NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, &RecordLevelAndCopy), NANOARROW_OK); + compressor->compression_type = NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME; + compressor->compression_level = 7; dst->size_bytes = 0; last_compression_level = 0; - ASSERT_EQ(compressor->compress( - compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, 7, - {{kUncompressed012}, sizeof(kUncompressed012)}, dst.get(), &error), + ASSERT_EQ(compressor->compress(compressor.get(), + {{kUncompressed012}, sizeof(kUncompressed012)}, + dst.get(), &error), NANOARROW_OK) << error.message; EXPECT_EQ(last_compression_level, 7); @@ -554,10 +556,9 @@ TEST(NanoarrowIpcTest, SerialCompressor) { ArrowIpcSerialCompressorSetFunction( compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, &FailCompress), NANOARROW_OK); - EXPECT_EQ(compressor->compress( - compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, - NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, - {{kUncompressed012}, sizeof(kUncompressed012)}, dst.get(), &error), + EXPECT_EQ(compressor->compress(compressor.get(), + {{kUncompressed012}, sizeof(kUncompressed012)}, + dst.get(), &error), EIO); EXPECT_STREQ(error.message, "FailCompress() failed"); } diff --git a/src/nanoarrow/ipc/encoder.c b/src/nanoarrow/ipc/encoder.c index a8bea61bf..6c29f3352 100644 --- a/src/nanoarrow/ipc/encoder.c +++ b/src/nanoarrow/ipc/encoder.c @@ -49,16 +49,9 @@ struct ArrowIpcEncoderPrivate { // Metadata to attach to the next encoded Message (in nanoarrow's packed // representation), or an empty buffer if the next Message has no metadata. struct ArrowBuffer message_metadata; - // Compression applied to the body buffers of subsequently encoded RecordBatches - enum ArrowIpcCompressionType codec; - // Compressor used when codec != NONE (release is NULL until one is needed) + // Compressor for the body buffers of subsequently encoded messages (release is + // NULL when they are not compressed) struct ArrowIpcCompressor compressor; - // Whether compressor was provided by ArrowIpcEncoderSetCompressor() - int custom_compressor; - // Compression level passed to the compressor when codec != NONE - int compression_level; - // The flatbuffer equivalent of codec (only meaningful when codec != NONE) - ns(CompressionType_enum_t) flatbuf_codec; }; ArrowErrorCode ArrowIpcEncoderInit(struct ArrowIpcEncoder* encoder) { @@ -79,11 +72,7 @@ ArrowErrorCode ArrowIpcEncoderInit(struct ArrowIpcEncoder* encoder) { ArrowBufferInit(&private->nodes); ArrowIpcDictionaryEncodingsInit(&private->dictionary_encodings); ArrowBufferInit(&private->message_metadata); - private->codec = NANOARROW_IPC_COMPRESSION_TYPE_NONE; private->compressor.release = NULL; - private->custom_compressor = 0; - private->compression_level = NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT; - private->flatbuf_codec = ns(CompressionType_LZ4_FRAME); return NANOARROW_OK; } @@ -150,7 +139,6 @@ ArrowErrorCode ArrowIpcEncoderSetCompressor(struct ArrowIpcEncoder* encoder, memcpy(&private->compressor, compressor, sizeof(struct ArrowIpcCompressor)); compressor->release = NULL; - private->custom_compressor = 1; return NANOARROW_OK; } @@ -161,15 +149,14 @@ ArrowErrorCode ArrowIpcEncoderSetCompression( struct ArrowIpcEncoderPrivate* private = (struct ArrowIpcEncoderPrivate*)encoder->private_data; - ns(CompressionType_enum_t) flatbuf_codec = ns(CompressionType_LZ4_FRAME); switch (compression_type) { case NANOARROW_IPC_COMPRESSION_TYPE_NONE: - break; + if (private->compressor.release != NULL) { + private->compressor.release(&private->compressor); + } + return NANOARROW_OK; case NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME: - flatbuf_codec = ns(CompressionType_LZ4_FRAME); - break; case NANOARROW_IPC_COMPRESSION_TYPE_ZSTD: - flatbuf_codec = ns(CompressionType_ZSTD); break; default: ArrowErrorSet(error, "Unknown compression type with value %d", @@ -177,42 +164,31 @@ ArrowErrorCode ArrowIpcEncoderSetCompression( return EINVAL; } - if (compression_type != NANOARROW_IPC_COMPRESSION_TYPE_NONE) { - // With the default compressor, fail now rather than when the first RecordBatch is - // encoded if this build does not support the codec or the level is out of range. - // A custom compressor may support other codecs and levels, so it is only checked - // when a RecordBatch is encoded. - if (!private->custom_compressor) { - int min_level; - int max_level; - if (ArrowIpcGetCompressionLevelRange(compression_type, &min_level, &max_level) != - NANOARROW_OK) { - ArrowErrorSet( - error, - "Compression type with value %d not supported by this build of nanoarrow", - (int)compression_type); - return ENOTSUP; - } - - if (compression_level < min_level || compression_level > max_level) { - ArrowErrorSet( - error, "Compression level %d is out of range for %s (expected %d to %d)", - compression_level, ArrowIpcCompressionTypeToString(compression_type), - min_level, max_level); - return EINVAL; - } - } + // Fail now rather than when the first message is encoded if this build does not + // support the codec or the level is out of range + int min_level; + int max_level; + if (ArrowIpcGetCompressionLevelRange(compression_type, &min_level, &max_level) != + NANOARROW_OK) { + ArrowErrorSet( + error, "Compression type with value %d not supported by this build of nanoarrow", + (int)compression_type); + return ENOTSUP; + } - if (private->compressor.release == NULL) { - NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowIpcSerialCompressor(&private->compressor), - error); - } + if (compression_level < min_level || compression_level > max_level) { + ArrowErrorSet(error, + "Compression level %d is out of range for %s (expected %d to %d)", + compression_level, ArrowIpcCompressionTypeToString(compression_type), + min_level, max_level); + return EINVAL; } - private->codec = compression_type; - private->flatbuf_codec = flatbuf_codec; - private->compression_level = compression_level; - return NANOARROW_OK; + struct ArrowIpcCompressor compressor; + NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowIpcSerialCompressor(&compressor), error); + compressor.compression_type = compression_type; + compressor.compression_level = compression_level; + return ArrowIpcEncoderSetCompressor(encoder, &compressor); } static ArrowErrorCode ArrowIpcEncoderWriteContinuationAndSize(struct ArrowBuffer* out, @@ -755,21 +731,29 @@ struct ArrowIpcBufferEncoder { int64_t body_length; }; +// The codec applied to message bodies (NONE when no compressor is set) +static enum ArrowIpcCompressionType ArrowIpcEncoderCodec( + struct ArrowIpcEncoderPrivate* private) { + if (private->compressor.release == NULL) { + return NANOARROW_IPC_COMPRESSION_TYPE_NONE; + } + return private->compressor.compression_type; +} + // Append buffer_view to body_buffer as a compressed IPC buffer: the uncompressed length // as a little-endian int64 followed by the compressed bytes. If compression does not // reduce the size, the buffer is stored uncompressed with a length prefix of -1 instead. static ArrowErrorCode ArrowIpcEncoderAppendCompressedBuffer( struct ArrowIpcEncoderPrivate* private, struct ArrowBufferView buffer_view, struct ArrowBuffer* body_buffer, struct ArrowError* error) { - NANOARROW_DCHECK(private->compressor.release != NULL); + NANOARROW_DCHECK(ArrowIpcEncoderCodec(private) != NANOARROW_IPC_COMPRESSION_TYPE_NONE); // placeholder for the prefix, then compress directly into the body int64_t prefix_offset = body_buffer->size_bytes; int64_t payload_offset = prefix_offset + (int64_t)sizeof(int64_t); NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowBufferAppendInt64(body_buffer, 0), error); - NANOARROW_RETURN_NOT_OK(private->compressor.compress( - &private->compressor, private->codec, private->compression_level, buffer_view, - body_buffer, error)); + NANOARROW_RETURN_NOT_OK(private->compressor.compress(&private->compressor, buffer_view, + body_buffer, error)); int64_t prefix = buffer_view.size_bytes; if (body_buffer->size_bytes - payload_offset >= buffer_view.size_bytes) { @@ -800,7 +784,8 @@ static ArrowErrorCode ArrowIpcEncoderBuildContiguousBodyBufferCallback( int64_t buffer_begin = _ArrowRoundUpToMultipleOf8(body_buffer->size_bytes); // Empty buffers are never compressed (nor length-prefixed), matching Arrow C++. int needs_compression = - private->codec != NANOARROW_IPC_COMPRESSION_TYPE_NONE && buffer_view.size_bytes > 0; + ArrowIpcEncoderCodec(private) != NANOARROW_IPC_COMPRESSION_TYPE_NONE && + buffer_view.size_bytes > 0; if (!needs_compression) { // Reserve the data and padding together to avoid growing the buffer twice. int64_t new_size = _ArrowRoundUpToMultipleOf8(buffer_begin + buffer_view.size_bytes); @@ -842,14 +827,25 @@ static ArrowErrorCode ArrowIpcEncoderBuildContiguousBodyBufferCallback( // built by the same buffer encoder, so both need this. static ArrowErrorCode ArrowIpcEncoderEncodeBodyCompression( struct ArrowIpcEncoderPrivate* private, struct ArrowError* error) { - if (private->codec == NANOARROW_IPC_COMPRESSION_TYPE_NONE) { - return NANOARROW_OK; + ns(CompressionType_enum_t) codec; + switch (ArrowIpcEncoderCodec(private)) { + case NANOARROW_IPC_COMPRESSION_TYPE_NONE: + return NANOARROW_OK; + case NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME: + codec = ns(CompressionType_LZ4_FRAME); + break; + case NANOARROW_IPC_COMPRESSION_TYPE_ZSTD: + codec = ns(CompressionType_ZSTD); + break; + default: + ArrowErrorSet(error, "Unknown compression type with value %d", + (int)private->compressor.compression_type); + return EINVAL; } flatcc_builder_t* builder = &private->builder; FLATCC_RETURN_UNLESS_0(RecordBatch_compression_start(builder), error); - FLATCC_RETURN_UNLESS_0(BodyCompression_codec_add(builder, private->flatbuf_codec), - error); + FLATCC_RETURN_UNLESS_0(BodyCompression_codec_add(builder, codec), error); FLATCC_RETURN_UNLESS_0( BodyCompression_method_add(builder, ns(BodyCompressionMethod_BUFFER)), error); FLATCC_RETURN_UNLESS_0(RecordBatch_compression_end(builder), error); diff --git a/src/nanoarrow/ipc/encoder_test.cc b/src/nanoarrow/ipc/encoder_test.cc index af78920c8..80efb9981 100644 --- a/src/nanoarrow/ipc/encoder_test.cc +++ b/src/nanoarrow/ipc/encoder_test.cc @@ -865,25 +865,20 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSetCompressor) { EXPECT_EQ(first_compressor->release, nullptr); EXPECT_EQ(compressor_release_calls, 0); - // A custom compressor that explicitly does not support LZ4 + // A custom compressor configured for LZ4 that explicitly does not support it nanoarrow::ipc::UniqueCompressor compressor; ASSERT_EQ(ArrowIpcSerialCompressor(compressor.get()), NANOARROW_OK); ASSERT_EQ(ArrowIpcSerialCompressorSetFunction( compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, nullptr), NANOARROW_OK); + compressor->compression_type = NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME; ASSERT_EQ(ArrowIpcEncoderSetCompressor(encoder.get(), compressor.get()), NANOARROW_OK); // The encoder took ownership of the compressor and released the previous one EXPECT_EQ(compressor->release, nullptr); EXPECT_EQ(compressor_release_calls, 1); - // With a custom compressor, neither codec support nor the level is checked until - // a batch is encoded - ASSERT_EQ(ArrowIpcEncoderSetCompression( - encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, 1000000, &error), - NANOARROW_OK) - << error.message; - + // With a custom compressor, support is not checked until a batch is encoded CompressibleRecordBatch batch; nanoarrow::UniqueBuffer body; EXPECT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), batch.array_view(), @@ -891,6 +886,18 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSetCompressor) { ENOTSUP); EXPECT_STREQ(error.message, "Compression type with value 1 not supported by this build of nanoarrow"); + + // NONE removes the custom compressor and batches are encoded uncompressed again + ASSERT_EQ( + ArrowIpcEncoderSetCompression(encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_NONE, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), + NANOARROW_OK) + << error.message; + body->size_bytes = 0; + EXPECT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), batch.array_view(), + body.get(), &error), + NANOARROW_OK) + << error.message; } // A stand-in compression function that records the level it was called with and @@ -918,11 +925,9 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderCompressionLevel) { ArrowIpcSerialCompressorSetFunction( compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, &RecordLevelAndCopy), NANOARROW_OK); + compressor->compression_type = NANOARROW_IPC_COMPRESSION_TYPE_ZSTD; + compressor->compression_level = 11; ASSERT_EQ(ArrowIpcEncoderSetCompressor(encoder.get(), compressor.get()), NANOARROW_OK); - ASSERT_EQ(ArrowIpcEncoderSetCompression( - encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, 11, &error), - NANOARROW_OK) - << error.message; CompressibleRecordBatch batch; ASSERT_EQ(ArrowIpcDecoderSetSchema(decoder.get(), batch.schema(), &error), NANOARROW_OK) diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index afc16860a..63165bb8d 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -442,19 +442,27 @@ ArrowIpcSerialDecompressorSetFunction(struct ArrowIpcDecompressor* decompressor, /// \brief A user-extensible compressor /// /// The ArrowIpcCompressor is the underlying object that enables buffer compression -/// in the ArrowIpcEncoder. An implementation of a compressor may support more than one -/// ArrowIpcCompressionType. +/// in the ArrowIpcEncoder. It carries the codec and level it compresses with, so +/// that an encoder needs nothing but a compressor to describe its compression. struct ArrowIpcCompressor { + /// \brief The codec this compressor applies + /// + /// NANOARROW_IPC_COMPRESSION_TYPE_NONE disables compression. + enum ArrowIpcCompressionType compression_type; + + /// \brief The level passed to the codec + /// + /// See ArrowIpcCompressFunction for its interpretation. + int compression_level; + /// \brief Compress a buffer /// /// Compresses src using compression_type at compression_level and appends the /// compressed bytes to dst. Any content already in dst must be preserved (i.e., - /// implementations may only append to dst). See ArrowIpcCompressFunction for the - /// interpretation of compression_level. + /// implementations may only append to dst). ArrowErrorCode (*compress)(struct ArrowIpcCompressor* compressor, - enum ArrowIpcCompressionType compression_type, - int compression_level, struct ArrowBufferView src, - struct ArrowBuffer* dst, struct ArrowError* error); + struct ArrowBufferView src, struct ArrowBuffer* dst, + struct ArrowError* error); /// \brief Release the compressor and any resources it may be holding /// @@ -495,6 +503,10 @@ NANOARROW_DLL ArrowIpcCompressFunction ArrowIpcGetZstdCompressionFunction(void); NANOARROW_DLL ArrowIpcCompressFunction ArrowIpcGetLZ4CompressionFunction(void); /// \brief An ArrowIpcCompressor implementation that performs compression in serial +/// +/// The compressor is initialized with NANOARROW_IPC_COMPRESSION_TYPE_NONE and +/// NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT; set its compression_type and +/// compression_level members to configure it. NANOARROW_DLL ArrowErrorCode ArrowIpcSerialCompressor(struct ArrowIpcCompressor* compressor); @@ -508,7 +520,8 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcSerialCompressorSetFunction( /// \brief Get the name of a compression type /// -/// Returns "none", "lz4", or "zstd", or NULL for an unknown compression type. +/// Returns "none", "lz4", or "zstd", or "" for an unknown +/// compression type (never NULL, so the result is safe to use in a format string). NANOARROW_DLL const char* ArrowIpcCompressionTypeToString( enum ArrowIpcCompressionType compression_type); @@ -984,34 +997,37 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcEncoderSetMessageMetadata(struct ArrowIpcEncoder* encoder, struct ArrowBuffer* metadata, struct ArrowError* error); -/// \brief Set the buffer compression used by subsequently encoded RecordBatch messages +/// \brief Compress the bodies of subsequently encoded messages with a built-in codec +/// +/// Installs an ArrowIpcSerialCompressor() configured with compression_type and +/// compression_level, replacing any compressor previously set with this function or +/// with ArrowIpcEncoderSetCompressor(). NANOARROW_IPC_COMPRESSION_TYPE_NONE removes the +/// compressor. The setting persists until it is changed and does not affect Schema +/// messages. /// -/// When compression_type is not NANOARROW_IPC_COMPRESSION_TYPE_NONE, the body buffers -/// of every RecordBatch encoded after this call are compressed with the given codec as -/// described by the Arrow IPC format: each non-empty buffer is written as its -/// uncompressed length (a little-endian int64) followed by the compressed bytes. -/// Buffers that do not shrink when compressed are written uncompressed with a length -/// prefix of -1, and empty buffers are written as-is. The setting persists until it is -/// changed and does not affect Schema messages. +/// The body buffers of every RecordBatch or DictionaryBatch encoded while a compressor +/// is set are compressed as described by the Arrow IPC format: each non-empty buffer +/// is written as its uncompressed length (a little-endian int64) followed by the +/// compressed bytes. Buffers that do not shrink when compressed are written +/// uncompressed with a length prefix of -1, and empty buffers are written as-is. /// -/// compression_level is passed to the compressor unchanged; see -/// ArrowIpcCompressFunction for its interpretation. Use -/// NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT for the codec's default level. +/// compression_level is passed to the codec unchanged; see ArrowIpcCompressFunction +/// for its interpretation and use NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT for the +/// codec's default level. /// /// Returns EINVAL for an unknown compression type or a compression_level outside the /// range reported by ArrowIpcGetCompressionLevelRange(), and ENOTSUP if the compression /// type is not supported by this build of nanoarrow (i.e., nanoarrow was not built with -/// NANOARROW_IPC_WITH_LZ4 or NANOARROW_IPC_WITH_ZSTD). If a custom compressor was set -/// with ArrowIpcEncoderSetCompressor(), neither is checked until a RecordBatch is -/// encoded. +/// NANOARROW_IPC_WITH_LZ4 or NANOARROW_IPC_WITH_ZSTD). NANOARROW_DLL ArrowErrorCode ArrowIpcEncoderSetCompression( struct ArrowIpcEncoder* encoder, enum ArrowIpcCompressionType compression_type, int compression_level, struct ArrowError* error); -/// \brief Set the compressor implementation used by this encoder +/// \brief Compress the bodies of subsequently encoded messages with a custom compressor /// -/// The encoder takes ownership of compressor. If this is not called, an -/// ArrowIpcSerialCompressor() is used when compression is first required. +/// The encoder takes ownership of compressor and replaces any compressor previously +/// set. The compressor's compression_type and compression_level members select the +/// codec and level, and they are not checked until a message is encoded. NANOARROW_DLL ArrowErrorCode ArrowIpcEncoderSetCompressor( struct ArrowIpcEncoder* encoder, struct ArrowIpcCompressor* compressor); @@ -1112,11 +1128,12 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcWriterInit( /// \brief Release all resources attached to a writer NANOARROW_DLL void ArrowIpcWriterReset(struct ArrowIpcWriter* writer); -/// \brief Set the buffer compression used for subsequently written record batches +/// \brief Compress the bodies of subsequently written batches with a built-in codec /// -/// See ArrowIpcEncoderSetCompression(). Compression applies to record batches written -/// after this call (in both stream and file mode) and may be changed between batches. -/// Use NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT for the codec's default level. +/// See ArrowIpcEncoderSetCompression(). Compression applies to record batches and +/// dictionary batches written after this call (in both stream and file mode) and may +/// be changed between batches. Use NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT for the +/// codec's default level. /// /// Returns EINVAL for an unknown compression type or an out-of-range compression_level, /// and ENOTSUP if the compression type is not supported by this build of nanoarrow. diff --git a/src/nanoarrow/nanoarrow_ipc.hpp b/src/nanoarrow/nanoarrow_ipc.hpp index 8afdac2bd..5fbd0d3d2 100644 --- a/src/nanoarrow/nanoarrow_ipc.hpp +++ b/src/nanoarrow/nanoarrow_ipc.hpp @@ -131,6 +131,8 @@ inline void release_pointer(struct ArrowIpcDecompressor* data) { template <> inline void init_pointer(struct ArrowIpcCompressor* data) { + data->compression_type = NANOARROW_IPC_COMPRESSION_TYPE_NONE; + data->compression_level = NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT; data->private_data = nullptr; data->release = nullptr; } From 81611bcccfc3059cf1c5830398527059ae929953 Mon Sep 17 00:00:00 2001 From: Pedro Holanda Date: Wed, 9 Sep 2026 12:42:15 +0200 Subject: [PATCH 12/13] refactor(ipc): make ArrowIpcCompressor two-pass and declare COMPRESSED_BODY in schemas --- src/nanoarrow/ipc/codecs.c | 45 +++-- src/nanoarrow/ipc/codecs_test.cc | 78 +++++--- src/nanoarrow/ipc/encoder.c | 315 +++++++++++++++++++----------- src/nanoarrow/ipc/encoder_test.cc | 70 ++++++- src/nanoarrow/ipc/reader.c | 6 - src/nanoarrow/ipc/writer_test.cc | 28 +++ src/nanoarrow/nanoarrow_ipc.h | 67 ++++--- src/nanoarrow/nanoarrow_ipc.hpp | 1 - 8 files changed, 416 insertions(+), 194 deletions(-) diff --git a/src/nanoarrow/ipc/codecs.c b/src/nanoarrow/ipc/codecs.c index a2c0d21e9..5d49122ec 100644 --- a/src/nanoarrow/ipc/codecs.c +++ b/src/nanoarrow/ipc/codecs.c @@ -355,11 +355,13 @@ ArrowErrorCode ArrowIpcSerialDecompressorSetFunction( struct ArrowIpcSerialCompressorPrivate { ArrowIpcCompressFunction compress_functions[3]; + int compression_level; }; -static ArrowErrorCode ArrowIpcSerialCompressorCompress( - struct ArrowIpcCompressor* compressor, struct ArrowBufferView src, - struct ArrowBuffer* dst, struct ArrowError* error) { +static ArrowErrorCode ArrowIpcSerialCompressorAdd(struct ArrowIpcCompressor* compressor, + struct ArrowBufferView src, + struct ArrowBuffer* dst, + struct ArrowError* error) { struct ArrowIpcSerialCompressorPrivate* private_data = (struct ArrowIpcSerialCompressorPrivate*)compressor->private_data; enum ArrowIpcCompressionType compression_type = compressor->compression_type; @@ -377,7 +379,17 @@ static ArrowErrorCode ArrowIpcSerialCompressorCompress( return ENOTSUP; } - NANOARROW_RETURN_NOT_OK(fn(src, compressor->compression_level, dst, error)); + // Compression happens synchronously, so there is never anything to wait for + NANOARROW_RETURN_NOT_OK(fn(src, private_data->compression_level, dst, error)); + return NANOARROW_OK; +} + +static ArrowErrorCode ArrowIpcSerialCompressorWait(struct ArrowIpcCompressor* compressor, + int64_t timeout_ms, + struct ArrowError* error) { + NANOARROW_UNUSED(compressor); + NANOARROW_UNUSED(timeout_ms); + NANOARROW_UNUSED(error); return NANOARROW_OK; } @@ -386,22 +398,33 @@ static void ArrowIpcSerialCompressorRelease(struct ArrowIpcCompressor* compresso compressor->release = NULL; } -ArrowErrorCode ArrowIpcSerialCompressor(struct ArrowIpcCompressor* compressor) { - compressor->compression_type = NANOARROW_IPC_COMPRESSION_TYPE_NONE; - compressor->compression_level = NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT; +ArrowErrorCode ArrowIpcSerialCompressor(struct ArrowIpcCompressor* compressor, + enum ArrowIpcCompressionType compression_type, + int compression_level) { compressor->release = NULL; - compressor->private_data = ArrowMalloc(sizeof(struct ArrowIpcSerialCompressorPrivate)); - if (compressor->private_data == NULL) { + if (compression_type != NANOARROW_IPC_COMPRESSION_TYPE_NONE && + !ArrowIpcCompressionTypeIsCodec(compression_type)) { + return EINVAL; + } + + struct ArrowIpcSerialCompressorPrivate* private_data = + (struct ArrowIpcSerialCompressorPrivate*)ArrowMalloc( + sizeof(struct ArrowIpcSerialCompressorPrivate)); + if (private_data == NULL) { return ENOMEM; } - memset(compressor->private_data, 0, sizeof(struct ArrowIpcSerialCompressorPrivate)); + memset(private_data, 0, sizeof(struct ArrowIpcSerialCompressorPrivate)); + private_data->compression_level = compression_level; + compressor->compression_type = compression_type; + compressor->private_data = private_data; ArrowIpcSerialCompressorSetFunction(compressor, NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, ArrowIpcGetZstdCompressionFunction()); ArrowIpcSerialCompressorSetFunction(compressor, NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, ArrowIpcGetLZ4CompressionFunction()); - compressor->compress = &ArrowIpcSerialCompressorCompress; + compressor->compress_add = &ArrowIpcSerialCompressorAdd; + compressor->compress_wait = &ArrowIpcSerialCompressorWait; compressor->release = &ArrowIpcSerialCompressorRelease; return NANOARROW_OK; } diff --git a/src/nanoarrow/ipc/codecs_test.cc b/src/nanoarrow/ipc/codecs_test.cc index f8f19f625..b6d2c89a7 100644 --- a/src/nanoarrow/ipc/codecs_test.cc +++ b/src/nanoarrow/ipc/codecs_test.cc @@ -484,29 +484,43 @@ TEST(NanoarrowIpcTest, SerialCompressor) { struct ArrowError error {}; nanoarrow::ipc::UniqueCompressor compressor; - ASSERT_EQ(ArrowIpcSerialCompressor(compressor.get()), NANOARROW_OK); + // An invalid compression type is rejected at construction + // NOLINTNEXTLINE(clang-analyzer-optin.core.EnumCastOutOfRange) + auto unknown_type = static_cast(3); + EXPECT_EQ(ArrowIpcSerialCompressor(compressor.get(), unknown_type, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT), + EINVAL); + EXPECT_EQ(compressor->release, nullptr); + + ASSERT_EQ( + ArrowIpcSerialCompressor(compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_NONE, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT), + NANOARROW_OK); EXPECT_EQ(compressor->compression_type, NANOARROW_IPC_COMPRESSION_TYPE_NONE); - EXPECT_EQ(compressor->compression_level, NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT); // Check the function setter error ASSERT_EQ(ArrowIpcSerialCompressorSetFunction( compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_NONE, nullptr), EINVAL); + // The serial compressor never waits and always succeeds when requested to + EXPECT_EQ(compressor->compress_wait(compressor.get(), 0, &error), NANOARROW_OK); + // NONE is not a codec that can be used to compress nanoarrow::UniqueBuffer dst; - EXPECT_EQ(compressor->compress(compressor.get(), {{nullptr}, 0}, dst.get(), &error), + EXPECT_EQ(compressor->compress_add(compressor.get(), {{nullptr}, 0}, dst.get(), &error), EINVAL); EXPECT_STREQ(error.message, "Unknown compression type with value 0"); // Check a compress for a supported codec if we have one (or for an error if we don't) compressor->compression_type = NANOARROW_IPC_COMPRESSION_TYPE_ZSTD; if (ArrowIpcGetZstdCompressionFunction() != nullptr) { - ASSERT_EQ(compressor->compress(compressor.get(), - {{kUncompressed012}, sizeof(kUncompressed012)}, - dst.get(), &error), + ASSERT_EQ(compressor->compress_add(compressor.get(), + {{kUncompressed012}, sizeof(kUncompressed012)}, + dst.get(), &error), NANOARROW_OK) << error.message; + ASSERT_EQ(compressor->compress_wait(compressor.get(), -1, &error), NANOARROW_OK); ASSERT_GT(dst->size_bytes, 0); uint8_t out[sizeof(kUncompressed012)]; @@ -517,8 +531,9 @@ TEST(NanoarrowIpcTest, SerialCompressor) { << error.message; EXPECT_TRUE(std::memcmp(out, kUncompressed012, sizeof(kUncompressed012)) == 0); } else { - EXPECT_EQ(compressor->compress(compressor.get(), {{nullptr}, 0}, dst.get(), &error), - ENOTSUP); + EXPECT_EQ( + compressor->compress_add(compressor.get(), {{nullptr}, 0}, dst.get(), &error), + ENOTSUP); EXPECT_STREQ( error.message, "Compression type with value 2 not supported by this build of nanoarrow"); @@ -528,37 +543,42 @@ TEST(NanoarrowIpcTest, SerialCompressor) { ASSERT_EQ(ArrowIpcSerialCompressorSetFunction( compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, nullptr), NANOARROW_OK); - EXPECT_EQ(compressor->compress(compressor.get(), {{nullptr}, 0}, dst.get(), &error), + EXPECT_EQ(compressor->compress_add(compressor.get(), {{nullptr}, 0}, dst.get(), &error), ENOTSUP); EXPECT_STREQ(error.message, "Compression type with value 2 not supported by this build of nanoarrow"); - // The compression level is passed through to the function for the codec - ASSERT_EQ(ArrowIpcSerialCompressorSetFunction(compressor.get(), - NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, - &RecordLevelAndCopy), - NANOARROW_OK); - compressor->compression_type = NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME; - compressor->compression_level = 7; - dst->size_bytes = 0; - last_compression_level = 0; - ASSERT_EQ(compressor->compress(compressor.get(), - {{kUncompressed012}, sizeof(kUncompressed012)}, - dst.get(), &error), - NANOARROW_OK) - << error.message; - EXPECT_EQ(last_compression_level, 7); - ASSERT_EQ(dst->size_bytes, static_cast(sizeof(kUncompressed012))); - EXPECT_EQ(std::memcmp(dst->data, kUncompressed012, sizeof(kUncompressed012)), 0); + // The compression level given at construction is passed to the function for the codec + for (int level : {NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, 7}) { + nanoarrow::ipc::UniqueCompressor leveled; + ASSERT_EQ(ArrowIpcSerialCompressor(leveled.get(), + NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, level), + NANOARROW_OK); + ASSERT_EQ( + ArrowIpcSerialCompressorSetFunction( + leveled.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, &RecordLevelAndCopy), + NANOARROW_OK); + dst->size_bytes = 0; + last_compression_level = -1; + ASSERT_EQ(leveled->compress_add(leveled.get(), + {{kUncompressed012}, sizeof(kUncompressed012)}, + dst.get(), &error), + NANOARROW_OK) + << error.message; + EXPECT_EQ(last_compression_level, level); + ASSERT_EQ(dst->size_bytes, static_cast(sizeof(kUncompressed012))); + EXPECT_EQ(std::memcmp(dst->data, kUncompressed012, sizeof(kUncompressed012)), 0); + } + compressor->compression_type = NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME; // Errors from the function for the codec are propagated ASSERT_EQ( ArrowIpcSerialCompressorSetFunction( compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, &FailCompress), NANOARROW_OK); - EXPECT_EQ(compressor->compress(compressor.get(), - {{kUncompressed012}, sizeof(kUncompressed012)}, - dst.get(), &error), + EXPECT_EQ(compressor->compress_add(compressor.get(), + {{kUncompressed012}, sizeof(kUncompressed012)}, + dst.get(), &error), EIO); EXPECT_STREQ(error.message, "FailCompress() failed"); } diff --git a/src/nanoarrow/ipc/encoder.c b/src/nanoarrow/ipc/encoder.c index 6c29f3352..c6188162a 100644 --- a/src/nanoarrow/ipc/encoder.c +++ b/src/nanoarrow/ipc/encoder.c @@ -52,6 +52,11 @@ struct ArrowIpcEncoderPrivate { // Compressor for the body buffers of subsequently encoded messages (release is // NULL when they are not compressed) struct ArrowIpcCompressor compressor; + // Views of the body buffers of the message being encoded, in body order + struct ArrowBuffer buffer_views; + // Compressed copies of those buffers (reused between messages) + struct ArrowBuffer* compressed_buffers; + int64_t n_compressed_buffers; }; ArrowErrorCode ArrowIpcEncoderInit(struct ArrowIpcEncoder* encoder) { @@ -73,6 +78,9 @@ ArrowErrorCode ArrowIpcEncoderInit(struct ArrowIpcEncoder* encoder) { ArrowIpcDictionaryEncodingsInit(&private->dictionary_encodings); ArrowBufferInit(&private->message_metadata); private->compressor.release = NULL; + ArrowBufferInit(&private->buffer_views); + private->compressed_buffers = NULL; + private->n_compressed_buffers = 0; return NANOARROW_OK; } @@ -89,6 +97,11 @@ void ArrowIpcEncoderReset(struct ArrowIpcEncoder* encoder) { if (private->compressor.release != NULL) { private->compressor.release(&private->compressor); } + ArrowBufferReset(&private->buffer_views); + for (int64_t i = 0; i < private->n_compressed_buffers; i++) { + ArrowBufferReset(&private->compressed_buffers[i]); + } + ArrowFree(private->compressed_buffers); ArrowFree(private); } memset(encoder, 0, sizeof(struct ArrowIpcEncoder)); @@ -126,6 +139,15 @@ ArrowErrorCode ArrowIpcEncoderSetMessageMetadata(struct ArrowIpcEncoder* encoder return NANOARROW_OK; } +// The codec applied to message bodies (NONE when no compressor is set) +static enum ArrowIpcCompressionType ArrowIpcEncoderCodec( + struct ArrowIpcEncoderPrivate* private) { + if (private->compressor.release == NULL) { + return NANOARROW_IPC_COMPRESSION_TYPE_NONE; + } + return private->compressor.compression_type; +} + ArrowErrorCode ArrowIpcEncoderSetCompressor(struct ArrowIpcEncoder* encoder, struct ArrowIpcCompressor* compressor) { NANOARROW_DCHECK(encoder != NULL && encoder->private_data != NULL && @@ -185,9 +207,8 @@ ArrowErrorCode ArrowIpcEncoderSetCompression( } struct ArrowIpcCompressor compressor; - NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowIpcSerialCompressor(&compressor), error); - compressor.compression_type = compression_type; - compressor.compression_level = compression_level; + NANOARROW_RETURN_NOT_OK_WITH_ERROR( + ArrowIpcSerialCompressor(&compressor, compression_type, compression_level), error); return ArrowIpcEncoderSetCompressor(encoder, &compressor); } @@ -632,7 +653,7 @@ static ArrowErrorCode ArrowIpcEncodeField( static ArrowErrorCode ArrowIpcEncodeSchema( flatcc_builder_t* builder, const struct ArrowSchema* schema, - const struct ArrowIpcDictionaryEncodings* dictionary_encodings, + const struct ArrowIpcDictionaryEncodings* dictionary_encodings, int compressed_body, struct ArrowError* error) { NANOARROW_DCHECK(schema->release != NULL); @@ -665,6 +686,11 @@ static ArrowErrorCode ArrowIpcEncodeSchema( FLATCC_RETURN_UNLESS_0(Schema_custom_metadata_end(builder), error); FLATCC_RETURN_UNLESS_0(Schema_features_start(builder), error); + if (compressed_body) { + // Declare that bodies in this stream or file may be compressed + ns(Feature_enum_t) feature = ns(Feature_COMPRESSED_BODY); + FLATCC_RETURN_IF_NULL(ns(Feature_vec_push(builder, &feature)), error); + } FLATCC_RETURN_UNLESS_0(Schema_features_end(builder), error); return NANOARROW_OK; @@ -695,8 +721,9 @@ ArrowErrorCode ArrowIpcEncoderEncodeSchema(struct ArrowIpcEncoder* encoder, ArrowIpcDictionaryEncodingsAppendSchema(&private->dictionary_encodings, schema), error); - NANOARROW_RETURN_NOT_OK( - ArrowIpcEncodeSchema(builder, schema, &private->dictionary_encodings, error)); + NANOARROW_RETURN_NOT_OK(ArrowIpcEncodeSchema( + builder, schema, &private->dictionary_encodings, + ArrowIpcEncoderCodec(private) != NANOARROW_IPC_COMPRESSION_TYPE_NONE, error)); FLATCC_RETURN_UNLESS_0(Message_header_Schema_end(builder), error); @@ -731,94 +758,156 @@ struct ArrowIpcBufferEncoder { int64_t body_length; }; -// The codec applied to message bodies (NONE when no compressor is set) -static enum ArrowIpcCompressionType ArrowIpcEncoderCodec( - struct ArrowIpcEncoderPrivate* private) { - if (private->compressor.release == NULL) { - return NANOARROW_IPC_COMPRESSION_TYPE_NONE; +static ArrowErrorCode ArrowIpcEncoderBuildContiguousBodyBufferCallback( + struct ArrowBufferView buffer_view, struct ArrowIpcEncoder* encoder, + struct ArrowIpcBufferEncoder* buffer_encoder, int64_t* offset, int64_t* length, + struct ArrowError* error) { + NANOARROW_UNUSED(encoder); + + struct ArrowBuffer* body_buffer = + (struct ArrowBuffer*)buffer_encoder->encode_buffer_state; + + int64_t old_size = body_buffer->size_bytes; + int64_t buffer_begin = _ArrowRoundUpToMultipleOf8(old_size); + int64_t buffer_end = buffer_begin + buffer_view.size_bytes; + int64_t new_size = _ArrowRoundUpToMultipleOf8(buffer_end); + + // reserve all the memory we'll need now + NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowBufferReserve(body_buffer, new_size - old_size), + error); + + // zero padding up to the start of the buffer + NANOARROW_ASSERT_OK(ArrowBufferAppendFill(body_buffer, 0, buffer_begin - old_size)); + + // store offset and length of the buffer + *offset = buffer_begin; + *length = buffer_view.size_bytes; + + NANOARROW_ASSERT_OK( + ArrowBufferAppend(body_buffer, buffer_view.data.data, buffer_view.size_bytes)); + + // zero padding after writing the buffer + NANOARROW_DCHECK(body_buffer->size_bytes == buffer_end); + NANOARROW_ASSERT_OK(ArrowBufferAppendFill(body_buffer, 0, new_size - buffer_end)); + + buffer_encoder->body_length = body_buffer->size_bytes; + return NANOARROW_OK; +} + +// Make sure there is a scratch ArrowBuffer for each of n compressed buffers +static ArrowErrorCode ArrowIpcEncoderReserveCompressedBuffers( + struct ArrowIpcEncoderPrivate* private, int64_t n) { + if (n <= private->n_compressed_buffers) { + return NANOARROW_OK; } - return private->compressor.compression_type; + + struct ArrowBuffer* buffers = + (struct ArrowBuffer*)ArrowMalloc(n * sizeof(struct ArrowBuffer)); + if (buffers == NULL) { + return ENOMEM; + } + + if (private->n_compressed_buffers > 0) { + memcpy(buffers, private->compressed_buffers, + private->n_compressed_buffers * sizeof(struct ArrowBuffer)); + ArrowFree(private->compressed_buffers); + } + for (int64_t i = private->n_compressed_buffers; i < n; i++) { + ArrowBufferInit(&buffers[i]); + } + + private->compressed_buffers = buffers; + private->n_compressed_buffers = n; + return NANOARROW_OK; } -// Append buffer_view to body_buffer as a compressed IPC buffer: the uncompressed length -// as a little-endian int64 followed by the compressed bytes. If compression does not -// reduce the size, the buffer is stored uncompressed with a length prefix of -1 instead. -static ArrowErrorCode ArrowIpcEncoderAppendCompressedBuffer( - struct ArrowIpcEncoderPrivate* private, struct ArrowBufferView buffer_view, - struct ArrowBuffer* body_buffer, struct ArrowError* error) { - NANOARROW_DCHECK(ArrowIpcEncoderCodec(private) != NANOARROW_IPC_COMPRESSION_TYPE_NONE); - - // placeholder for the prefix, then compress directly into the body - int64_t prefix_offset = body_buffer->size_bytes; - int64_t payload_offset = prefix_offset + (int64_t)sizeof(int64_t); - NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowBufferAppendInt64(body_buffer, 0), error); - NANOARROW_RETURN_NOT_OK(private->compressor.compress(&private->compressor, buffer_view, - body_buffer, error)); - - int64_t prefix = buffer_view.size_bytes; - if (body_buffer->size_bytes - payload_offset >= buffer_view.size_bytes) { - body_buffer->size_bytes = payload_offset; - NANOARROW_RETURN_NOT_OK_WITH_ERROR( - ArrowBufferAppend(body_buffer, buffer_view.data.data, buffer_view.size_bytes), - error); - prefix = -1; +// Replace the collected buffer views of the message being encoded with views of their +// compressed form: the uncompressed length as a little-endian int64 followed by the +// compressed bytes. Buffers that do not shrink are stored uncompressed with a prefix of +// -1 and empty buffers are left as they are, matching Arrow C++. All buffers are queued +// with the compressor before waiting, so that a compressor may compress them in +// parallel. +static ArrowErrorCode ArrowIpcEncoderCompressBuffers( + struct ArrowIpcEncoderPrivate* private, struct ArrowError* error) { + struct ArrowBufferView* views = (struct ArrowBufferView*)private->buffer_views.data; + int64_t n_views = private->buffer_views.size_bytes / (int64_t)sizeof(*views); + NANOARROW_RETURN_NOT_OK_WITH_ERROR( + ArrowIpcEncoderReserveCompressedBuffers(private, n_views), error); + + for (int64_t i = 0; i < n_views; i++) { + if (views[i].size_bytes == 0) { + continue; + } + + // placeholder for the prefix, then the compressed bytes + struct ArrowBuffer* dst = &private->compressed_buffers[i]; + NANOARROW_ASSERT_OK(ArrowBufferResize(dst, 0, 0)); + NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowBufferAppendInt64(dst, 0), error); + int result = + private->compressor.compress_add(&private->compressor, views[i], dst, error); + if (result != NANOARROW_OK) { + // don't leave queued work referring to our buffers behind + struct ArrowError ignored; + NANOARROW_UNUSED( + private->compressor.compress_wait(&private->compressor, -1, &ignored)); + return result; + } } - // the prefix is always little endian - if (ArrowIpcSystemEndianness() == NANOARROW_IPC_ENDIANNESS_BIG) { - prefix = (int64_t)bswap64((uint64_t)prefix); + NANOARROW_RETURN_NOT_OK( + private->compressor.compress_wait(&private->compressor, -1, error)); + + for (int64_t i = 0; i < n_views; i++) { + if (views[i].size_bytes == 0) { + continue; + } + + // if compression did not reduce the size, store the buffer uncompressed instead + // (signalled to the reader by a prefix of -1) + struct ArrowBuffer* dst = &private->compressed_buffers[i]; + int64_t prefix = views[i].size_bytes; + if (dst->size_bytes - (int64_t)sizeof(int64_t) >= views[i].size_bytes) { + dst->size_bytes = sizeof(int64_t); + NANOARROW_RETURN_NOT_OK_WITH_ERROR( + ArrowBufferAppend(dst, views[i].data.data, views[i].size_bytes), error); + prefix = -1; + } + + // the prefix is always little endian + if (ArrowIpcSystemEndianness() == NANOARROW_IPC_ENDIANNESS_BIG) { + prefix = (int64_t)bswap64((uint64_t)prefix); + } + memcpy(dst->data, &prefix, sizeof(int64_t)); + + views[i].data.data = dst->data; + views[i].size_bytes = dst->size_bytes; } - memcpy(body_buffer->data + prefix_offset, &prefix, sizeof(int64_t)); + return NANOARROW_OK; } -static ArrowErrorCode ArrowIpcEncoderBuildContiguousBodyBufferCallback( - struct ArrowBufferView buffer_view, struct ArrowIpcEncoder* encoder, - struct ArrowIpcBufferEncoder* buffer_encoder, int64_t* offset, int64_t* length, +// Encode the collected buffer views of the message being encoded (compressed first if +// a compressor is set) with the buffer encoder, recording their offsets and lengths +static ArrowErrorCode ArrowIpcEncoderEncodeBuffers( + struct ArrowIpcEncoder* encoder, struct ArrowIpcBufferEncoder* buffer_encoder, struct ArrowError* error) { struct ArrowIpcEncoderPrivate* private = (struct ArrowIpcEncoderPrivate*)encoder->private_data; - struct ArrowBuffer* body_buffer = - (struct ArrowBuffer*)buffer_encoder->encode_buffer_state; - int64_t buffer_begin = _ArrowRoundUpToMultipleOf8(body_buffer->size_bytes); - // Empty buffers are never compressed (nor length-prefixed), matching Arrow C++. - int needs_compression = - ArrowIpcEncoderCodec(private) != NANOARROW_IPC_COMPRESSION_TYPE_NONE && - buffer_view.size_bytes > 0; - if (!needs_compression) { - // Reserve the data and padding together to avoid growing the buffer twice. - int64_t new_size = _ArrowRoundUpToMultipleOf8(buffer_begin + buffer_view.size_bytes); - NANOARROW_RETURN_NOT_OK_WITH_ERROR( - ArrowBufferReserve(body_buffer, new_size - body_buffer->size_bytes), error); + if (ArrowIpcEncoderCodec(private) != NANOARROW_IPC_COMPRESSION_TYPE_NONE) { + NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderCompressBuffers(private, error)); } - // zero padding up to the start of the buffer - NANOARROW_RETURN_NOT_OK_WITH_ERROR( - ArrowBufferAppendFill(body_buffer, 0, buffer_begin - body_buffer->size_bytes), - error); - - if (needs_compression) { - NANOARROW_RETURN_NOT_OK( - ArrowIpcEncoderAppendCompressedBuffer(private, buffer_view, body_buffer, error)); - } else { + struct ArrowBufferView* views = (struct ArrowBufferView*)private->buffer_views.data; + int64_t n_views = private->buffer_views.size_bytes / (int64_t)sizeof(*views); + for (int64_t i = 0; i < n_views; i++) { + struct ns(Buffer) buffer; + NANOARROW_RETURN_NOT_OK(buffer_encoder->encode_buffer( + views[i], encoder, buffer_encoder, &buffer.offset, &buffer.length, error)); NANOARROW_RETURN_NOT_OK_WITH_ERROR( - ArrowBufferAppend(body_buffer, buffer_view.data.data, buffer_view.size_bytes), - error); + ArrowBufferAppend(&private->buffers, &buffer, sizeof(buffer)), error); } - // store offset and length (including any prefix) of the buffer - *offset = buffer_begin; - *length = body_buffer->size_bytes - buffer_begin; - - // zero padding after writing the buffer - int64_t buffer_end = body_buffer->size_bytes; - NANOARROW_RETURN_NOT_OK_WITH_ERROR( - ArrowBufferAppendFill(body_buffer, 0, - _ArrowRoundUpToMultipleOf8(buffer_end) - buffer_end), - error); - - buffer_encoder->body_length = body_buffer->size_bytes; return NANOARROW_OK; } @@ -852,10 +941,25 @@ static ArrowErrorCode ArrowIpcEncoderEncodeBodyCompression( return NANOARROW_OK; } +// Collect the node and buffer views of an array to be encoded +static ArrowErrorCode ArrowIpcEncoderCollectArray(struct ArrowIpcEncoderPrivate* private, + const struct ArrowArrayView* array_view, + struct ArrowError* error) { + struct ns(FieldNode) node = {array_view->length, array_view->null_count}; + NANOARROW_RETURN_NOT_OK_WITH_ERROR( + ArrowBufferAppend(&private->nodes, &node, sizeof(node)), error); + for (int64_t b = 0; b < array_view->array->n_buffers; ++b) { + NANOARROW_RETURN_NOT_OK_WITH_ERROR( + ArrowBufferAppend(&private->buffer_views, &array_view->buffer_views[b], + sizeof(struct ArrowBufferView)), + error); + } + return NANOARROW_OK; +} + static ArrowErrorCode ArrowIpcEncoderEncodeRecordBatchImpl( - struct ArrowIpcEncoder* encoder, struct ArrowIpcBufferEncoder* buffer_encoder, - const struct ArrowArrayView* array_view, struct ArrowBuffer* buffers, - struct ArrowBuffer* nodes, struct ArrowError* error) { + struct ArrowIpcEncoderPrivate* private, const struct ArrowArrayView* array_view, + struct ArrowError* error) { if (array_view->offset != 0) { ArrowErrorSet(error, "Cannot encode arrays with nonzero offset"); return ENOTSUP; @@ -863,29 +967,15 @@ static ArrowErrorCode ArrowIpcEncoderEncodeRecordBatchImpl( if (array_view->dictionary != NULL) { // Values live in a separate DictionaryBatch message per the Arrow IPC spec; - // the parent's index node + buffers were already emitted by the caller loop, + // the parent's index node + buffers were already collected by the caller loop, // so stop recursing here. return NANOARROW_OK; } for (int64_t c = 0; c < array_view->n_children; ++c) { const struct ArrowArrayView* child = array_view->children[c]; - - struct ns(FieldNode) node = {child->length, child->null_count}; - NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowBufferAppend(nodes, &node, sizeof(node)), - error); - - for (int64_t b = 0; b < child->array->n_buffers; ++b) { - struct ns(Buffer) buffer; - NANOARROW_RETURN_NOT_OK( - buffer_encoder->encode_buffer(child->buffer_views[b], encoder, buffer_encoder, - &buffer.offset, &buffer.length, error)); - NANOARROW_RETURN_NOT_OK_WITH_ERROR( - ArrowBufferAppend(buffers, &buffer, sizeof(buffer)), error); - } - - NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderEncodeRecordBatchImpl( - encoder, buffer_encoder, child, buffers, nodes, error)); + NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderCollectArray(private, child, error)); + NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderEncodeRecordBatchImpl(private, child, error)); } return NANOARROW_OK; } @@ -923,8 +1013,10 @@ static ArrowErrorCode ArrowIpcEncoderEncodeRecordBatch( NANOARROW_ASSERT_OK(ArrowBufferResize(&private->buffers, 0, 0)); NANOARROW_ASSERT_OK(ArrowBufferResize(&private->nodes, 0, 0)); - NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderEncodeRecordBatchImpl( - encoder, buffer_encoder, array_view, &private->buffers, &private->nodes, error)); + NANOARROW_ASSERT_OK(ArrowBufferResize(&private->buffer_views, 0, 0)); + NANOARROW_RETURN_NOT_OK( + ArrowIpcEncoderEncodeRecordBatchImpl(private, array_view, error)); + NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderEncodeBuffers(encoder, buffer_encoder, error)); FLATCC_RETURN_UNLESS_0(RecordBatch_nodes_create( // builder, (struct ns(FieldNode)*)private->nodes.data, @@ -987,22 +1079,14 @@ static ArrowErrorCode ArrowIpcEncoderEncodeDictionaryBatch( NANOARROW_ASSERT_OK(ArrowBufferResize(&private->buffers, 0, 0)); NANOARROW_ASSERT_OK(ArrowBufferResize(&private->nodes, 0, 0)); + NANOARROW_ASSERT_OK(ArrowBufferResize(&private->buffer_views, 0, 0)); - // The values array is a single top-level column. Emit the top-level node + - // buffers here, then descend into any nested children. - struct ns(FieldNode) top_node = {values_view->length, values_view->null_count}; - NANOARROW_RETURN_NOT_OK_WITH_ERROR( - ArrowBufferAppend(&private->nodes, &top_node, sizeof(top_node)), error); - for (int64_t b = 0; b < values_view->array->n_buffers; ++b) { - struct ns(Buffer) buffer; - NANOARROW_RETURN_NOT_OK(buffer_encoder->encode_buffer( - values_view->buffer_views[b], encoder, buffer_encoder, &buffer.offset, - &buffer.length, error)); - NANOARROW_RETURN_NOT_OK_WITH_ERROR( - ArrowBufferAppend(&private->buffers, &buffer, sizeof(buffer)), error); - } - NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderEncodeRecordBatchImpl( - encoder, buffer_encoder, values_view, &private->buffers, &private->nodes, error)); + // The values array is a single top-level column: collect it, then descend into any + // nested children. + NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderCollectArray(private, values_view, error)); + NANOARROW_RETURN_NOT_OK( + ArrowIpcEncoderEncodeRecordBatchImpl(private, values_view, error)); + NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderEncodeBuffers(encoder, buffer_encoder, error)); FLATCC_RETURN_UNLESS_0( RecordBatch_nodes_create(builder, (struct ns(FieldNode)*)private->nodes.data, @@ -1067,8 +1151,9 @@ ArrowErrorCode ArrowIpcEncoderEncodeFooter(struct ArrowIpcEncoder* encoder, FLATCC_RETURN_UNLESS_0(Footer_version_add(builder, ns(MetadataVersion_V5)), error); FLATCC_RETURN_UNLESS_0(Footer_schema_start(builder), error); - NANOARROW_RETURN_NOT_OK( - ArrowIpcEncodeSchema(builder, &footer->schema, &footer->dictionaries, error)); + NANOARROW_RETURN_NOT_OK(ArrowIpcEncodeSchema( + builder, &footer->schema, &footer->dictionaries, + ArrowIpcEncoderCodec(private) != NANOARROW_IPC_COMPRESSION_TYPE_NONE, error)); FLATCC_RETURN_UNLESS_0(Footer_schema_end(builder), error); const struct ArrowIpcFileBlock* blocks = diff --git a/src/nanoarrow/ipc/encoder_test.cc b/src/nanoarrow/ipc/encoder_test.cc index 80efb9981..886dff542 100644 --- a/src/nanoarrow/ipc/encoder_test.cc +++ b/src/nanoarrow/ipc/encoder_test.cc @@ -856,7 +856,10 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSetCompressor) { // A compressor whose release we can observe nanoarrow::ipc::UniqueCompressor first_compressor; - ASSERT_EQ(ArrowIpcSerialCompressor(first_compressor.get()), NANOARROW_OK); + ASSERT_EQ(ArrowIpcSerialCompressor(first_compressor.get(), + NANOARROW_IPC_COMPRESSION_TYPE_NONE, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT), + NANOARROW_OK); original_compressor_release = first_compressor->release; first_compressor->release = &CountingCompressorRelease; compressor_release_calls = 0; @@ -867,11 +870,13 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSetCompressor) { // A custom compressor configured for LZ4 that explicitly does not support it nanoarrow::ipc::UniqueCompressor compressor; - ASSERT_EQ(ArrowIpcSerialCompressor(compressor.get()), NANOARROW_OK); + ASSERT_EQ( + ArrowIpcSerialCompressor(compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT), + NANOARROW_OK); ASSERT_EQ(ArrowIpcSerialCompressorSetFunction( compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, nullptr), NANOARROW_OK); - compressor->compression_type = NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME; ASSERT_EQ(ArrowIpcEncoderSetCompressor(encoder.get(), compressor.get()), NANOARROW_OK); // The encoder took ownership of the compressor and released the previous one @@ -920,13 +925,13 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderCompressionLevel) { struct ArrowError error; nanoarrow::ipc::UniqueCompressor compressor; - ASSERT_EQ(ArrowIpcSerialCompressor(compressor.get()), NANOARROW_OK); + ASSERT_EQ( + ArrowIpcSerialCompressor(compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, 11), + NANOARROW_OK); ASSERT_EQ( ArrowIpcSerialCompressorSetFunction( compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, &RecordLevelAndCopy), NANOARROW_OK); - compressor->compression_type = NANOARROW_IPC_COMPRESSION_TYPE_ZSTD; - compressor->compression_level = 11; ASSERT_EQ(ArrowIpcEncoderSetCompressor(encoder.get(), compressor.get()), NANOARROW_OK); CompressibleRecordBatch batch; @@ -1192,3 +1197,56 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderCompressedDictionaryBatchZstd) { } TestCompressedDictionaryBatch(NANOARROW_IPC_COMPRESSION_TYPE_ZSTD); } + +// Schemas encoded while a compressor is set declare the COMPRESSED_BODY feature +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSchemaDeclaresCompression) { + nanoarrow::ipc::UniqueEncoder encoder; + ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); + nanoarrow::ipc::UniqueDecoder decoder; + ASSERT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK); + SimpleRecordBatch batch; + struct ArrowError error; + + auto encode_and_decode_schema = [&](nanoarrow::UniqueBuffer& message) { + message->size_bytes = 0; + ASSERT_EQ(ArrowIpcEncoderEncodeSchema(encoder.get(), batch.schema(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ( + ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/true, message.get()), + NANOARROW_OK); + struct ArrowBufferView message_view = {{message->data}, message->size_bytes}; + ASSERT_EQ(ArrowIpcDecoderVerifyHeader(decoder.get(), message_view, &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcDecoderDecodeHeader(decoder.get(), message_view, &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(decoder->message_type, NANOARROW_IPC_MESSAGE_TYPE_SCHEMA); + }; + + // Without a compressor no feature is declared + nanoarrow::UniqueBuffer message; + ASSERT_NO_FATAL_FAILURE(encode_and_decode_schema(message)); + EXPECT_EQ(decoder->feature_flags & NANOARROW_IPC_FEATURE_COMPRESSED_BODY, 0); + + // Any compressor (encoding a schema never runs it) declares the feature + nanoarrow::ipc::UniqueCompressor compressor; + ASSERT_EQ( + ArrowIpcSerialCompressor(compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT), + NANOARROW_OK); + ASSERT_EQ(ArrowIpcEncoderSetCompressor(encoder.get(), compressor.get()), NANOARROW_OK); + ASSERT_NO_FATAL_FAILURE(encode_and_decode_schema(message)); + EXPECT_EQ(decoder->feature_flags & NANOARROW_IPC_FEATURE_COMPRESSED_BODY, + NANOARROW_IPC_FEATURE_COMPRESSED_BODY); + + // Removing the compressor removes the declaration again + ASSERT_EQ( + ArrowIpcEncoderSetCompression(encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_NONE, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), + NANOARROW_OK) + << error.message; + ASSERT_NO_FATAL_FAILURE(encode_and_decode_schema(message)); + EXPECT_EQ(decoder->feature_flags & NANOARROW_IPC_FEATURE_COMPRESSED_BODY, 0); +} diff --git a/src/nanoarrow/ipc/reader.c b/src/nanoarrow/ipc/reader.c index 0b99b828e..c699ca331 100644 --- a/src/nanoarrow/ipc/reader.c +++ b/src/nanoarrow/ipc/reader.c @@ -390,12 +390,6 @@ static int ArrowIpcArrayStreamReaderReadSchemaIfNeeded( } // ...or if it uses features we don't support - if (private_data->decoder.feature_flags & NANOARROW_IPC_FEATURE_COMPRESSED_BODY) { - ArrowErrorSet(&private_data->error, - "This stream uses unsupported feature COMPRESSED_BODY"); - return EINVAL; - } - if (private_data->decoder.feature_flags & NANOARROW_IPC_FEATURE_DICTIONARY_REPLACEMENT) { ArrowErrorSet(&private_data->error, diff --git a/src/nanoarrow/ipc/writer_test.cc b/src/nanoarrow/ipc/writer_test.cc index adcc0615f..5f12d27c2 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -515,6 +515,30 @@ static void CheckCompressibleBatch(const struct ArrowBuffer* output, int64_t off EXPECT_EQ(eos->release, nullptr); } +// Check whether the schema message (of a stream) or the footer (of a file) declares +// the COMPRESSED_BODY feature +static void CheckDeclaresCompression(const struct ArrowBuffer* output, bool as_file, + bool expected) { + struct ArrowError error; + nanoarrow::ipc::UniqueDecoder decoder; + ASSERT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK); + struct ArrowBufferView view = {{output->data}, output->size_bytes}; + if (as_file) { + ASSERT_EQ(ArrowIpcDecoderVerifyFooter(decoder.get(), view, &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcDecoderDecodeFooter(decoder.get(), view, &error), NANOARROW_OK) + << error.message; + } else { + ASSERT_EQ(ArrowIpcDecoderVerifyHeader(decoder.get(), view, &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcDecoderDecodeHeader(decoder.get(), view, &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(decoder->message_type, NANOARROW_IPC_MESSAGE_TYPE_SCHEMA); + } + EXPECT_EQ((decoder->feature_flags & NANOARROW_IPC_FEATURE_COMPRESSED_BODY) != 0, + expected); +} + static void TestCompressedWriting(enum ArrowIpcCompressionType codec) { for (bool as_file : {false, true}) { SCOPED_TRACE(as_file ? "file" : "stream"); @@ -535,6 +559,10 @@ static void TestCompressedWriting(enum ArrowIpcCompressionType codec) { int64_t offset = as_file ? sizeof(NANOARROW_IPC_FILE_PADDED_MAGIC) : 0; ASSERT_NO_FATAL_FAILURE(CheckCompressibleBatch(compressed.get(), offset)); ASSERT_NO_FATAL_FAILURE(CheckCompressibleBatch(accelerated.get(), offset)); + + // The schema message (or the file footer) declares that bodies are compressed + ASSERT_NO_FATAL_FAILURE(CheckDeclaresCompression(compressed.get(), as_file, true)); + ASSERT_NO_FATAL_FAILURE(CheckDeclaresCompression(uncompressed.get(), as_file, false)); } } diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index 63165bb8d..6ab7ae483 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -442,33 +442,42 @@ ArrowIpcSerialDecompressorSetFunction(struct ArrowIpcDecompressor* decompressor, /// \brief A user-extensible compressor /// /// The ArrowIpcCompressor is the underlying object that enables buffer compression -/// in the ArrowIpcEncoder. It carries the codec and level it compresses with, so -/// that an encoder needs nothing but a compressor to describe its compression. +/// in the ArrowIpcEncoder. Its structure allows it to be backed by a multithreaded +/// implementation; however, this is not required and the default implementation does +/// not implement this. The encoder queues every buffer of a message with compress_add() +/// and then calls compress_wait() before it assembles the message body. struct ArrowIpcCompressor { /// \brief The codec this compressor applies /// - /// NANOARROW_IPC_COMPRESSION_TYPE_NONE disables compression. + /// The encoder declares this codec in the messages it encodes, so implementations + /// must compress with it. NANOARROW_IPC_COMPRESSION_TYPE_NONE disables compression. + /// Any other parameters (e.g., a compression level) are part of the implementation. enum ArrowIpcCompressionType compression_type; - /// \brief The level passed to the codec + /// \brief Queue a buffer for compression /// - /// See ArrowIpcCompressFunction for its interpretation. - int compression_level; - - /// \brief Compress a buffer + /// Compresses src using compression_type and appends the compressed bytes to dst, + /// preserving any content already in dst. The content of dst is undefined until the + /// next call to compress_wait() returns NANOARROW_OK, and the caller must not use + /// src or dst until then. + ArrowErrorCode (*compress_add)(struct ArrowIpcCompressor* compressor, + struct ArrowBufferView src, struct ArrowBuffer* dst, + struct ArrowError* error); + + /// \brief Wait for any unfinished calls to compress_add to complete /// - /// Compresses src using compression_type at compression_level and appends the - /// compressed bytes to dst. Any content already in dst must be preserved (i.e., - /// implementations may only append to dst). - ArrowErrorCode (*compress)(struct ArrowIpcCompressor* compressor, - struct ArrowBufferView src, struct ArrowBuffer* dst, - struct ArrowError* error); + /// Returns NANOARROW_OK if all pending calls completed. Returns ETIMEDOUT if not all + /// remaining calls completed within timeout_ms (a negative timeout waits + /// indefinitely). + ArrowErrorCode (*compress_wait)(struct ArrowIpcCompressor* compressor, + int64_t timeout_ms, struct ArrowError* error); /// \brief Release the compressor and any resources it may be holding /// + /// Implementations must wait for or cancel any queued work before returning. /// Release callback implementations must set the release member to NULL. /// Callers must check that the release callback is not NULL before calling - /// compress() or release(). + /// compress_add(), compress_wait(), or release(). void (*release)(struct ArrowIpcCompressor* compressor); /// \brief Implementation-specific opaque data @@ -504,11 +513,13 @@ NANOARROW_DLL ArrowIpcCompressFunction ArrowIpcGetLZ4CompressionFunction(void); /// \brief An ArrowIpcCompressor implementation that performs compression in serial /// -/// The compressor is initialized with NANOARROW_IPC_COMPRESSION_TYPE_NONE and -/// NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT; set its compression_type and -/// compression_level members to configure it. -NANOARROW_DLL ArrowErrorCode -ArrowIpcSerialCompressor(struct ArrowIpcCompressor* compressor); +/// The compressor compresses with compression_type at compression_level (see +/// ArrowIpcCompressFunction for the interpretation of the level). Returns EINVAL if +/// compression_type is not a valid compression type; whether it is supported by this +/// build of nanoarrow is not checked until a buffer is compressed. +NANOARROW_DLL ArrowErrorCode ArrowIpcSerialCompressor( + struct ArrowIpcCompressor* compressor, enum ArrowIpcCompressionType compression_type, + int compression_level); /// \brief Override the ArrowIpcCompressFunction used for a specific compression type /// @@ -1002,8 +1013,9 @@ ArrowIpcEncoderSetMessageMetadata(struct ArrowIpcEncoder* encoder, /// Installs an ArrowIpcSerialCompressor() configured with compression_type and /// compression_level, replacing any compressor previously set with this function or /// with ArrowIpcEncoderSetCompressor(). NANOARROW_IPC_COMPRESSION_TYPE_NONE removes the -/// compressor. The setting persists until it is changed and does not affect Schema -/// messages. +/// compressor. The setting persists until it is changed. Schema messages (and file +/// footers) encoded while a compressor is set declare the COMPRESSED_BODY feature, so +/// compression should be set before the schema is encoded. /// /// The body buffers of every RecordBatch or DictionaryBatch encoded while a compressor /// is set are compressed as described by the Arrow IPC format: each non-empty buffer @@ -1026,8 +1038,10 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcEncoderSetCompression( /// \brief Compress the bodies of subsequently encoded messages with a custom compressor /// /// The encoder takes ownership of compressor and replaces any compressor previously -/// set. The compressor's compression_type and compression_level members select the -/// codec and level, and they are not checked until a message is encoded. +/// set. The compressor's compression_type member selects the codec (any level is +/// part of the compressor's own configuration, like the one given to +/// ArrowIpcSerialCompressor()), and its support is not checked until a message is +/// encoded. NANOARROW_DLL ArrowErrorCode ArrowIpcEncoderSetCompressor( struct ArrowIpcEncoder* encoder, struct ArrowIpcCompressor* compressor); @@ -1132,8 +1146,9 @@ NANOARROW_DLL void ArrowIpcWriterReset(struct ArrowIpcWriter* writer); /// /// See ArrowIpcEncoderSetCompression(). Compression applies to record batches and /// dictionary batches written after this call (in both stream and file mode) and may -/// be changed between batches. Use NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT for the -/// codec's default level. +/// be changed between batches. Set it before writing the schema so that the stream +/// (or the file footer) declares the COMPRESSED_BODY feature. Use +/// NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT for the codec's default level. /// /// Returns EINVAL for an unknown compression type or an out-of-range compression_level, /// and ENOTSUP if the compression type is not supported by this build of nanoarrow. diff --git a/src/nanoarrow/nanoarrow_ipc.hpp b/src/nanoarrow/nanoarrow_ipc.hpp index 5fbd0d3d2..9629eba31 100644 --- a/src/nanoarrow/nanoarrow_ipc.hpp +++ b/src/nanoarrow/nanoarrow_ipc.hpp @@ -132,7 +132,6 @@ inline void release_pointer(struct ArrowIpcDecompressor* data) { template <> inline void init_pointer(struct ArrowIpcCompressor* data) { data->compression_type = NANOARROW_IPC_COMPRESSION_TYPE_NONE; - data->compression_level = NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT; data->private_data = nullptr; data->release = nullptr; } From fd1023e76d7ae652655fa5d198dc3accafde9e6a Mon Sep 17 00:00:00 2001 From: Pedro Holanda Date: Wed, 9 Sep 2026 14:00:38 +0200 Subject: [PATCH 13/13] fix(ipc): declare COMPRESSED_BODY in footers after compressed bodies and harden the encoder --- src/nanoarrow/ipc/codecs.c | 6 +- src/nanoarrow/ipc/encoder.c | 35 +++- src/nanoarrow/ipc/encoder_test.cc | 311 ++++++++++++++++++++++++++++++ src/nanoarrow/ipc/writer_test.cc | 53 +++++ src/nanoarrow/nanoarrow_ipc.h | 28 +-- 5 files changed, 414 insertions(+), 19 deletions(-) diff --git a/src/nanoarrow/ipc/codecs.c b/src/nanoarrow/ipc/codecs.c index 5d49122ec..e5cd0b312 100644 --- a/src/nanoarrow/ipc/codecs.c +++ b/src/nanoarrow/ipc/codecs.c @@ -149,9 +149,8 @@ static ArrowErrorCode ArrowIpcCompressLZ4(struct ArrowBufferView src, } // Default preferences except for the compression level (no content size, no - // checksums). - // This produces a single complete frame, which is what ArrowIpcDecompressLZ4() - // and Arrow C++ expect. + // checksums). This produces a single complete frame, which is what + // ArrowIpcDecompressLZ4() and Arrow C++ expect. LZ4F_preferences_t prefs; memset(&prefs, 0, sizeof(prefs)); prefs.compressionLevel = compression_level; @@ -195,6 +194,7 @@ ArrowErrorCode ArrowIpcGetCompressionLevelRange( enum ArrowIpcCompressionType compression_type, int* min_level_out, int* max_level_out) { NANOARROW_DCHECK(min_level_out != NULL && max_level_out != NULL); + // (unused when neither codec is built in) NANOARROW_UNUSED(min_level_out); NANOARROW_UNUSED(max_level_out); diff --git a/src/nanoarrow/ipc/encoder.c b/src/nanoarrow/ipc/encoder.c index c6188162a..ab7f0b60d 100644 --- a/src/nanoarrow/ipc/encoder.c +++ b/src/nanoarrow/ipc/encoder.c @@ -16,6 +16,7 @@ // under the License. #include +#include #include #include @@ -52,6 +53,8 @@ struct ArrowIpcEncoderPrivate { // Compressor for the body buffers of subsequently encoded messages (release is // NULL when they are not compressed) struct ArrowIpcCompressor compressor; + // Whether compression was declared or encoded since the last Schema message + int has_compressed_body; // Views of the body buffers of the message being encoded, in body order struct ArrowBuffer buffer_views; // Compressed copies of those buffers (reused between messages) @@ -78,6 +81,7 @@ ArrowErrorCode ArrowIpcEncoderInit(struct ArrowIpcEncoder* encoder) { ArrowIpcDictionaryEncodingsInit(&private->dictionary_encodings); ArrowBufferInit(&private->message_metadata); private->compressor.release = NULL; + private->has_compressed_body = 0; ArrowBufferInit(&private->buffer_views); private->compressed_buffers = NULL; private->n_compressed_buffers = 0; @@ -732,6 +736,8 @@ ArrowErrorCode ArrowIpcEncoderEncodeSchema(struct ArrowIpcEncoder* encoder, FLATCC_RETURN_UNLESS_0(Message_bodyLength_add(builder, 0), error); FLATCC_RETURN_IF_NULL(ns(Message_end_as_root(builder)), error); + private->has_compressed_body = + ArrowIpcEncoderCodec(private) != NANOARROW_IPC_COMPRESSION_TYPE_NONE; return NANOARROW_OK; } @@ -834,6 +840,8 @@ static ArrowErrorCode ArrowIpcEncoderCompressBuffers( NANOARROW_RETURN_NOT_OK_WITH_ERROR( ArrowIpcEncoderReserveCompressedBuffers(private, n_views), error); + // Allocate every prefix before queueing work so allocation failures cannot leave + // jobs referring to the caller's source buffers or our scratch buffers. for (int64_t i = 0; i < n_views; i++) { if (views[i].size_bytes == 0) { continue; @@ -843,8 +851,15 @@ static ArrowErrorCode ArrowIpcEncoderCompressBuffers( struct ArrowBuffer* dst = &private->compressed_buffers[i]; NANOARROW_ASSERT_OK(ArrowBufferResize(dst, 0, 0)); NANOARROW_RETURN_NOT_OK_WITH_ERROR(ArrowBufferAppendInt64(dst, 0), error); - int result = - private->compressor.compress_add(&private->compressor, views[i], dst, error); + } + + for (int64_t i = 0; i < n_views; i++) { + if (views[i].size_bytes == 0) { + continue; + } + + int result = private->compressor.compress_add(&private->compressor, views[i], + &private->compressed_buffers[i], error); if (result != NANOARROW_OK) { // don't leave queued work referring to our buffers behind struct ArrowError ignored; @@ -862,9 +877,18 @@ static ArrowErrorCode ArrowIpcEncoderCompressBuffers( continue; } + // a compressor that produced nothing violated its contract; without this check the + // buffer would be written with a length prefix and no payload + struct ArrowBuffer* dst = &private->compressed_buffers[i]; + if (dst->size_bytes == (int64_t)sizeof(int64_t)) { + ArrowErrorSet(error, + "Compressor produced no output for a buffer of %" PRId64 " bytes", + views[i].size_bytes); + return EIO; + } + // if compression did not reduce the size, store the buffer uncompressed instead // (signalled to the reader by a prefix of -1) - struct ArrowBuffer* dst = &private->compressed_buffers[i]; int64_t prefix = views[i].size_bytes; if (dst->size_bytes - (int64_t)sizeof(int64_t) >= views[i].size_bytes) { dst->size_bytes = sizeof(int64_t); @@ -938,6 +962,7 @@ static ArrowErrorCode ArrowIpcEncoderEncodeBodyCompression( FLATCC_RETURN_UNLESS_0( BodyCompression_method_add(builder, ns(BodyCompressionMethod_BUFFER)), error); FLATCC_RETURN_UNLESS_0(RecordBatch_compression_end(builder), error); + private->has_compressed_body = 1; return NANOARROW_OK; } @@ -1153,7 +1178,9 @@ ArrowErrorCode ArrowIpcEncoderEncodeFooter(struct ArrowIpcEncoder* encoder, FLATCC_RETURN_UNLESS_0(Footer_schema_start(builder), error); NANOARROW_RETURN_NOT_OK(ArrowIpcEncodeSchema( builder, &footer->schema, &footer->dictionaries, - ArrowIpcEncoderCodec(private) != NANOARROW_IPC_COMPRESSION_TYPE_NONE, error)); + private->has_compressed_body || + ArrowIpcEncoderCodec(private) != NANOARROW_IPC_COMPRESSION_TYPE_NONE, + error)); FLATCC_RETURN_UNLESS_0(Footer_schema_end(builder), error); const struct ArrowIpcFileBlock* blocks = diff --git a/src/nanoarrow/ipc/encoder_test.cc b/src/nanoarrow/ipc/encoder_test.cc index 886dff542..6c79babd7 100644 --- a/src/nanoarrow/ipc/encoder_test.cc +++ b/src/nanoarrow/ipc/encoder_test.cc @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +#include #include #include @@ -1016,6 +1017,161 @@ static struct ArrowBufferAllocator FailingAllocator(FailingAllocatorState* state return allocator; } +// Defer all work until Wait(), as a compressor backed by a thread pool could do. +// Copying forces the uncompressed fallback and works without either codec built in. +struct DeferredCompressor { + struct Job { + struct ArrowBufferView src; + struct ArrowBuffer* dst; + }; + std::vector pending; + std::vector destinations; + size_t max_pending = 0; + int adds = 0; + int waits = 0; + int fail_on_add = 0; + bool fail_wait = false; + bool produce_nothing = false; + + static ArrowErrorCode Add(struct ArrowIpcCompressor* compressor, + struct ArrowBufferView src, struct ArrowBuffer* dst, + struct ArrowError* error) { + auto* state = static_cast(compressor->private_data); + if (++state->adds == state->fail_on_add) { + ArrowErrorSet(error, "Deferred add failed"); + return EIO; + } + state->pending.push_back({src, dst}); + state->destinations.push_back(dst); + if (state->pending.size() > state->max_pending) { + state->max_pending = state->pending.size(); + } + return NANOARROW_OK; + } + + static ArrowErrorCode Wait(struct ArrowIpcCompressor* compressor, int64_t timeout_ms, + struct ArrowError* error) { + EXPECT_LT(timeout_ms, 0); + auto* state = static_cast(compressor->private_data); + ++state->waits; + int result = NANOARROW_OK; + for (const auto& job : state->pending) { + if (result == NANOARROW_OK && !state->produce_nothing) { + result = ArrowBufferAppend(job.dst, job.src.data.data, job.src.size_bytes); + } + } + // Complete or cancel every job, including when reporting an error. + state->pending.clear(); + if (state->fail_wait) { + ArrowErrorSet(error, "Deferred wait failed"); + return EIO; + } + return result; + } + + static void Release(struct ArrowIpcCompressor* compressor) { + auto* state = static_cast(compressor->private_data); + state->pending.clear(); + compressor->release = nullptr; + } + + struct ArrowIpcCompressor MakeCompressor() { + struct ArrowIpcCompressor compressor {}; + compressor.compression_type = NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME; + compressor.compress_add = &Add; + compressor.compress_wait = &Wait; + compressor.release = &Release; + compressor.private_data = this; + return compressor; + } +}; + +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderDeferredCompressionAllocationFailure) { + struct ArrowError error; + CompressibleRecordBatch batch; + DeferredCompressor state; + FailingAllocatorState allocator_state{0, 1}; + nanoarrow::ipc::UniqueEncoder encoder; + ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); + auto compressor = state.MakeCompressor(); + ASSERT_EQ(ArrowIpcEncoderSetCompressor(encoder.get(), &compressor), NANOARROW_OK); + + nanoarrow::UniqueBuffer body, message; + ASSERT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), batch.array_view(), + body.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcEncoderFinalizeBuffer(encoder.get(), true, message.get()), + NANOARROW_OK); + EXPECT_TRUE(state.pending.empty()); + EXPECT_GT(state.max_pending, 1); + EXPECT_EQ(state.waits, 1); + + nanoarrow::ipc::UniqueDecoder decoder; + ASSERT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK); + ASSERT_EQ(ArrowIpcDecoderSetSchema(decoder.get(), batch.schema(), &error), + NANOARROW_OK); + struct ArrowBufferView message_view = {{message->data}, message->size_bytes}; + ASSERT_EQ(ArrowIpcDecoderVerifyHeader(decoder.get(), message_view, &error), + NANOARROW_OK); + ASSERT_EQ(ArrowIpcDecoderDecodeHeader(decoder.get(), message_view, &error), + NANOARROW_OK); + nanoarrow::UniqueArray decoded; + ASSERT_EQ( + ArrowIpcDecoderDecodeArray(decoder.get(), {{body->data}, body->size_bytes}, -1, + decoded.get(), NANOARROW_VALIDATION_LEVEL_FULL, &error), + NANOARROW_OK) + << error.message; + nanoarrow::UniqueArrayView decoded_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(decoded_view.get(), batch.schema(), &error), + NANOARROW_OK); + ASSERT_EQ(ArrowArrayViewSetArray(decoded_view.get(), decoded.get(), &error), + NANOARROW_OK); + int is_equal = 0; + ASSERT_EQ(ArrowArrayViewCompare(decoded_view.get(), batch.array_view(), + NANOARROW_COMPARE_IDENTICAL, &is_equal, &error), + NANOARROW_OK); + EXPECT_EQ(is_equal, 1) << error.message; + + // Scratch buffers are reused (by index) for the next message, which is what keeps the + // pointers captured above valid. Make a later prefix allocation fail, after an + // earlier buffer could have been queued with the compressor. + ASSERT_GT(state.destinations.size(), 1); + struct ArrowBuffer* failing_buffer = state.destinations[1]; + ArrowBufferReset(failing_buffer); + ASSERT_EQ(ArrowBufferSetAllocator(failing_buffer, FailingAllocator(&allocator_state)), + NANOARROW_OK); + body->size_bytes = 0; + EXPECT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), batch.array_view(), + body.get(), &error), + ENOMEM); + EXPECT_EQ(allocator_state.calls, 1); + EXPECT_TRUE(state.pending.empty()); +} + +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderDeferredCompressionErrors) { + for (bool fail_wait : {false, true}) { + SCOPED_TRACE(fail_wait ? "wait error" : "add error"); + struct ArrowError error; + CompressibleRecordBatch batch; + DeferredCompressor state; + state.fail_wait = fail_wait; + state.fail_on_add = fail_wait ? 0 : 2; + nanoarrow::ipc::UniqueEncoder encoder; + ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); + auto compressor = state.MakeCompressor(); + ASSERT_EQ(ArrowIpcEncoderSetCompressor(encoder.get(), &compressor), NANOARROW_OK); + nanoarrow::UniqueBuffer body; + EXPECT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), batch.array_view(), + body.get(), &error), + EIO); + EXPECT_STREQ(error.message, + fail_wait ? "Deferred wait failed" : "Deferred add failed"); + EXPECT_TRUE(state.pending.empty()); + EXPECT_EQ(state.waits, 1); + } +} + // Encode a batch with a body allocator that fails on the fail_on-th allocation, for // every fail_on until encoding succeeds, so that each allocation site reports ENOMEM static void TestEncodeAllocationFailures(enum ArrowIpcCompressionType codec) { @@ -1250,3 +1406,158 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSchemaDeclaresCompression) { ASSERT_NO_FATAL_FAILURE(encode_and_decode_schema(message)); EXPECT_EQ(decoder->feature_flags & NANOARROW_IPC_FEATURE_COMPRESSED_BODY, 0); } + +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderFooterCompressionHistory) { + for (bool dictionary_batch : {false, true}) { + SCOPED_TRACE(dictionary_batch ? "dictionary batch" : "record batch"); + struct ArrowError error; + CompressibleRecordBatch batch; + DeferredCompressor state; + nanoarrow::ipc::UniqueEncoder encoder; + ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); + auto compressor = state.MakeCompressor(); + ASSERT_EQ(ArrowIpcEncoderSetCompressor(encoder.get(), &compressor), NANOARROW_OK); + + // Low-level callers can encode a body without first encoding a Schema message. + nanoarrow::UniqueBuffer body, message; + if (dictionary_batch) { + ASSERT_EQ(ArrowIpcEncoderEncodeSimpleDictionaryBatch( + encoder.get(), 0, false, batch.array_view()->children[0], body.get(), + &error), + NANOARROW_OK); + } else { + ASSERT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), batch.array_view(), + body.get(), &error), + NANOARROW_OK); + } + ASSERT_EQ( + ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/true, message.get()), + NANOARROW_OK); + ASSERT_EQ( + ArrowIpcEncoderSetCompression(encoder.get(), NANOARROW_IPC_COMPRESSION_TYPE_NONE, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), + NANOARROW_OK); + + nanoarrow::ipc::UniqueFooter footer; + ASSERT_EQ(ArrowSchemaDeepCopy(batch.schema(), &footer->schema), NANOARROW_OK); + auto check_footer = [&](bool expected) { + ASSERT_EQ(ArrowIpcEncoderEncodeFooter(encoder.get(), footer.get(), &error), + NANOARROW_OK); + nanoarrow::UniqueBuffer buffer; + ASSERT_EQ(ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/false, + buffer.get()), + NANOARROW_OK); + int32_t footer_size = static_cast(buffer->size_bytes); + if (ArrowIpcSystemEndianness() == NANOARROW_IPC_ENDIANNESS_BIG) { + footer_size = static_cast(bswap32(static_cast(footer_size))); + } + ASSERT_EQ(ArrowBufferAppendInt32(buffer.get(), footer_size), NANOARROW_OK); + ASSERT_EQ(ArrowBufferAppend(buffer.get(), "ARROW1", 6), NANOARROW_OK); + + nanoarrow::ipc::UniqueDecoder decoder; + ASSERT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK); + struct ArrowBufferView view = {{buffer->data}, buffer->size_bytes}; + ASSERT_EQ(ArrowIpcDecoderVerifyFooter(decoder.get(), view, &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcDecoderDecodeFooter(decoder.get(), view, &error), NANOARROW_OK) + << error.message; + EXPECT_EQ((decoder->feature_flags & NANOARROW_IPC_FEATURE_COMPRESSED_BODY) != 0, + expected); + }; + ASSERT_NO_FATAL_FAILURE(check_footer(true)); + + // A new schema starts a new file's history on the same encoder. + ASSERT_EQ(ArrowIpcEncoderEncodeSchema(encoder.get(), batch.schema(), &error), + NANOARROW_OK); + message->size_bytes = 0; + ASSERT_EQ( + ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/true, message.get()), + NANOARROW_OK); + ASSERT_NO_FATAL_FAILURE(check_footer(false)); + } +} + +// The scratch buffers for compressed bodies grow when a message has more buffers than +// any encoded before it; the existing ones are moved and stay usable +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderCompressedBuffersGrow) { + struct ArrowError error; + CompressibleRecordBatch batch; + DeferredCompressor state; + nanoarrow::ipc::UniqueEncoder encoder; + ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); + auto compressor = state.MakeCompressor(); + ASSERT_EQ(ArrowIpcEncoderSetCompressor(encoder.get(), &compressor), NANOARROW_OK); + + // A DictionaryBatch of a single int32 column needs two scratch buffers... + nanoarrow::UniqueBuffer body, message; + ASSERT_EQ( + ArrowIpcEncoderEncodeSimpleDictionaryBatch( + encoder.get(), 0, false, batch.array_view()->children[0], body.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ( + ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/true, message.get()), + NANOARROW_OK); + size_t n_small = state.destinations.size(); + EXPECT_GT(n_small, 0); + + // ...and the RecordBatch of all three columns needs more + body->size_bytes = 0; + message->size_bytes = 0; + ASSERT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), batch.array_view(), + body.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ( + ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/true, message.get()), + NANOARROW_OK); + EXPECT_GT(state.destinations.size() - n_small, n_small); + + nanoarrow::ipc::UniqueDecoder decoder; + ASSERT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK); + ASSERT_EQ(ArrowIpcDecoderSetSchema(decoder.get(), batch.schema(), &error), + NANOARROW_OK); + struct ArrowBufferView message_view = {{message->data}, message->size_bytes}; + ASSERT_EQ(ArrowIpcDecoderVerifyHeader(decoder.get(), message_view, &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcDecoderDecodeHeader(decoder.get(), message_view, &error), + NANOARROW_OK) + << error.message; + nanoarrow::UniqueArray decoded; + ASSERT_EQ( + ArrowIpcDecoderDecodeArray(decoder.get(), {{body->data}, body->size_bytes}, -1, + decoded.get(), NANOARROW_VALIDATION_LEVEL_FULL, &error), + NANOARROW_OK) + << error.message; + nanoarrow::UniqueArrayView decoded_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(decoded_view.get(), batch.schema(), &error), + NANOARROW_OK); + ASSERT_EQ(ArrowArrayViewSetArray(decoded_view.get(), decoded.get(), &error), + NANOARROW_OK); + int is_equal = 0; + ASSERT_EQ(ArrowArrayViewCompare(decoded_view.get(), batch.array_view(), + NANOARROW_COMPARE_IDENTICAL, &is_equal, &error), + NANOARROW_OK); + EXPECT_EQ(is_equal, 1) << error.message; +} + +// A compressor that reports success without producing output is an error rather +// than a buffer with a length prefix and no payload +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderCompressorWithoutOutput) { + struct ArrowError error; + CompressibleRecordBatch batch; + DeferredCompressor state; + state.produce_nothing = true; + nanoarrow::ipc::UniqueEncoder encoder; + ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); + auto compressor = state.MakeCompressor(); + ASSERT_EQ(ArrowIpcEncoderSetCompressor(encoder.get(), &compressor), NANOARROW_OK); + + nanoarrow::UniqueBuffer body; + EXPECT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), batch.array_view(), + body.get(), &error), + EIO); + EXPECT_THAT(error.message, + ::testing::StartsWith("Compressor produced no output for a buffer of")); +} diff --git a/src/nanoarrow/ipc/writer_test.cc b/src/nanoarrow/ipc/writer_test.cc index 5f12d27c2..3b08d0c89 100644 --- a/src/nanoarrow/ipc/writer_test.cc +++ b/src/nanoarrow/ipc/writer_test.cc @@ -566,6 +566,59 @@ static void TestCompressedWriting(enum ArrowIpcCompressionType codec) { } } +TEST(NanoarrowIpcWriter, FileRetainsCompressionDeclaration) { + if (ArrowIpcGetLZ4CompressionFunction() == nullptr && + ArrowIpcGetZstdCompressionFunction() == nullptr) { + GTEST_SKIP() << "nanoarrow_ipc not built with NANOARROW_IPC_WITH_LZ4 or " + "NANOARROW_IPC_WITH_ZSTD"; + } + + for (auto codec : + {NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME, NANOARROW_IPC_COMPRESSION_TYPE_ZSTD}) { + int min_level, max_level; + if (ArrowIpcGetCompressionLevelRange(codec, &min_level, &max_level) == ENOTSUP) { + continue; + } + SCOPED_TRACE(ArrowIpcCompressionTypeToString(codec)); + for (bool write_compressed_batch : {false, true}) { + SCOPED_TRACE(write_compressed_batch ? "compressed batch" : "schema only"); + struct ArrowError error; + nanoarrow::UniqueSchema schema; + nanoarrow::UniqueArray array; + ASSERT_NO_FATAL_FAILURE(InitCompressibleBatch(schema.get(), array.get())); + nanoarrow::UniqueArrayView view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(view.get(), schema.get(), &error), + NANOARROW_OK); + ASSERT_EQ(ArrowArrayViewSetArray(view.get(), array.get(), &error), NANOARROW_OK); + nanoarrow::UniqueBuffer output; + nanoarrow::ipc::UniqueOutputStream stream; + ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream.get(), output.get()), NANOARROW_OK); + nanoarrow::ipc::UniqueWriter writer; + ASSERT_EQ(ArrowIpcWriterInit(writer.get(), stream.get()), NANOARROW_OK); + ASSERT_EQ(ArrowIpcWriterSetCompression( + writer.get(), codec, NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), + NANOARROW_OK); + ASSERT_EQ(ArrowIpcWriterStartFile(writer.get(), &error), NANOARROW_OK); + ASSERT_EQ(ArrowIpcWriterWriteSchema(writer.get(), schema.get(), &error), + NANOARROW_OK); + if (write_compressed_batch) { + ASSERT_EQ(ArrowIpcWriterWriteArrayView(writer.get(), view.get(), &error), + NANOARROW_OK); + } + ASSERT_EQ( + ArrowIpcWriterSetCompression(writer.get(), NANOARROW_IPC_COMPRESSION_TYPE_NONE, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error), + NANOARROW_OK); + ASSERT_EQ(ArrowIpcWriterWriteArrayView(writer.get(), view.get(), &error), + NANOARROW_OK); + ASSERT_EQ(ArrowIpcWriterWriteArrayView(writer.get(), nullptr, &error), + NANOARROW_OK); + ASSERT_EQ(ArrowIpcWriterFinalizeFile(writer.get(), &error), NANOARROW_OK); + ASSERT_NO_FATAL_FAILURE(CheckDeclaresCompression(output.get(), true, true)); + } + } +} + TEST(NanoarrowIpcWriter, CompressedWritingLZ4) { if (ArrowIpcGetLZ4CompressionFunction() == nullptr) { GTEST_SKIP() << "nanoarrow_ipc not built with NANOARROW_IPC_WITH_LZ4"; diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index 6ab7ae483..8a29cf5fb 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -457,18 +457,19 @@ struct ArrowIpcCompressor { /// \brief Queue a buffer for compression /// /// Compresses src using compression_type and appends the compressed bytes to dst, - /// preserving any content already in dst. The content of dst is undefined until the - /// next call to compress_wait() returns NANOARROW_OK, and the caller must not use - /// src or dst until then. + /// preserving any content already in dst. The caller must keep src and dst valid + /// and must not access them until the queued work has completed or been cancelled. + /// The content of dst is only valid after compress_wait() returns NANOARROW_OK. ArrowErrorCode (*compress_add)(struct ArrowIpcCompressor* compressor, struct ArrowBufferView src, struct ArrowBuffer* dst, struct ArrowError* error); /// \brief Wait for any unfinished calls to compress_add to complete /// - /// Returns NANOARROW_OK if all pending calls completed. Returns ETIMEDOUT if not all - /// remaining calls completed within timeout_ms (a negative timeout waits - /// indefinitely). + /// Returns NANOARROW_OK if all pending calls completed successfully. Returns + /// ETIMEDOUT if not all remaining calls completed within timeout_ms. A negative + /// timeout waits indefinitely and must complete or cancel all queued work before + /// returning, including when returning an error. ArrowErrorCode (*compress_wait)(struct ArrowIpcCompressor* compressor, int64_t timeout_ms, struct ArrowError* error); @@ -1013,9 +1014,11 @@ ArrowIpcEncoderSetMessageMetadata(struct ArrowIpcEncoder* encoder, /// Installs an ArrowIpcSerialCompressor() configured with compression_type and /// compression_level, replacing any compressor previously set with this function or /// with ArrowIpcEncoderSetCompressor(). NANOARROW_IPC_COMPRESSION_TYPE_NONE removes the -/// compressor. The setting persists until it is changed. Schema messages (and file -/// footers) encoded while a compressor is set declare the COMPRESSED_BODY feature, so -/// compression should be set before the schema is encoded. +/// compressor. The setting persists until it is changed. Schema messages encoded while +/// a compressor is set declare the COMPRESSED_BODY feature, so compression should be +/// set before the schema is encoded. File footers also declare this feature if +/// compression was declared or used since the most recent Schema message, even if the +/// compressor has since been removed. /// /// The body buffers of every RecordBatch or DictionaryBatch encoded while a compressor /// is set are compressed as described by the Arrow IPC format: each non-empty buffer @@ -1146,9 +1149,10 @@ NANOARROW_DLL void ArrowIpcWriterReset(struct ArrowIpcWriter* writer); /// /// See ArrowIpcEncoderSetCompression(). Compression applies to record batches and /// dictionary batches written after this call (in both stream and file mode) and may -/// be changed between batches. Set it before writing the schema so that the stream -/// (or the file footer) declares the COMPRESSED_BODY feature. Use -/// NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT for the codec's default level. +/// be changed between batches. Set it before writing the schema so that the Schema +/// message declares the COMPRESSED_BODY feature (a file footer declares it whenever +/// compression was declared or used). Use NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT for +/// the codec's default level. /// /// Returns EINVAL for an unknown compression type or an out-of-range compression_level, /// and ENOTSUP if the compression type is not supported by this build of nanoarrow.