Simplify

Author: Hind-M

include/sparrow_ipc/deserialize_primitive_array.hpp

@@ -38,6 +38,7 @@ namespace sparrow_ipc
         const auto compression = record_batch.compression();
         std::vector<std::vector<std::uint8_t>> decompressed_buffers;
 
+        // TODO do not decompress validity buffers?
         auto validity_buffer_span = utils::get_and_decompress_buffer(record_batch, body, buffer_index, compression, decompressed_buffers);
 
         const auto [bitmap_ptr, null_count] = utils::get_bitmap_pointer_and_null_count(validity_buffer_span, record_batch.length());

include/sparrow_ipc/serialize_utils.hpp

@@ -128,13 +128,11 @@ namespace sparrow_ipc
      *
      * @param record_batch The record batch to serialize.
      * @param compression_type The compression algorithm to use (e.g., LZ4_FRAME, ZSTD).
-     * @return A std::pair containing:
-     *         - first: A vector of bytes representing the complete compressed message body.
-     *         - second: A vector of FlatBuffer Buffer objects describing the offset and
-     *                   size of each buffer within the compressed body.
+     * @return A vector of FlatBuffer Buffer objects describing the offset and
+     *         size of each buffer within the compressed body.
      */
-    [[nodiscard]] SPARROW_IPC_API std::pair<std::vector<uint8_t>, std::vector<org::apache::arrow::flatbuf::Buffer>>
-    generate_compressed_body_and_buffers(const sparrow::record_batch& record_batch, const org::apache::arrow::flatbuf::CompressionType compression_type);
+    [[nodiscard]] SPARROW_IPC_API std::vector<org::apache::arrow::flatbuf::Buffer>
+    generate_compressed_buffers(const sparrow::record_batch& record_batch, const org::apache::arrow::flatbuf::CompressionType compression_type);
 
     /**
      * @brief Fills the body vector with serialized data from an arrow proxy and its children.
@@ -150,8 +148,9 @@ namespace sparrow_ipc
      *
      * @param arrow_proxy The arrow proxy containing buffers and potential child proxies to serialize
      * @param stream The output stream where the serialized body data will be written
+     * @param compression The compression type to use when serializing
      */
-    SPARROW_IPC_API void fill_body(const sparrow::arrow_proxy& arrow_proxy, any_output_stream& stream);
+    SPARROW_IPC_API void fill_body(const sparrow::arrow_proxy& arrow_proxy, any_output_stream& stream, std::optional<org::apache::arrow::flatbuf::CompressionType> compression = std::nullopt);
 
     /**
      * @brief Generates a serialized body from a record batch.
@@ -162,8 +161,9 @@ namespace sparrow_ipc
      *
      * @param record_batch The record batch containing columns to be serialized
      * @param stream The output stream where the serialized body will be written
+     * @param compression The compression type to use when serializing
      */
-    SPARROW_IPC_API void generate_body(const sparrow::record_batch& record_batch, any_output_stream& stream);
+    SPARROW_IPC_API void generate_body(const sparrow::record_batch& record_batch, any_output_stream& stream, std::optional<org::apache::arrow::flatbuf::CompressionType> compression = std::nullopt);
 
     /**
      * @brief Calculates the total size of the body section for an Arrow array.
@@ -173,9 +173,10 @@ namespace sparrow_ipc
      * buffer size is aligned to 8-byte boundaries as required by the Arrow format.
      *
      * @param arrow_proxy The Arrow array proxy containing buffers and child arrays
+     * @param compression The compression type to use when serializing
      * @return int64_t The total aligned size in bytes of all buffers in the array hierarchy
      */
