You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by pa...@apache.org on 2026/09/09 17:10:28 UTC
(arrow-nanoarrow) branch main updated: feat: support writing LZ4 and ZSTD compressed IPC streams and files (#935)
This is an automated email from the ASF dual-hosted git repository.
paleolimbot pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-nanoarrow.git
The following commit(s) were added to refs/heads/main by this push:
new 66579451 feat: support writing LZ4 and ZSTD compressed IPC streams and files (#935)
66579451 is described below
commit 665794510011c07d4a6932b248608c2c02113796
Author: Pedro Holanda <pe...@gmail.com>
AuthorDate: Wed Sep 9 19:10:21 2026 +0200
feat: support writing LZ4 and ZSTD compressed IPC streams and files (#935)
This PR adds LZ4 and ZSTD compression to nanoarrow’s IPC encoder and
writer, allowing applications to write compressed Arrow streams and
files. The motivation comes from the duckdb-arrow extension, where we
want nanoarrow to handle compression while the extension takes care of
the SQL options. To support this, callers can select a codec and
compression level, query the supported level range, and convert between
codec names and enum values without needing to include the compression
libraries’ headers themselves.
Compression happens per buffer within each record batch, following the
Arrow IPC format. If compressing a buffer does not make it smaller, we
store it uncompressed, and empty buffers remain empty. The
implementation provides a serial compressor by default, with an
interface for applications to supply their own. Unsupported codecs and
out-of-range levels are rejected when configuring the built-in
compressor, so callers can report these errors before writing data. The
added tests cover compression and decompression roundtrips, compressed
streams and files, compression levels, and builds with either codec
unavailable.
---
CMakeLists.txt | 18 +-
src/nanoarrow/ipc/codecs.c | 276 +++++++++-
src/nanoarrow/ipc/codecs_test.cc | 397 ++++++++++++++
src/nanoarrow/ipc/decoder.c | 59 ++-
src/nanoarrow/ipc/encoder.c | 356 +++++++++++--
src/nanoarrow/ipc/encoder_test.cc | 1060 +++++++++++++++++++++++++++++++++++++
src/nanoarrow/ipc/files_test.cc | 78 ++-
src/nanoarrow/ipc/reader.c | 6 -
src/nanoarrow/ipc/writer.c | 11 +
src/nanoarrow/ipc/writer_test.cc | 318 +++++++++++
src/nanoarrow/nanoarrow_ipc.h | 195 +++++++
src/nanoarrow/nanoarrow_ipc.hpp | 23 +
12 files changed, 2683 insertions(+), 114 deletions(-)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 6335b902..848b933e 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.c b/src/nanoarrow/ipc/codecs.c
index ab297989..e5cd0b31 100644
--- a/src/nanoarrow/ipc/codecs.c
+++ b/src/nanoarrow/ipc/codecs.c
@@ -16,6 +16,7 @@
// under the License.
#include <inttypes.h>
+#include <limits.h>
#include "nanoarrow/nanoarrow_ipc.h"
@@ -44,6 +45,32 @@ static ArrowErrorCode ArrowIpcDecompressZstd(struct ArrowBufferView src, uint8_t
return NANOARROW_OK;
}
+
+static ArrowErrorCode ArrowIpcCompressZstd(struct ArrowBufferView src,
+ 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)) {
+ 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, compression_level);
+ 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 +81,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 <lz4.h>
#include <lz4frame.h>
@@ -103,6 +138,40 @@ static ArrowErrorCode ArrowIpcDecompressLZ4(struct ArrowBufferView src, uint8_t*
NANOARROW_UNUSED(LZ4F_freeDecompressionContext(ctx));
return NANOARROW_OK;
}
+
+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() 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),
+ 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 +182,100 @@ ArrowIpcDecompressFunction ArrowIpcGetLZ4DecompressionFunction(void) {
#endif
}
+ArrowIpcCompressFunction ArrowIpcGetLZ4CompressionFunction(void) {
+#if defined(NANOARROW_IPC_WITH_LZ4)
+ return &ArrowIpcCompressLZ4;
+#else
+ return NULL;
+#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);
+ // (unused when neither codec is built in)
+ 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) {
+ switch (compression_type) {
+ case NANOARROW_IPC_COMPRESSION_TYPE_ZSTD:
+ case NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME:
+ return 1;
+ default:
+ return 0;
+ }
+}
+
+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 "<unknown compression type>";
+ }
+}
+
+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];
};
@@ -124,18 +287,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 +345,100 @@ 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];
+ int compression_level;
+};
+
+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;
+
+ 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;
+ }
+
+ // 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;
+}
+
+static void ArrowIpcSerialCompressorRelease(struct ArrowIpcCompressor* compressor) {
+ ArrowFree(compressor->private_data);
+ compressor->release = NULL;
+}
+
+ArrowErrorCode ArrowIpcSerialCompressor(struct ArrowIpcCompressor* compressor,
+ enum ArrowIpcCompressionType compression_type,
+ int compression_level) {
+ compressor->release = 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(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_add = &ArrowIpcSerialCompressorAdd;
+ compressor->compress_wait = &ArrowIpcSerialCompressorWait;
+ compressor->release = &ArrowIpcSerialCompressorRelease;
+ 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 76278174..b6d2c89a 100644
--- a/src/nanoarrow/ipc/codecs_test.cc
+++ b/src/nanoarrow/ipc/codecs_test.cc
@@ -16,6 +16,9 @@
// under the License.
#include <cstring>
+#include <limits>
+#include <string>
+#include <vector>
#include <gmock/gmock-matchers.h>
#include <gtest/gtest.h>
@@ -152,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);
@@ -185,3 +195,390 @@ TEST(NanoarrowIpcTest, SerialDecompressor) {
EXPECT_STREQ(error.message,
"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");
+ // 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<enum ArrowIpcCompressionType>(3);
+ EXPECT_STREQ(ArrowIpcCompressionTypeToString(unknown_type),
+ "<unknown compression type>");
+
+ 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("<unknown compression type>", &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);
+}
+
+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<enum ArrowIpcCompressionType>(3);
+ 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.
+static int64_t TestCompressRoundtrip(ArrowIpcCompressFunction compress,
+ ArrowIpcDecompressFunction decompress,
+ int compression_level,
+ const std::vector<uint8_t>& 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<int64_t>(input.size())},
+ compression_level, 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<uint8_t> 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<int64_t>(output.size()), &error),
+ NANOARROW_OK)
+ << error.message;
+ EXPECT_EQ(output, input);
+
+ return compressed->size_bytes - existing_size;
+}
+
+static std::vector<uint8_t> CompressibleInput(size_t n) {
+ std::vector<uint8_t> out(n);
+ for (size_t i = 0; i < n; i++) {
+ out[i] = static_cast<uint8_t>(i % 7);
+ }
+ return out;
+}
+
+// Check compress/decompress on empty, small, and multi-block inputs at each level
+static void TestCompressionFunctions(ArrowIpcCompressFunction compress,
+ ArrowIpcDecompressFunction decompress,
+ const std::vector<int>& compression_levels) {
+ ASSERT_NE(compress, nullptr);
+ 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, {});
+ TestCompressRoundtrip(
+ compress, decompress, level,
+ std::vector<uint8_t>(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<int64_t>(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) {
+#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";
+ }
+ // Default, a negative (fast) level, the lowest regular level, and a high level
+ TestCompressionFunctions(ArrowIpcGetZstdCompressionFunction(),
+ ArrowIpcGetZstdDecompressionFunction(),
+ {NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, -5, 1, 19});
+}
+
+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";
+ }
+ // Default (fast), acceleration, the last fast level, and LZ4HC levels
+ TestCompressionFunctions(ArrowIpcGetLZ4CompressionFunction(),
+ ArrowIpcGetLZ4DecompressionFunction(),
+ {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<int>::min(), std::numeric_limits<int>::min() + 1,
+ std::numeric_limits<int>::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);
+ }
+}
+
+// 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<FailingAllocatorState*>(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;
+
+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) {
+ struct ArrowError error {};
+ nanoarrow::ipc::UniqueCompressor compressor;
+
+ // An invalid compression type is rejected at construction
+ // NOLINTNEXTLINE(clang-analyzer-optin.core.EnumCastOutOfRange)
+ auto unknown_type = static_cast<enum ArrowIpcCompressionType>(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);
+
+ // 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_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_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)];
+ 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_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");
+ }
+
+ // 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_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 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<int64_t>(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_add(compressor.get(),
+ {{kUncompressed012}, sizeof(kUncompressed012)},
+ dst.get(), &error),
+ EIO);
+ EXPECT_STREQ(error.message, "FailCompress() failed");
+}
diff --git a/src/nanoarrow/ipc/decoder.c b/src/nanoarrow/ipc/decoder.c
index b13d0bd6..00dbcf2c 100644
--- a/src/nanoarrow/ipc/decoder.c
+++ b/src/nanoarrow/ipc/decoder.c
@@ -1390,8 +1390,35 @@ static int ArrowIpcDecoderDecodeSchemaHeader(struct ArrowIpcDecoder* decoder,
return NANOARROW_OK;
}
+// Decode the BodyCompression of a RecordBatch (which is also nested in every
+// DictionaryBatch) into codec_out
+static ArrowErrorCode ArrowIpcDecoderDecodeBodyCompression(
+ ns(RecordBatch_table_t) batch, enum ArrowIpcCompressionType* codec_out,
+ struct ArrowError* error) {
+ if (batch == NULL || !ns(RecordBatch_compression_is_present(batch))) {
+ *codec_out = NANOARROW_IPC_COMPRESSION_TYPE_NONE;
+ return NANOARROW_OK;
+ }
+
+ ns(BodyCompression_table_t) compression = ns(RecordBatch_compression(batch));
+ ns(CompressionType_enum_t) codec = ns(BodyCompression_codec(compression));
+ switch (codec) {
+ case ns(CompressionType_LZ4_FRAME):
+ *codec_out = NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME;
+ return NANOARROW_OK;
+ case ns(CompressionType_ZSTD):
+ *codec_out = NANOARROW_IPC_COMPRESSION_TYPE_ZSTD;
+ return NANOARROW_OK;
+ default:
+ ArrowErrorSet(error, "Unrecognized RecordBatch BodyCompression codec value: %d",
+ (int)codec);
+ return EINVAL;
+ }
+}
+
static int ArrowIpcDecoderDecodeDictionaryBatchHeader(
- struct ArrowIpcDecoder* decoder, flatbuffers_generic_t message_header) {
+ struct ArrowIpcDecoder* decoder, flatbuffers_generic_t message_header,
+ struct ArrowError* error) {
struct ArrowIpcDecoderPrivate* private_data =
(struct ArrowIpcDecoderPrivate*)decoder->private_data;
@@ -1399,6 +1426,10 @@ static int ArrowIpcDecoderDecodeDictionaryBatchHeader(
private_data->dictionary.id = ns(DictionaryBatch_id(dictionary));
private_data->dictionary.is_delta = ns(DictionaryBatch_isDelta(dictionary));
+ // The dictionary values are a RecordBatch with its own compression setting
+ NANOARROW_RETURN_NOT_OK(ArrowIpcDecoderDecodeBodyCompression(
+ ns(DictionaryBatch_data(dictionary)), &decoder->codec, error));
+
decoder->dictionary = &private_data->dictionary;
return NANOARROW_OK;
}
@@ -1437,24 +1468,8 @@ static int ArrowIpcDecoderDecodeRecordBatchHeader(struct ArrowIpcDecoder* decode
return EINVAL;
}
- if (ns(RecordBatch_compression_is_present(batch))) {
- ns(BodyCompression_table_t) compression = ns(RecordBatch_compression(batch));
- ns(CompressionType_enum_t) codec = ns(BodyCompression_codec(compression));
- switch (codec) {
- case ns(CompressionType_LZ4_FRAME):
- decoder->codec = NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME;
- break;
- case ns(CompressionType_ZSTD):
- decoder->codec = NANOARROW_IPC_COMPRESSION_TYPE_ZSTD;
- break;
- default:
- ArrowErrorSet(error, "Unrecognized RecordBatch BodyCompression codec value: %d",
- (int)codec);
- return EINVAL;
- }
- } else {
- decoder->codec = NANOARROW_IPC_COMPRESSION_TYPE_NONE;
- }
+ NANOARROW_RETURN_NOT_OK(
+ ArrowIpcDecoderDecodeBodyCompression(batch, &decoder->codec, error));
// Copying field node and buffer information is separate so as only to pay for the
// nodes that are actually accessed.
@@ -1721,7 +1736,7 @@ ArrowErrorCode ArrowIpcDecoderDecodeHeader(struct ArrowIpcDecoder* decoder,
break;
case ns(MessageHeader_DictionaryBatch):
NANOARROW_RETURN_NOT_OK(
- ArrowIpcDecoderDecodeDictionaryBatchHeader(decoder, message_header));
+ ArrowIpcDecoderDecodeDictionaryBatchHeader(decoder, message_header, error));
break;
case ns(MessageHeader_RecordBatch):
NANOARROW_RETURN_NOT_OK(
@@ -2811,8 +2826,10 @@ static ArrowErrorCode ArrowIpcDecoderDecodeDictionaryInternal(
(struct ArrowIpcDecoderPrivate*)dictionary->decoder.private_data;
dictionary->decoder.message_type = NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH;
dictionary_decoder_private_data->last_message = record_batch;
- // Transfer the endianness setting so that buffers are byte-swapped if needed
+ // Transfer the endianness and compression settings so that buffers are byte-swapped
+ // and decompressed if needed (the nested decoder uses a default decompressor)
dictionary_decoder_private_data->endianness = private_data->endianness;
+ dictionary->decoder.codec = decoder->codec;
struct ArrowArrayView* array_view;
NANOARROW_RETURN_NOT_OK(ArrowIpcDecoderDecodeArrayViewInternal(
diff --git a/src/nanoarrow/ipc/encoder.c b/src/nanoarrow/ipc/encoder.c
index d1435f49..ab7f0b60 100644
--- a/src/nanoarrow/ipc/encoder.c
+++ b/src/nanoarrow/ipc/encoder.c
@@ -16,6 +16,7 @@
// under the License.
#include <errno.h>
+#include <inttypes.h>
#include <stdio.h>
#include <string.h>
@@ -49,6 +50,16 @@ 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;
+ // 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)
+ struct ArrowBuffer* compressed_buffers;
+ int64_t n_compressed_buffers;
};
ArrowErrorCode ArrowIpcEncoderInit(struct ArrowIpcEncoder* encoder) {
@@ -69,6 +80,11 @@ ArrowErrorCode ArrowIpcEncoderInit(struct ArrowIpcEncoder* encoder) {
ArrowBufferInit(&private->nodes);
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;
return NANOARROW_OK;
}
@@ -82,6 +98,14 @@ 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);
+ }
+ 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));
@@ -119,6 +143,79 @@ 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 &&
+ 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;
+ return NANOARROW_OK;
+}
+
+ArrowErrorCode ArrowIpcEncoderSetCompression(
+ struct ArrowIpcEncoder* encoder, enum ArrowIpcCompressionType compression_type,
+ int compression_level, struct ArrowError* error) {
+ NANOARROW_DCHECK(encoder != NULL && encoder->private_data != NULL);
+ struct ArrowIpcEncoderPrivate* private =
+ (struct ArrowIpcEncoderPrivate*)encoder->private_data;
+
+ switch (compression_type) {
+ case NANOARROW_IPC_COMPRESSION_TYPE_NONE:
+ if (private->compressor.release != NULL) {
+ private->compressor.release(&private->compressor);
+ }
+ return NANOARROW_OK;
+ case NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME:
+ case NANOARROW_IPC_COMPRESSION_TYPE_ZSTD:
+ break;
+ default:
+ ArrowErrorSet(error, "Unknown compression type with value %d",
+ (int)compression_type);
+ 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 (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;
+ }
+
+ struct ArrowIpcCompressor compressor;
+ NANOARROW_RETURN_NOT_OK_WITH_ERROR(
+ ArrowIpcSerialCompressor(&compressor, compression_type, compression_level), error);
+ return ArrowIpcEncoderSetCompressor(encoder, &compressor);
+}
+
static ArrowErrorCode ArrowIpcEncoderWriteContinuationAndSize(struct ArrowBuffer* out,
size_t size) {
_NANOARROW_CHECK_UPPER_LIMIT(size, INT32_MAX);
@@ -560,7 +657,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);
@@ -593,6 +690,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;
@@ -623,8 +725,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);
@@ -633,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;
}
@@ -695,10 +800,191 @@ static ArrowErrorCode ArrowIpcEncoderBuildContiguousBodyBufferCallback(
return NANOARROW_OK;
}
-static ArrowErrorCode ArrowIpcEncoderEncodeRecordBatchImpl(
+// 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;
+ }
+
+ 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;
+}
+
+// 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);
+
+ // 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;
+ }
+
+ // 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);
+ }
+
+ 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;
+ NANOARROW_UNUSED(
+ private->compressor.compress_wait(&private->compressor, -1, &ignored));
+ return result;
+ }
+ }
+
+ 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;
+ }
+
+ // 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)
+ 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;
+ }
+
+ return NANOARROW_OK;
+}
+
+// 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,
- const struct ArrowArrayView* array_view, struct ArrowBuffer* buffers,
- struct ArrowBuffer* nodes, struct ArrowError* error) {
+ struct ArrowError* error) {
+ struct ArrowIpcEncoderPrivate* private =
+ (struct ArrowIpcEncoderPrivate*)encoder->private_data;
+
+ if (ArrowIpcEncoderCodec(private) != NANOARROW_IPC_COMPRESSION_TYPE_NONE) {
+ NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderCompressBuffers(private, error));
+ }
+
+ 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(&private->buffers, &buffer, sizeof(buffer)), error);
+ }
+
+ return NANOARROW_OK;
+}
+
+// Add the BodyCompression table to the RecordBatch currently being built, if any
+// compression is enabled. Bodies of RecordBatch and DictionaryBatch messages are both
+// built by the same buffer encoder, so both need this.
+static ArrowErrorCode ArrowIpcEncoderEncodeBodyCompression(
+ struct ArrowIpcEncoderPrivate* private, struct ArrowError* error) {
+ 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, codec), error);
+ 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;
+}
+
+// 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 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;
@@ -706,29 +992,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;
}
@@ -762,10 +1034,14 @@ 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);
+ NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderEncodeBodyCompression(private, error));
+
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,
@@ -824,25 +1100,18 @@ static ArrowErrorCode ArrowIpcEncoderEncodeDictionaryBatch(
FLATCC_RETURN_UNLESS_0(DictionaryBatch_id_add(builder, dictionary_id), error);
FLATCC_RETURN_UNLESS_0(DictionaryBatch_data_start(builder), error);
FLATCC_RETURN_UNLESS_0(RecordBatch_length_add(builder, values_view->length), error);
+ NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderEncodeBodyCompression(private, error));
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,
@@ -907,8 +1176,11 @@ 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,
+ 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 4fa0e4fe..6c79babd 100644
--- a/src/nanoarrow/ipc/encoder_test.cc
+++ b/src/nanoarrow/ipc/encoder_test.cc
@@ -15,8 +15,10 @@
// specific language governing permissions and limitations
// under the License.
+#include <gmock/gmock-matchers.h>
#include <gtest/gtest.h>
+#include <cstring>
#include <string>
#include <utility>
#include <vector>
@@ -501,3 +503,1061 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderDictionaryBatch) {
EXPECT_GT(message_buffer->size_bytes, 8);
EXPECT_EQ(message_buffer->size_bytes % 8, 0);
}
+
+// 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 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<int64_t>(bswap64(static_cast<uint64_t>(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, NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &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;
+ 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(), NANOARROW_IPC_COMPRESSION_TYPE_NONE,
+ NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &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, 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<int32_t> 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<int*>(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";
+ }
+ 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;
+
+ // 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<enum ArrowIpcCompressionType>(3);
+ EXPECT_EQ(
+ ArrowIpcEncoderSetCompression(encoder.get(), unknown_type,
+ NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error),
+ EINVAL);
+ EXPECT_STREQ(error.message, "Unknown compression type with value 3");
+
+ // NONE is always supported
+ 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,
+ 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,
+ 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,
+ NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error),
+ NANOARROW_OK)
+ << error.message;
+#else
+ 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
+}
+
+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_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;
+ 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 configured for LZ4 that explicitly does not support it
+ nanoarrow::ipc::UniqueCompressor compressor;
+ 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);
+
+ 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, support is not checked until a batch is encoded
+ 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");
+
+ // 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
+// "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_IPC_COMPRESSION_TYPE_ZSTD, 11),
+ NANOARROW_OK);
+ ASSERT_EQ(
+ ArrowIpcSerialCompressorSetFunction(
+ compressor.get(), NANOARROW_IPC_COMPRESSION_TYPE_ZSTD, &RecordLevelAndCopy),
+ NANOARROW_OK);
+ ASSERT_EQ(ArrowIpcEncoderSetCompressor(encoder.get(), compressor.get()), NANOARROW_OK);
+
+ 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;
+}
+
+// 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<FailingAllocatorState*>(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;
+}
+
+// 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<Job> pending;
+ std::vector<struct ArrowBuffer*> 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<DeferredCompressor*>(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<DeferredCompressor*>(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<DeferredCompressor*>(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) {
+ 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);
+}
+
+// DictionaryBatch bodies are compressed like RecordBatch bodies and must declare it
+static void TestCompressedDictionaryBatch(enum ArrowIpcCompressionType codec) {
+ struct ArrowError error;
+
+ // A dictionary-encoded int32 -> utf8 field, which gets dictionary id 0
+ 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);
+ ASSERT_EQ(ArrowSchemaAllocateDictionary(schema->children[0]), NANOARROW_OK);
+ ASSERT_EQ(
+ ArrowSchemaInitFromType(schema->children[0]->dictionary, NANOARROW_TYPE_STRING),
+ NANOARROW_OK);
+
+ struct ArrowIpcDictionaryEncodings encodings;
+ ArrowIpcDictionaryEncodingsInit(&encodings);
+ struct ArrowIpcDictionaryEncoding encoding;
+ encoding.id = 0;
+ encoding.kind = NANOARROW_IPC_DICTIONARY_KIND_DENSE_ARRAY;
+ encoding.schema = schema->children[0];
+ ASSERT_EQ(ArrowIpcDictionaryEncodingsAppend(&encodings, encoding), NANOARROW_OK);
+
+ // Repetitive values so that the dictionary body actually compresses
+ nanoarrow::UniqueArray values;
+ ASSERT_EQ(
+ ArrowArrayInitFromSchema(values.get(), schema->children[0]->dictionary, &error),
+ NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(ArrowArrayStartAppending(values.get()), NANOARROW_OK);
+ for (int i = 0; i < 1024; i++) {
+ std::string value = "value-" + std::to_string(i % 4);
+ ASSERT_EQ(ArrowArrayAppendString(values.get(), ArrowCharView(value.c_str())),
+ NANOARROW_OK);
+ }
+ ASSERT_EQ(ArrowArrayFinishBuildingDefault(values.get(), &error), NANOARROW_OK)
+ << error.message;
+ nanoarrow::UniqueArrayView values_view;
+ ASSERT_EQ(ArrowArrayViewInitFromSchema(values_view.get(),
+ schema->children[0]->dictionary, &error),
+ NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(ArrowArrayViewSetArray(values_view.get(), values.get(), &error), NANOARROW_OK)
+ << error.message;
+
+ // Encode the DictionaryBatch uncompressed (for reference) and compressed
+ nanoarrow::ipc::UniqueEncoder encoder;
+ ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK);
+ nanoarrow::UniqueBuffer uncompressed_body, message, body;
+ ASSERT_EQ(ArrowIpcEncoderEncodeSimpleDictionaryBatch(encoder.get(), /*dictionary_id=*/0,
+ /*is_delta=*/0, values_view.get(),
+ uncompressed_body.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(
+ ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/true, message.get()),
+ NANOARROW_OK);
+ message->size_bytes = 0;
+
+ ASSERT_EQ(ArrowIpcEncoderSetCompression(
+ encoder.get(), codec, NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error),
+ NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(ArrowIpcEncoderEncodeSimpleDictionaryBatch(encoder.get(), /*dictionary_id=*/0,
+ /*is_delta=*/0, values_view.get(),
+ body.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(
+ ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/true, message.get()),
+ NANOARROW_OK);
+ EXPECT_LT(body->size_bytes, uncompressed_body->size_bytes);
+
+ // The header is a DictionaryBatch whose values declare the codec
+ nanoarrow::ipc::UniqueDecoder decoder;
+ ASSERT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK);
+ ASSERT_EQ(ArrowIpcDecoderSetEndianness(decoder.get(), ArrowIpcSystemEndianness()),
+ 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;
+ EXPECT_EQ(decoder->message_type, NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH);
+ EXPECT_EQ(decoder->codec, codec);
+ ASSERT_NE(decoder->dictionary, nullptr);
+ EXPECT_EQ(decoder->dictionary->id, 0);
+ EXPECT_EQ(decoder->body_size_bytes, body->size_bytes);
+
+ // The values decode to the original array
+ struct ArrowIpcDictionaries dictionaries;
+ ASSERT_EQ(ArrowIpcDictionariesInit(&dictionaries, &encodings, &error), NANOARROW_OK)
+ << error.message;
+ struct ArrowBufferView body_view = {{body->data}, body->size_bytes};
+ int result = ArrowIpcDecoderDecodeDictionary(
+ decoder.get(), body_view, NANOARROW_VALIDATION_LEVEL_FULL, &dictionaries, &error);
+ EXPECT_EQ(result, NANOARROW_OK) << error.message;
+ if (result == NANOARROW_OK) {
+ const struct ArrowArray* decoded = nullptr;
+ ASSERT_EQ(ArrowIpcDictionariesFindCurrentValue(&dictionaries, 0, &decoded, &error),
+ NANOARROW_OK)
+ << error.message;
+ nanoarrow::UniqueArrayView decoded_view;
+ ASSERT_EQ(ArrowArrayViewInitFromSchema(decoded_view.get(),
+ schema->children[0]->dictionary, &error),
+ NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(ArrowArrayViewSetArray(decoded_view.get(), decoded, &error), NANOARROW_OK)
+ << error.message;
+ int is_equal = 0;
+ ASSERT_EQ(ArrowArrayViewCompare(decoded_view.get(), values_view.get(),
+ NANOARROW_COMPARE_IDENTICAL, &is_equal, &error),
+ NANOARROW_OK);
+ EXPECT_EQ(is_equal, 1) << error.message;
+ }
+
+ ArrowIpcDictionariesReset(&dictionaries);
+ ArrowIpcDictionaryEncodingsReset(&encodings);
+}
+
+TEST(NanoarrowIpcTest, NanoarrowIpcEncoderCompressedDictionaryBatchLZ4) {
+ if (ArrowIpcGetLZ4CompressionFunction() == nullptr) {
+ GTEST_SKIP() << "nanoarrow_ipc not built with NANOARROW_IPC_WITH_LZ4";
+ }
+ TestCompressedDictionaryBatch(NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME);
+}
+
+TEST(NanoarrowIpcTest, NanoarrowIpcEncoderCompressedDictionaryBatchZstd) {
+ if (ArrowIpcGetZstdCompressionFunction() == nullptr) {
+ GTEST_SKIP() << "nanoarrow_ipc not built with NANOARROW_IPC_WITH_ZSTD";
+ }
+ 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);
+}
+
+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<int32_t>(buffer->size_bytes);
+ if (ArrowIpcSystemEndianness() == NANOARROW_IPC_ENDIANNESS_BIG) {
+ footer_size = static_cast<int32_t>(bswap32(static_cast<uint32_t>(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/files_test.cc b/src/nanoarrow/ipc/files_test.cc
index 46c7203c..73fc9724 100644
--- a/src/nanoarrow/ipc/files_test.cc
+++ b/src/nanoarrow/ipc/files_test.cc
@@ -18,6 +18,9 @@
#include <errno.h>
#include <fstream>
#include <sstream>
+#include <string>
+#include <utility>
+#include <vector>
#include <zlib.h>
@@ -212,6 +215,7 @@ class TestFile {
ArrowErrorCode WriteNanoarrowStream(const nanoarrow::UniqueSchema& schema,
const std::vector<nanoarrow::UniqueArray>& arrays,
+ enum ArrowIpcCompressionType codec,
struct ArrowBuffer* buffer,
struct ArrowError* error) {
nanoarrow::ipc::UniqueOutputStream output_stream;
@@ -219,6 +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, NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, error));
nanoarrow::UniqueArrayView array_view;
NANOARROW_RETURN_NOT_OK(
@@ -259,14 +265,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<std::pair<std::string, enum ArrowIpcCompressionType>> 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<nanoarrow::UniqueBuffer> 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 +302,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<nanoarrow::UniqueArray> 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<nanoarrow::UniqueArray> 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/reader.c b/src/nanoarrow/ipc/reader.c
index 0b99b828..c699ca33 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.c b/src/nanoarrow/ipc/writer.c
index 00ea93ec..bb8c7971 100644
--- a/src/nanoarrow/ipc/writer.c
+++ b/src/nanoarrow/ipc/writer.c
@@ -226,6 +226,17 @@ void ArrowIpcWriterReset(struct ArrowIpcWriter* writer) {
memset(writer, 0, sizeof(struct ArrowIpcWriter));
}
+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,
+ compression_level, 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 c30dd603..3b08d0c8 100644
--- a/src/nanoarrow/ipc/writer_test.cc
+++ b/src/nanoarrow/ipc/writer_test.cc
@@ -405,3 +405,321 @@ TEST(NanoarrowIpcWriter, RoundtripDictionaryStream) {
<< error.message;
EXPECT_EQ(roundtrip_array->release, nullptr);
}
+
+// 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) using codec and
+// compression_level
+static void WriteCompressibleBatch(enum ArrowIpcCompressionType codec,
+ int compression_level, 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, compression_level, &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<struct ArrowIpcWriterPrivate*>(writer->private_data);
+ ASSERT_EQ(p->footer.record_batch_blocks.size_bytes, sizeof(struct ArrowIpcFileBlock));
+ auto* block =
+ reinterpret_cast<struct ArrowIpcFileBlock*>(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);
+}
+
+// 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");
+
+ 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));
+
+ // 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));
+ }
+}
+
+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";
+ }
+ 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;
+ // 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<enum ArrowIpcCompressionType>(3);
+ EXPECT_EQ(ArrowIpcWriterSetCompression(writer.get(), unknown_type,
+ NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error),
+ EINVAL);
+ EXPECT_STREQ(error.message, "Unknown compression type with value 3");
+
+#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
+}
+
+// A dictionary-encoded column written with compression: the DictionaryBatch is
+// compressed too and the reader decodes it
+TEST(NanoarrowIpcWriter, CompressedDictionaryStream) {
+ if (ArrowIpcGetLZ4CompressionFunction() == nullptr) {
+ GTEST_SKIP() << "nanoarrow_ipc not built with NANOARROW_IPC_WITH_LZ4";
+ }
+ struct ArrowError error;
+
+ nanoarrow::UniqueSchema schema;
+ nanoarrow::UniqueArray array;
+ ASSERT_NO_FATAL_FAILURE(MakeDictionaryStructArray(array.get(), schema.get()));
+ nanoarrow::UniqueArrayStream array_stream;
+ ASSERT_EQ(ArrowBasicArrayStreamInit(array_stream.get(), schema.get(), 1), NANOARROW_OK);
+ ArrowBasicArrayStreamSetArray(array_stream.get(), 0, array.get());
+
+ nanoarrow::UniqueBuffer output;
+ nanoarrow::ipc::UniqueOutputStream out_stream;
+ ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(out_stream.get(), output.get()), NANOARROW_OK);
+ nanoarrow::ipc::UniqueWriter writer;
+ ASSERT_EQ(ArrowIpcWriterInit(writer.get(), out_stream.get()), NANOARROW_OK);
+ ASSERT_EQ(
+ ArrowIpcWriterSetCompression(writer.get(), NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME,
+ NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, &error),
+ NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(ArrowIpcWriterWriteArrayStream(writer.get(), array_stream.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+
+ nanoarrow::ipc::UniqueInputStream input;
+ ASSERT_EQ(ArrowIpcInputStreamInitBuffer(input.get(), output.get()), NANOARROW_OK);
+ nanoarrow::UniqueArrayStream reader;
+ ASSERT_EQ(ArrowIpcArrayStreamReaderInit(reader.get(), input.get(), nullptr),
+ NANOARROW_OK);
+
+ nanoarrow::UniqueSchema roundtrip_schema;
+ ASSERT_EQ(ArrowArrayStreamGetSchema(reader.get(), roundtrip_schema.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+ nanoarrow::UniqueArray roundtrip_array;
+ ASSERT_EQ(ArrowArrayStreamGetNext(reader.get(), roundtrip_array.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(roundtrip_array->length, 3);
+
+ nanoarrow::UniqueArrayView view;
+ ASSERT_EQ(ArrowArrayViewInitFromSchema(view.get(), roundtrip_schema.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(ArrowArrayViewSetArray(view.get(), roundtrip_array.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+ struct ArrowArrayView* indices_view = view->children[0];
+ struct ArrowArrayView* values_view = indices_view->dictionary;
+ ASSERT_NE(values_view, nullptr);
+ EXPECT_EQ(ArrowArrayViewGetIntUnsafe(indices_view, 0), 0);
+ EXPECT_EQ(ArrowArrayViewGetIntUnsafe(indices_view, 1), 1);
+ EXPECT_EQ(ArrowArrayViewGetIntUnsafe(indices_view, 2), 0);
+ struct ArrowStringView v0 = ArrowArrayViewGetStringUnsafe(values_view, 0);
+ struct ArrowStringView v1 = ArrowArrayViewGetStringUnsafe(values_view, 1);
+ EXPECT_EQ(std::string(v0.data, v0.size_bytes), "foo");
+ EXPECT_EQ(std::string(v1.data, v1.size_bytes), "bar");
+}
diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h
index 0523ac38..8a29cf5f 100644
--- a/src/nanoarrow/nanoarrow_ipc.h
+++ b/src/nanoarrow/nanoarrow_ipc.h
@@ -33,6 +33,20 @@
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 ArrowIpcCompressionTypeToString \
+ 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 \
@@ -92,6 +106,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 \
@@ -108,6 +126,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 \
@@ -416,6 +436,127 @@ 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
+/// 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
+ ///
+ /// 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 Queue a buffer for compression
+ ///
+ /// Compresses src using compression_type and appends the compressed bytes to dst,
+ /// 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 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);
+
+ /// \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_add(), compress_wait(), or release().
+ void (*release)(struct ArrowIpcCompressor* compressor);
+
+ /// \brief Implementation-specific opaque data
+ void* private_data;
+};
+
+/// \brief A self-contained compression function
+///
+/// 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 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,
+ 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
+///
+/// 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
+///
+/// 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 Get the name of a compression type
+///
+/// Returns "none", "lz4", or "zstd", or "<unknown compression type>" 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);
+
+/// \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 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,
@@ -868,6 +1009,45 @@ NANOARROW_DLL ArrowErrorCode
ArrowIpcEncoderSetMessageMetadata(struct ArrowIpcEncoder* encoder,
struct ArrowBuffer* metadata, struct ArrowError* error);
+/// \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. 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
+/// 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 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).
+NANOARROW_DLL ArrowErrorCode ArrowIpcEncoderSetCompression(
+ struct ArrowIpcEncoder* encoder, enum ArrowIpcCompressionType compression_type,
+ int compression_level, struct ArrowError* error);
+
+/// \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 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);
+
/// \brief Encode an ArrowSchema
///
/// Returns ENOMEM if allocation fails, NANOARROW_OK otherwise.
@@ -965,6 +1145,21 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcWriterInit(
/// \brief Release all resources attached to a writer
NANOARROW_DLL void ArrowIpcWriterReset(struct ArrowIpcWriter* writer);
+/// \brief Compress the bodies of subsequently written batches with a built-in codec
+///
+/// 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 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.
+NANOARROW_DLL ArrowErrorCode ArrowIpcWriterSetCompression(
+ struct ArrowIpcWriter* writer, enum ArrowIpcCompressionType compression_type,
+ int compression_level, struct ArrowError* error);
+
/// \brief Write a schema to the output byte stream
///
/// Errors are propagated from the underlying encoder and output byte stream.
diff --git a/src/nanoarrow/nanoarrow_ipc.hpp b/src/nanoarrow/nanoarrow_ipc.hpp
index e6b23924..9629eba3 100644
--- a/src/nanoarrow/nanoarrow_ipc.hpp
+++ b/src/nanoarrow/nanoarrow_ipc.hpp
@@ -129,6 +129,26 @@ inline void release_pointer(struct ArrowIpcDecompressor* data) {
}
}
+template <>
+inline void init_pointer(struct ArrowIpcCompressor* data) {
+ data->compression_type = NANOARROW_IPC_COMPRESSION_TYPE_NONE;
+ 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 +235,9 @@ using UniqueEncoder = internal::Unique<struct ArrowIpcEncoder>;
/// \brief Class wrapping a unique struct ArrowIpcDecompressor
using UniqueDecompressor = internal::Unique<struct ArrowIpcDecompressor>;
+/// \brief Class wrapping a unique struct ArrowIpcCompressor
+using UniqueCompressor = internal::Unique<struct ArrowIpcCompressor>;
+
/// \brief Class wrapping a unique struct ArrowIpcInputStream
using UniqueInputStream = internal::Unique<struct ArrowIpcInputStream>;