diff --git a/CMakeLists.txt b/CMakeLists.txt index 6335b9029..848b933e8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -247,8 +247,21 @@ if(NANOARROW_IPC) # This could be configurable if shared zstd is a must if(TARGET zstd::libzstd_static) set(NANOARROW_IPC_EXTRA_LIBS zstd::libzstd_static) - else() + elseif(TARGET zstd::libzstd) set(NANOARROW_IPC_EXTRA_LIBS zstd::libzstd) + else() + set(NANOARROW_IPC_EXTRA_LIBS zstd::libzstd_shared) + endif() + + # A static libzstd from a system package manager is usually not built with -fPIC + # and can't be linked into a shared object, so the shared nanoarrow_ipc library + # links the shared libzstd instead + if(TARGET zstd::libzstd_shared) + set(NANOARROW_IPC_EXTRA_LIBS_SHARED zstd::libzstd_shared) + elseif(TARGET zstd::libzstd) + set(NANOARROW_IPC_EXTRA_LIBS_SHARED zstd::libzstd) + else() + set(NANOARROW_IPC_EXTRA_LIBS_SHARED zstd::libzstd_static) endif() endif() @@ -276,6 +289,7 @@ if(NANOARROW_IPC) set(NANOARROW_IPC_EXTRA_FLAGS ${NANOARROW_IPC_EXTRA_FLAGS} "-DNANOARROW_IPC_WITH_LZ4") set(NANOARROW_IPC_EXTRA_LIBS ${NANOARROW_IPC_EXTRA_LIBS} lz4::lz4) + set(NANOARROW_IPC_EXTRA_LIBS_SHARED ${NANOARROW_IPC_EXTRA_LIBS_SHARED} lz4::lz4) endif() if(NOT NANOARROW_BUNDLE) @@ -294,7 +308,7 @@ if(NANOARROW_IPC) add_library(nanoarrow_ipc_shared SHARED ${NANOARROW_IPC_BUILD_SOURCES}) target_link_libraries(nanoarrow_ipc_shared - PRIVATE flatccrt ${NANOARROW_IPC_EXTRA_LIBS} + PRIVATE flatccrt ${NANOARROW_IPC_EXTRA_LIBS_SHARED} PUBLIC nanoarrow_shared nanoarrow_coverage_config) install(FILES src/nanoarrow/nanoarrow_ipc.h src/nanoarrow/nanoarrow_ipc.hpp diff --git a/src/nanoarrow/ipc/codecs.c b/src/nanoarrow/ipc/codecs.c index ab2979893..e5cd0b312 100644 --- a/src/nanoarrow/ipc/codecs.c +++ b/src/nanoarrow/ipc/codecs.c @@ -16,6 +16,7 @@ // under the License. #include +#include #include "nanoarrow/nanoarrow_ipc.h" @@ -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 #include @@ -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 ""; + } +} + +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 76278174a..b6d2c89a7 100644 --- a/src/nanoarrow/ipc/codecs_test.cc +++ b/src/nanoarrow/ipc/codecs_test.cc @@ -16,6 +16,9 @@ // under the License. #include +#include +#include +#include #include #include @@ -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(3); + EXPECT_STREQ(ArrowIpcCompressionTypeToString(unknown_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("", &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(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& input) { + struct ArrowError error {}; + nanoarrow::UniqueBuffer compressed; + + // Content already in dst must be preserved (compress functions only append) + const char* existing = "existing"; + const int64_t existing_size = 8; + EXPECT_EQ(ArrowBufferAppend(compressed.get(), existing, existing_size), NANOARROW_OK); + + EXPECT_EQ(compress({{input.data()}, static_cast(input.size())}, + 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 output(input.size()); + struct ArrowBufferView compressed_view = {{compressed->data + existing_size}, + compressed->size_bytes - existing_size}; + EXPECT_EQ(decompress(compressed_view, output.data(), + static_cast(output.size()), &error), + NANOARROW_OK) + << error.message; + EXPECT_EQ(output, input); + + return compressed->size_bytes - existing_size; +} + +static std::vector CompressibleInput(size_t n) { + std::vector out(n); + for (size_t i = 0; i < n; i++) { + out[i] = static_cast(i % 7); + } + return out; +} + +// Check compress/decompress on empty, small, and multi-block inputs at each level +static void TestCompressionFunctions(ArrowIpcCompressFunction compress, + ArrowIpcDecompressFunction decompress, + const std::vector& 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(kUncompressed012, + kUncompressed012 + sizeof(kUncompressed012))); + + // Large enough to span several blocks; a repetitive input must actually shrink + int64_t compressed_size = TestCompressRoundtrip(compress, decompress, level, input); + EXPECT_LT(compressed_size, static_cast(input.size() / 10)); + 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::min(), std::numeric_limits::min() + 1, + std::numeric_limits::min() + 2}) { + SCOPED_TRACE("compression level " + std::to_string(level)); + // The most negative levels must saturate at maximum acceleration rather than + // overflow and fall back to the default compression level. + EXPECT_EQ(TestCompressRoundtrip(compress, decompress, level, input), + accelerated_size); + } +} + +// An allocator whose reallocate() fails on the fail_on-th call (1-based) and otherwise +// delegates to the default allocator +struct FailingAllocatorState { + int calls; + int fail_on; +}; + +static uint8_t* FailingReallocate(struct ArrowBufferAllocator* allocator, uint8_t* ptr, + int64_t old_size, int64_t new_size) { + auto* state = static_cast(allocator->private_data); + auto default_allocator = ArrowBufferAllocatorDefault(); + if (++state->calls == state->fail_on) { + // nanoarrow discards the buffer on failure, so the old allocation is freed here + default_allocator.free(&default_allocator, ptr, old_size); + return nullptr; + } + return default_allocator.reallocate(&default_allocator, ptr, old_size, new_size); +} + +static void FailingFree(struct ArrowBufferAllocator* allocator, uint8_t* ptr, + int64_t size) { + NANOARROW_UNUSED(allocator); + auto default_allocator = ArrowBufferAllocatorDefault(); + default_allocator.free(&default_allocator, ptr, size); +} + +static struct ArrowBufferAllocator FailingAllocator(FailingAllocatorState* state) { + struct ArrowBufferAllocator allocator = ArrowBufferAllocatorDefault(); + allocator.reallocate = &FailingReallocate; + allocator.free = &FailingFree; + allocator.private_data = state; + return allocator; +} + +TEST(NanoarrowIpcTest, CompressAllocationFailure) { + struct ArrowError error {}; + for (auto compress : + {ArrowIpcGetLZ4CompressionFunction(), ArrowIpcGetZstdCompressionFunction()}) { + if (compress == nullptr) { + continue; + } + + FailingAllocatorState state{0, 1}; + nanoarrow::UniqueBuffer dst; + ASSERT_EQ(ArrowBufferSetAllocator(dst.get(), FailingAllocator(&state)), NANOARROW_OK); + EXPECT_EQ(compress({{kUncompressed012}, sizeof(kUncompressed012)}, + NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, dst.get(), &error), + ENOMEM); + EXPECT_THAT(error.message, ::testing::HasSubstr("ArrowBufferReserve")); + EXPECT_EQ(state.calls, 1); + } +} + +// A stand-in compression function that always fails +static ArrowErrorCode FailCompress(struct ArrowBufferView src, int compression_level, + struct ArrowBuffer* dst, struct ArrowError* error) { + NANOARROW_UNUSED(src); + NANOARROW_UNUSED(compression_level); + NANOARROW_UNUSED(dst); + ArrowErrorSet(error, "FailCompress() failed"); + return EIO; +} + +// A stand-in compression function that records the level it was called with and +// "compresses" by copying +static int last_compression_level = 0; + +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(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(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 b13d0bd69..00dbcf2c7 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 d1435f491..ab7f0b60d 100644 --- a/src/nanoarrow/ipc/encoder.c +++ b/src/nanoarrow/ipc/encoder.c @@ -16,6 +16,7 @@ // under the License. #include +#include #include #include @@ -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 4fa0e4fee..6c79babd7 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 #include +#include #include #include #include @@ -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(bswap64(static_cast(value))); + } + return value; +} + +static void TestCompressedRecordBatchRoundtrip(enum ArrowIpcCompressionType codec) { + nanoarrow::ipc::UniqueEncoder encoder; + ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); + nanoarrow::ipc::UniqueDecoder decoder; + ASSERT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK); + + CompressibleRecordBatch batch; + struct ArrowError error; + ASSERT_EQ(ArrowIpcDecoderSetSchema(decoder.get(), batch.schema(), &error), NANOARROW_OK) + << error.message; + + // Encode without compression for reference + nanoarrow::UniqueBuffer uncompressed_message, uncompressed_body; + ASSERT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), batch.array_view(), + uncompressed_body.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/true, + uncompressed_message.get()), + NANOARROW_OK); + + ASSERT_EQ(ArrowIpcEncoderSetCompression( + encoder.get(), codec, 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 values(1025, 42); + struct ArrowError error; + nanoarrow::UniqueSchema schema; + ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_STRUCT), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaAllocateChildren(schema.get(), 1), NANOARROW_OK); + ASSERT_EQ(ArrowSchemaInitFromType(schema->children[0], NANOARROW_TYPE_INT32), + NANOARROW_OK); + nanoarrow::UniqueArray array; + ASSERT_EQ(ArrowArrayInitFromSchema(array.get(), schema.get(), &error), NANOARROW_OK); + ASSERT_EQ(ArrowArrayStartAppending(array.get()), NANOARROW_OK); + for (int32_t value : values) { + ASSERT_EQ(ArrowArrayAppendInt(array->children[0], value), NANOARROW_OK); + ASSERT_EQ(ArrowArrayFinishElement(array.get()), NANOARROW_OK); + } + ASSERT_EQ(ArrowArrayFinishBuildingDefault(array.get(), &error), NANOARROW_OK); + nanoarrow::UniqueArrayView array_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(array_view.get(), schema.get(), &error), + NANOARROW_OK); + ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), array.get(), &error), NANOARROW_OK); + + int allocations = 0; + auto allocator = ArrowBufferAllocatorDefault(); + allocator.private_data = &allocations; + allocator.reallocate = [](struct ArrowBufferAllocator* allocator, uint8_t* ptr, + int64_t old_size, int64_t new_size) { + ++*static_cast(allocator->private_data); + auto default_allocator = ArrowBufferAllocatorDefault(); + return default_allocator.reallocate(&default_allocator, ptr, old_size, new_size); + }; + nanoarrow::UniqueBuffer body; + ASSERT_EQ(ArrowBufferSetAllocator(body.get(), allocator), NANOARROW_OK); + + ASSERT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), array_view.get(), + body.get(), &error), + NANOARROW_OK) + << error.message; + EXPECT_EQ(allocations, 1); + ASSERT_EQ(body->size_bytes, 4104); + EXPECT_EQ(body->capacity_bytes, 4104); + EXPECT_EQ(std::memcmp(body->data, values.data(), 4100), 0); + EXPECT_EQ(std::memcmp(body->data + 4100, "\0\0\0\0", 4), 0); +} + +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderCompressedRecordBatchLZ4) { + if (ArrowIpcGetLZ4CompressionFunction() == nullptr) { + GTEST_SKIP() << "nanoarrow_ipc not built with NANOARROW_IPC_WITH_LZ4"; + } + 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(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(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 pending; + std::vector destinations; + size_t max_pending = 0; + int adds = 0; + int waits = 0; + int fail_on_add = 0; + bool fail_wait = false; + bool produce_nothing = false; + + static ArrowErrorCode Add(struct ArrowIpcCompressor* compressor, + struct ArrowBufferView src, struct ArrowBuffer* dst, + struct ArrowError* error) { + auto* state = static_cast(compressor->private_data); + if (++state->adds == state->fail_on_add) { + ArrowErrorSet(error, "Deferred add failed"); + return EIO; + } + state->pending.push_back({src, dst}); + state->destinations.push_back(dst); + if (state->pending.size() > state->max_pending) { + state->max_pending = state->pending.size(); + } + return NANOARROW_OK; + } + + static ArrowErrorCode Wait(struct ArrowIpcCompressor* compressor, int64_t timeout_ms, + struct ArrowError* error) { + EXPECT_LT(timeout_ms, 0); + auto* state = static_cast(compressor->private_data); + ++state->waits; + int result = NANOARROW_OK; + for (const auto& job : state->pending) { + if (result == NANOARROW_OK && !state->produce_nothing) { + result = ArrowBufferAppend(job.dst, job.src.data.data, job.src.size_bytes); + } + } + // Complete or cancel every job, including when reporting an error. + state->pending.clear(); + if (state->fail_wait) { + ArrowErrorSet(error, "Deferred wait failed"); + return EIO; + } + return result; + } + + static void Release(struct ArrowIpcCompressor* compressor) { + auto* state = static_cast(compressor->private_data); + state->pending.clear(); + compressor->release = nullptr; + } + + struct ArrowIpcCompressor MakeCompressor() { + struct ArrowIpcCompressor compressor {}; + compressor.compression_type = NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME; + compressor.compress_add = &Add; + compressor.compress_wait = &Wait; + compressor.release = &Release; + compressor.private_data = this; + return compressor; + } +}; + +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderDeferredCompressionAllocationFailure) { + struct ArrowError error; + CompressibleRecordBatch batch; + DeferredCompressor state; + FailingAllocatorState allocator_state{0, 1}; + nanoarrow::ipc::UniqueEncoder encoder; + ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); + auto compressor = state.MakeCompressor(); + ASSERT_EQ(ArrowIpcEncoderSetCompressor(encoder.get(), &compressor), NANOARROW_OK); + + nanoarrow::UniqueBuffer body, message; + ASSERT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), batch.array_view(), + body.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcEncoderFinalizeBuffer(encoder.get(), true, message.get()), + NANOARROW_OK); + EXPECT_TRUE(state.pending.empty()); + EXPECT_GT(state.max_pending, 1); + EXPECT_EQ(state.waits, 1); + + nanoarrow::ipc::UniqueDecoder decoder; + ASSERT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK); + ASSERT_EQ(ArrowIpcDecoderSetSchema(decoder.get(), batch.schema(), &error), + NANOARROW_OK); + struct ArrowBufferView message_view = {{message->data}, message->size_bytes}; + ASSERT_EQ(ArrowIpcDecoderVerifyHeader(decoder.get(), message_view, &error), + NANOARROW_OK); + ASSERT_EQ(ArrowIpcDecoderDecodeHeader(decoder.get(), message_view, &error), + NANOARROW_OK); + nanoarrow::UniqueArray decoded; + ASSERT_EQ( + ArrowIpcDecoderDecodeArray(decoder.get(), {{body->data}, body->size_bytes}, -1, + decoded.get(), NANOARROW_VALIDATION_LEVEL_FULL, &error), + NANOARROW_OK) + << error.message; + nanoarrow::UniqueArrayView decoded_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(decoded_view.get(), batch.schema(), &error), + NANOARROW_OK); + ASSERT_EQ(ArrowArrayViewSetArray(decoded_view.get(), decoded.get(), &error), + NANOARROW_OK); + int is_equal = 0; + ASSERT_EQ(ArrowArrayViewCompare(decoded_view.get(), batch.array_view(), + NANOARROW_COMPARE_IDENTICAL, &is_equal, &error), + NANOARROW_OK); + EXPECT_EQ(is_equal, 1) << error.message; + + // Scratch buffers are reused (by index) for the next message, which is what keeps the + // pointers captured above valid. Make a later prefix allocation fail, after an + // earlier buffer could have been queued with the compressor. + ASSERT_GT(state.destinations.size(), 1); + struct ArrowBuffer* failing_buffer = state.destinations[1]; + ArrowBufferReset(failing_buffer); + ASSERT_EQ(ArrowBufferSetAllocator(failing_buffer, FailingAllocator(&allocator_state)), + NANOARROW_OK); + body->size_bytes = 0; + EXPECT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), batch.array_view(), + body.get(), &error), + ENOMEM); + EXPECT_EQ(allocator_state.calls, 1); + EXPECT_TRUE(state.pending.empty()); +} + +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderDeferredCompressionErrors) { + for (bool fail_wait : {false, true}) { + SCOPED_TRACE(fail_wait ? "wait error" : "add error"); + struct ArrowError error; + CompressibleRecordBatch batch; + DeferredCompressor state; + state.fail_wait = fail_wait; + state.fail_on_add = fail_wait ? 0 : 2; + nanoarrow::ipc::UniqueEncoder encoder; + ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); + auto compressor = state.MakeCompressor(); + ASSERT_EQ(ArrowIpcEncoderSetCompressor(encoder.get(), &compressor), NANOARROW_OK); + nanoarrow::UniqueBuffer body; + EXPECT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), batch.array_view(), + body.get(), &error), + EIO); + EXPECT_STREQ(error.message, + fail_wait ? "Deferred wait failed" : "Deferred add failed"); + EXPECT_TRUE(state.pending.empty()); + EXPECT_EQ(state.waits, 1); + } +} + +// Encode a batch with a body allocator that fails on the fail_on-th allocation, for +// every fail_on until encoding succeeds, so that each allocation site reports ENOMEM +static void TestEncodeAllocationFailures(enum ArrowIpcCompressionType codec) { + 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(buffer->size_bytes); + if (ArrowIpcSystemEndianness() == NANOARROW_IPC_ENDIANNESS_BIG) { + footer_size = static_cast(bswap32(static_cast(footer_size))); + } + ASSERT_EQ(ArrowBufferAppendInt32(buffer.get(), footer_size), NANOARROW_OK); + ASSERT_EQ(ArrowBufferAppend(buffer.get(), "ARROW1", 6), NANOARROW_OK); + + nanoarrow::ipc::UniqueDecoder decoder; + ASSERT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK); + struct ArrowBufferView view = {{buffer->data}, buffer->size_bytes}; + ASSERT_EQ(ArrowIpcDecoderVerifyFooter(decoder.get(), view, &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcDecoderDecodeFooter(decoder.get(), view, &error), NANOARROW_OK) + << error.message; + EXPECT_EQ((decoder->feature_flags & NANOARROW_IPC_FEATURE_COMPRESSED_BODY) != 0, + expected); + }; + ASSERT_NO_FATAL_FAILURE(check_footer(true)); + + // A new schema starts a new file's history on the same encoder. + ASSERT_EQ(ArrowIpcEncoderEncodeSchema(encoder.get(), batch.schema(), &error), + NANOARROW_OK); + message->size_bytes = 0; + ASSERT_EQ( + ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/true, message.get()), + NANOARROW_OK); + ASSERT_NO_FATAL_FAILURE(check_footer(false)); + } +} + +// The scratch buffers for compressed bodies grow when a message has more buffers than +// any encoded before it; the existing ones are moved and stay usable +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderCompressedBuffersGrow) { + struct ArrowError error; + CompressibleRecordBatch batch; + DeferredCompressor state; + nanoarrow::ipc::UniqueEncoder encoder; + ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); + auto compressor = state.MakeCompressor(); + ASSERT_EQ(ArrowIpcEncoderSetCompressor(encoder.get(), &compressor), NANOARROW_OK); + + // A DictionaryBatch of a single int32 column needs two scratch buffers... + nanoarrow::UniqueBuffer body, message; + ASSERT_EQ( + ArrowIpcEncoderEncodeSimpleDictionaryBatch( + encoder.get(), 0, false, batch.array_view()->children[0], body.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ( + ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/true, message.get()), + NANOARROW_OK); + size_t n_small = state.destinations.size(); + EXPECT_GT(n_small, 0); + + // ...and the RecordBatch of all three columns needs more + body->size_bytes = 0; + message->size_bytes = 0; + ASSERT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), batch.array_view(), + body.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ( + ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/true, message.get()), + NANOARROW_OK); + EXPECT_GT(state.destinations.size() - n_small, n_small); + + nanoarrow::ipc::UniqueDecoder decoder; + ASSERT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK); + ASSERT_EQ(ArrowIpcDecoderSetSchema(decoder.get(), batch.schema(), &error), + NANOARROW_OK); + struct ArrowBufferView message_view = {{message->data}, message->size_bytes}; + ASSERT_EQ(ArrowIpcDecoderVerifyHeader(decoder.get(), message_view, &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowIpcDecoderDecodeHeader(decoder.get(), message_view, &error), + NANOARROW_OK) + << error.message; + nanoarrow::UniqueArray decoded; + ASSERT_EQ( + ArrowIpcDecoderDecodeArray(decoder.get(), {{body->data}, body->size_bytes}, -1, + decoded.get(), NANOARROW_VALIDATION_LEVEL_FULL, &error), + NANOARROW_OK) + << error.message; + nanoarrow::UniqueArrayView decoded_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(decoded_view.get(), batch.schema(), &error), + NANOARROW_OK); + ASSERT_EQ(ArrowArrayViewSetArray(decoded_view.get(), decoded.get(), &error), + NANOARROW_OK); + int is_equal = 0; + ASSERT_EQ(ArrowArrayViewCompare(decoded_view.get(), batch.array_view(), + NANOARROW_COMPARE_IDENTICAL, &is_equal, &error), + NANOARROW_OK); + EXPECT_EQ(is_equal, 1) << error.message; +} + +// A compressor that reports success without producing output is an error rather +// than a buffer with a length prefix and no payload +TEST(NanoarrowIpcTest, NanoarrowIpcEncoderCompressorWithoutOutput) { + struct ArrowError error; + CompressibleRecordBatch batch; + DeferredCompressor state; + state.produce_nothing = true; + nanoarrow::ipc::UniqueEncoder encoder; + ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK); + auto compressor = state.MakeCompressor(); + ASSERT_EQ(ArrowIpcEncoderSetCompressor(encoder.get(), &compressor), NANOARROW_OK); + + nanoarrow::UniqueBuffer body; + EXPECT_EQ(ArrowIpcEncoderEncodeSimpleRecordBatch(encoder.get(), batch.array_view(), + body.get(), &error), + EIO); + EXPECT_THAT(error.message, + ::testing::StartsWith("Compressor produced no output for a buffer of")); +} diff --git a/src/nanoarrow/ipc/files_test.cc b/src/nanoarrow/ipc/files_test.cc index 46c7203c1..73fc9724f 100644 --- a/src/nanoarrow/ipc/files_test.cc +++ b/src/nanoarrow/ipc/files_test.cc @@ -18,6 +18,9 @@ #include #include #include +#include +#include +#include #include @@ -212,6 +215,7 @@ class TestFile { ArrowErrorCode WriteNanoarrowStream(const nanoarrow::UniqueSchema& schema, const std::vector& arrays, + enum ArrowIpcCompressionType codec, struct ArrowBuffer* buffer, struct ArrowError* error) { nanoarrow::ipc::UniqueOutputStream output_stream; @@ -219,6 +223,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> codecs = { + {"uncompressed", NANOARROW_IPC_COMPRESSION_TYPE_NONE}}; + if (ArrowIpcGetLZ4CompressionFunction() != nullptr) { + codecs.emplace_back("lz4", NANOARROW_IPC_COMPRESSION_TYPE_LZ4_FRAME); + } + if (ArrowIpcGetZstdCompressionFunction() != nullptr) { + codecs.emplace_back("zstd", NANOARROW_IPC_COMPRESSION_TYPE_ZSTD); + } + + std::vector roundtripped(codecs.size()); if (write_supported_) { - ASSERT_EQ(WriteNanoarrowStream(schema, arrays, roundtripped.get(), &error), - NANOARROW_OK) - << error.message; + for (size_t i = 0; i < codecs.size(); i++) { + SCOPED_TRACE("Write the " + codecs[i].first + " stream using nanoarrow"); + ASSERT_EQ(WriteNanoarrowStream(schema, arrays, codecs[i].second, + roundtripped[i].get(), &error), + NANOARROW_OK) + << error.message; + } } // Read the same file with Arrow C++ @@ -283,28 +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 roundtripped_arrays; - { - SCOPED_TRACE("Read the roundtripped buffer using nanoarrow"); - nanoarrow::UniqueArrayStream array_stream; - ASSERT_EQ(GetArrowArrayStreamIPC(roundtripped.get(), array_stream.get(), &error), - NANOARROW_OK); - ASSERT_EQ(ReadArrowArrayStreamIPC(array_stream.get(), roundtripped_schema.get(), - &roundtripped_arrays, &error), - NANOARROW_OK); - - AssertEqualsTable(std::move(roundtripped_schema), std::move(roundtripped_arrays), - maybe_table_arrow.ValueUnsafe()); + AssertEqualsTable(maybe_table_roundtripped.ValueUnsafe(), + maybe_table_arrow.ValueUnsafe()); + } + + nanoarrow::UniqueSchema roundtripped_schema; + std::vector roundtripped_arrays; + { + SCOPED_TRACE("Read the roundtripped buffer using nanoarrow"); + nanoarrow::UniqueArrayStream array_stream; + ASSERT_EQ( + GetArrowArrayStreamIPC(roundtripped[i].get(), array_stream.get(), &error), + NANOARROW_OK); + ASSERT_EQ(ReadArrowArrayStreamIPC(array_stream.get(), roundtripped_schema.get(), + &roundtripped_arrays, &error), + NANOARROW_OK); + + AssertEqualsTable(std::move(roundtripped_schema), std::move(roundtripped_arrays), + maybe_table_arrow.ValueUnsafe()); + } } } diff --git a/src/nanoarrow/ipc/reader.c b/src/nanoarrow/ipc/reader.c index 0b99b828e..c699ca331 100644 --- a/src/nanoarrow/ipc/reader.c +++ b/src/nanoarrow/ipc/reader.c @@ -390,12 +390,6 @@ static int ArrowIpcArrayStreamReaderReadSchemaIfNeeded( } // ...or if it uses features we don't support - if (private_data->decoder.feature_flags & NANOARROW_IPC_FEATURE_COMPRESSED_BODY) { - ArrowErrorSet(&private_data->error, - "This stream uses unsupported feature COMPRESSED_BODY"); - return EINVAL; - } - if (private_data->decoder.feature_flags & NANOARROW_IPC_FEATURE_DICTIONARY_REPLACEMENT) { ArrowErrorSet(&private_data->error, diff --git a/src/nanoarrow/ipc/writer.c b/src/nanoarrow/ipc/writer.c index 00ea93ecc..bb8c79710 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 c30dd6032..3b08d0c89 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(writer->private_data); + ASSERT_EQ(p->footer.record_batch_blocks.size_bytes, sizeof(struct ArrowIpcFileBlock)); + auto* block = + reinterpret_cast(p->footer.record_batch_blocks.data); + EXPECT_EQ(block->body_length, p->body_buffer.size_bytes); + ASSERT_EQ(ArrowIpcWriterFinalizeFile(writer.get(), &error), NANOARROW_OK) + << error.message; + } +} + +// Read the stream starting at offset back with the array stream reader and check that +// the values match what InitCompressibleBatch() produced +static void CheckCompressibleBatch(const struct ArrowBuffer* output, int64_t offset) { + struct ArrowError error; + + nanoarrow::UniqueBuffer input_buffer; + ASSERT_EQ(ArrowBufferAppend(input_buffer.get(), output->data + offset, + output->size_bytes - offset), + NANOARROW_OK); + nanoarrow::ipc::UniqueInputStream input; + ASSERT_EQ(ArrowIpcInputStreamInitBuffer(input.get(), input_buffer.get()), NANOARROW_OK); + nanoarrow::UniqueArrayStream stream; + ASSERT_EQ(ArrowIpcArrayStreamReaderInit(stream.get(), input.get(), nullptr), + NANOARROW_OK); + + nanoarrow::UniqueSchema schema; + ASSERT_EQ(ArrowArrayStreamGetSchema(stream.get(), schema.get(), &error), NANOARROW_OK) + << error.message; + EXPECT_STREQ(schema->format, "+s"); + + nanoarrow::UniqueArray array; + ASSERT_EQ(ArrowArrayStreamGetNext(stream.get(), array.get(), &error), NANOARROW_OK) + << error.message; + ASSERT_EQ(array->length, kCompressibleBatchLength); + + nanoarrow::UniqueArrayView array_view; + ASSERT_EQ(ArrowArrayViewInitFromSchema(array_view.get(), schema.get(), &error), + NANOARROW_OK) + << error.message; + ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), array.get(), &error), NANOARROW_OK) + << error.message; + for (int64_t i = 0; i < kCompressibleBatchLength; i++) { + ASSERT_EQ(ArrowArrayViewGetIntUnsafe(array_view->children[0], i), i % 8); + } + + nanoarrow::UniqueArray eos; + ASSERT_EQ(ArrowArrayStreamGetNext(stream.get(), eos.get(), &error), NANOARROW_OK) + << error.message; + EXPECT_EQ(eos->release, nullptr); +} + +// 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(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 0523ac38f..8a29cf5fb 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 "" 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 e6b23924f..9629eba31 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; /// \brief Class wrapping a unique struct ArrowIpcDecompressor using UniqueDecompressor = internal::Unique; +/// \brief Class wrapping a unique struct ArrowIpcCompressor +using UniqueCompressor = internal::Unique; + /// \brief Class wrapping a unique struct ArrowIpcInputStream using UniqueInputStream = internal::Unique;