-    [[nodiscard]] SPARROW_IPC_API int64_t calculate_body_size(const sparrow::arrow_proxy& arrow_proxy);
+    [[nodiscard]] SPARROW_IPC_API int64_t calculate_body_size(const sparrow::arrow_proxy& arrow_proxy, std::optional<org::apache::arrow::flatbuf::CompressionType> compression = std::nullopt);
 
     /**
      * @brief Calculates the total body size of a record batch by summing the body sizes of all its columns.
@@ -185,9 +186,10 @@ namespace sparrow_ipc
      * the total memory required for the serialized data content of the record batch.
      *
      * @param record_batch The sparrow record batch containing columns to calculate size for
+     * @param compression The compression type to use when serializing
      * @return int64_t The total body size in bytes of all columns in the record batch
      */
-    [[nodiscard]] SPARROW_IPC_API int64_t calculate_body_size(const sparrow::record_batch& record_batch);
+    [[nodiscard]] SPARROW_IPC_API int64_t calculate_body_size(const sparrow::record_batch& record_batch, std::optional<org::apache::arrow::flatbuf::CompressionType> compression = std::nullopt);
 
     SPARROW_IPC_API std::vector<sparrow::data_type> get_column_dtypes(const sparrow::record_batch& rb);
 }

src/flatbuffer_utils.cpp

@@ -583,7 +583,7 @@ namespace sparrow_ipc
             0   // TODO :variadic buffer Counts
         );
 
-        const int64_t body_size = body_size_override.value_or(calculate_body_size(record_batch));
+        const int64_t body_size = body_size_override.value_or(calculate_body_size(record_batch, compression));
         const auto record_batch_message_offset = org::apache::arrow::flatbuf::CreateMessage(
             record_batch_builder,
             org::apache::arrow::flatbuf::MetadataVersion::V5,

src/serialize.cpp

@@ -28,15 +28,14 @@ namespace sparrow_ipc
         if (compression.has_value())
         {
             // TODO Handle this inside get_record_batch_message_builder
-            auto [compressed_body, compressed_buffers] = generate_compressed_body_and_buffers(record_batch, compression.value());
-            common_serialize(get_record_batch_message_builder(record_batch, compression, compressed_body.size(), &compressed_buffers), stream);
-            // TODO Use something equivalent to generate_body (stream wise, handling children etc)
-            stream.write(std::span(compressed_body.data(), compressed_body.size()));
+            auto compressed_buffers = generate_compressed_buffers(record_batch, compression.value());
+            auto body_size_override = calculate_body_size(record_batch, compression);
+            common_serialize(get_record_batch_message_builder(record_batch, compression, body_size_override, &compressed_buffers), stream);
         }
         else
         {
-            common_serialize(get_record_batch_message_builder(record_batch, compression), stream);
-            generate_body(record_batch, stream);
+            common_serialize(get_record_batch_message_builder(record_batch, compression, std::nullopt, nullptr), stream);
         }
+        generate_body(record_batch, stream, compression);
     }
 }

src/serialize_utils.cpp

@@ -7,52 +7,69 @@
 
 namespace sparrow_ipc
 {
-    void fill_body(const sparrow::arrow_proxy& arrow_proxy, any_output_stream& stream)
+    void fill_body(const sparrow::arrow_proxy& arrow_proxy, any_output_stream& stream, std::optional<org::apache::arrow::flatbuf::CompressionType> compression)
     {
-        for (const auto& buffer : arrow_proxy.buffers())
-        {
-            stream.write(buffer);
+        std::for_each(arrow_proxy.buffers().begin(), arrow_proxy.buffers().end(), [&](const auto& buffer) {
+            if (compression.has_value())
+            {
+                auto compressed_buffer_with_header = compress(compression.value(), std::span<const uint8_t>(buffer.data(), buffer.size()));
+                stream.write(std::span(compressed_buffer_with_header.data(), compressed_buffer_with_header.size()));
+            }
+            else
+            {
+                stream.write(buffer);
+            }
             stream.add_padding();
-        }
-        for (const auto& child : arrow_proxy.children())
-        {
-            fill_body(child, stream);
-        }
+        });
+
+        std::for_each(arrow_proxy.children().begin(), arrow_proxy.children().end(), [&](const auto& child) {
+            fill_body(child, stream, compression);
+        });
     }
 
-    void generate_body(const sparrow::record_batch& record_batch, any_output_stream& stream)
+    void generate_body(const sparrow::record_batch& record_batch, any_output_stream& stream, std::optional<org::apache::arrow::flatbuf::CompressionType> compression)
     {
-        for (const auto& column : record_batch.columns())
-        {
+        std::for_each(record_batch.columns().begin(), record_batch.columns().end(), [&](const auto& column) {
             const auto& arrow_proxy = sparrow::detail::array_access::get_arrow_proxy(column);
-            fill_body(arrow_proxy, stream);
-        }
+            fill_body(arrow_proxy, stream, compression);
+        });
     }
 
-    int64_t calculate_body_size(const sparrow::arrow_proxy& arrow_proxy)
+    int64_t calculate_body_size(const sparrow::arrow_proxy& arrow_proxy, std::optional<org::apache::arrow::flatbuf::CompressionType> compression)
     {
         int64_t total_size = 0;
-        for (const auto& buffer : arrow_proxy.buffers())
+        if (compression.has_value())
+        {
+            for (const auto& buffer : arrow_proxy.buffers())
+            {
+                total_size += utils::align_to_8(compress(compression.value(), std::span<const uint8_t>(buffer.data(), buffer.size())).size());
+            }
+        }
+        else
         {
-            total_size += utils::align_to_8(buffer.size());
+            for (const auto& buffer : arrow_proxy.buffers())
+            {
+                total_size += utils::align_to_8(buffer.size());
+            }
         }
+
         for (const auto& child : arrow_proxy.children())
         {
-            total_size += calculate_body_size(child);
+            total_size += calculate_body_size(child, compression);
         }
         return total_size;
     }
 
-    int64_t calculate_body_size(const sparrow::record_batch& record_batch)
+    int64_t calculate_body_size(const sparrow::record_batch& record_batch, std::optional<org::apache::arrow::flatbuf::CompressionType> compression)
     {
         return std::accumulate(
             record_batch.columns().begin(),
             record_batch.columns().end(),
             int64_t{0},
-            [](int64_t acc, const sparrow::array& arr)
+            [&](int64_t acc, const sparrow::array& arr)
             {
                 const auto& arrow_proxy = sparrow::detail::array_access::get_arrow_proxy(arr);
-                return acc + calculate_body_size(arrow_proxy);
+                return acc + calculate_body_size(arrow_proxy, compression);
             }
         );
     }
@@ -78,18 +95,7 @@ namespace sparrow_ipc
         flatbuffers::FlatBufferBuilder record_batch_builder = get_record_batch_message_builder(record_batch, compression);
         const flatbuffers::uoffset_t record_batch_len = record_batch_builder.GetSize();
 
-        std::size_t actual_body_size = 0;
-        if (compression.has_value())
-        {
-            // If compressed, the body size is the sum of compressed buffer sizes + original size prefixes + padding
-            auto [compressed_body, compressed_buffers] = generate_compressed_body_and_buffers(record_batch, compression.value());
-            actual_body_size = compressed_body.size();
-        }
-        else
-        {
-            // If not compressed, the body size is the sum of uncompressed buffer sizes with padding
-            actual_body_size = static_cast<std::size_t>(calculate_body_size(record_batch));
-        }
+        const std::size_t actual_body_size = static_cast<std::size_t>(calculate_body_size(record_batch, compression));
 
         // Calculate total size:
         // - Continuation bytes (4)
@@ -103,10 +109,9 @@ namespace sparrow_ipc
         return metadata_size + actual_body_size;
     }
 
-    std::pair<std::vector<uint8_t>, std::vector<org::apache::arrow::flatbuf::Buffer>>
-    generate_compressed_body_and_buffers(const sparrow::record_batch& record_batch, const org::apache::arrow::flatbuf::CompressionType compression_type)
+    std::vector<org::apache::arrow::flatbuf::Buffer>
+    generate_compressed_buffers(const sparrow::record_batch& record_batch, const org::apache::arrow::flatbuf::CompressionType compression_type)
     {
-        std::vector<uint8_t> compressed_body;
         std::vector<org::apache::arrow::flatbuf::Buffer> compressed_buffers;
         int64_t current_offset = 0;
 
@@ -115,24 +120,13 @@ namespace sparrow_ipc
             const auto& arrow_proxy = sparrow::detail::array_access::get_arrow_proxy(column);
             for (const auto& buffer : arrow_proxy.buffers())
             {
-                // Compress the buffer. The returned buffer already has the correct size header.
                 std::vector<uint8_t> compressed_buffer_with_header = compress(compression_type, std::span<const uint8_t>(buffer.data(), buffer.size()));
-
                 const size_t aligned_chunk_size = utils::align_to_8(compressed_buffer_with_header.size());
-                const size_t padding_needed = aligned_chunk_size - compressed_buffer_with_header.size();
-
-                // Write compressed data with header
-                compressed_body.insert(compressed_body.end(), compressed_buffer_with_header.begin(), compressed_buffer_with_header.end());
-
-                // Add padding
-                compressed_body.insert(compressed_body.end(), padding_needed, 0);
-
-                // Update compressed buffer metadata
                 compressed_buffers.emplace_back(current_offset, aligned_chunk_size);
                 current_offset += aligned_chunk_size;
             }
         }
-        return {compressed_body, compressed_buffers};
+        return compressed_buffers;
     }
 
     std::vector<sparrow::data_type> get_column_dtypes(const sparrow::record_batch& rb)

tests/test_serialize_utils.cpp

@@ -40,20 +40,43 @@ namespace sparrow_ipc
             }
         }
 
+        // TODO after the used fcts are stable regarding compression, add tests for fcts having it as an additional argument
+        // cf. fill_body example
         TEST_CASE("fill_body")
         {
-            SUBCASE("Simple primitive array")
+            SUBCASE("Simple primitive array (uncompressed)")
             {
                 auto array = sp::primitive_array<int32_t>({1, 2, 3, 4, 5});
                 auto proxy = sp::detail::array_access::get_arrow_proxy(array);
                 std::vector<uint8_t> body;
                 sparrow_ipc::memory_output_stream stream(body);
                 sparrow_ipc::any_output_stream astream(stream);
-                fill_body(proxy, astream);
+                fill_body(proxy, astream, std::nullopt);
                 CHECK_GT(body.size(), 0);
                 // Body size should be aligned
                 CHECK_EQ(body.size() % 8, 0);
             }
+
+            SUBCASE("Simple primitive array (compressible)")
+            {
+                std::vector<int32_t> data(1000, 12345); // Repeating values, should be very compressible
+                auto array = sp::primitive_array<int32_t>(data);
+                auto proxy = sp::detail::array_access::get_arrow_proxy(array);
+
+                // Compressed
+                std::vector<uint8_t> body_compressed;
+                sparrow_ipc::memory_output_stream stream_compressed(body_compressed);
+                sparrow_ipc::any_output_stream astream_compressed(stream_compressed);
+                fill_body(proxy, astream_compressed, org::apache::arrow::flatbuf::CompressionType::LZ4_FRAME);
+
+                // Uncompressed
+                std::vector<uint8_t> body_uncompressed;
+                sparrow_ipc::memory_output_stream stream_uncompressed(body_uncompressed);
+                sparrow_ipc::any_output_stream astream_uncompressed(stream_uncompressed);
+                fill_body(proxy, astream_uncompressed, std::nullopt);
+                // Check that compressed size is smaller than uncompressed size
+                CHECK_LT(body_compressed.size(), body_uncompressed.size());
+            }
         }
 
         TEST_CASE("generate_body")