Refactor and add some tests

Author: Hind-M

include/sparrow_ipc/flatbuffer_utils.hpp

@@ -6,6 +6,7 @@
 #include <sparrow/record_batch.hpp>
 
 #include "sparrow_ipc/compression.hpp"
+#include "sparrow_ipc/utils.hpp"
 
 namespace sparrow_ipc
 {
@@ -166,6 +167,42 @@ namespace sparrow_ipc
     [[nodiscard]] std::vector<org::apache::arrow::flatbuf::FieldNode>
     create_fieldnodes(const sparrow::record_batch& record_batch);
 
+    namespace details
+    {
+        template <typename Func>
+        void fill_buffers_impl(
+            const sparrow::arrow_proxy& arrow_proxy,
+            std::vector<org::apache::arrow::flatbuf::Buffer>& flatbuf_buffers,
+            int64_t& offset,
+            Func&& get_buffer_size
+        )
+        {
+            const auto& buffers = arrow_proxy.buffers();
+            for (const auto& buffer : buffers)
+            {
+                int64_t size = get_buffer_size(buffer);
+                flatbuf_buffers.emplace_back(offset, size);
+                offset += utils::align_to_8(size);
+            }
+            for (const auto& child : arrow_proxy.children())
+            {
+                fill_buffers_impl(child, flatbuf_buffers, offset, get_buffer_size);
+            }
+        }
+
+        template <typename Func>
+        std::vector<org::apache::arrow::flatbuf::Buffer> get_buffers_impl(const sparrow::record_batch& record_batch, Func&& fill_buffers_func)
+        {
+            std::vector<org::apache::arrow::flatbuf::Buffer> buffers;
+            int64_t offset = 0;
+            for (const auto& column : record_batch.columns())
+            {
+                const auto& arrow_proxy = sparrow::detail::array_access::get_arrow_proxy(column);
+                fill_buffers_func(arrow_proxy, buffers, offset);
+            }
+            return buffers;
+        }
+    } // namespace details
 
     /**
      * @brief Recursively fills a vector of FlatBuffer Buffer objects with buffer information from an Arrow
@@ -208,20 +245,39 @@ namespace sparrow_ipc
     get_buffers(const sparrow::record_batch& record_batch);
 
     /**
-     * @brief Generates the compressed message body and buffer metadata for a record batch.
+     * @brief Recursively populates a vector with compressed buffer metadata from an Arrow proxy.
      *
-     * This function traverses the record batch, compresses each buffer using the specified
-     * compression algorithm, and constructs the message body. For each compressed buffer,
-     * it is prefixed by its 8-byte uncompressed size. Padding is added after each
-     * compressed buffer to ensure 8-byte alignment.
+     * This function traverses the Arrow proxy and its children, compressing each buffer and recording
+     * its metadata (offset and size) in the provided vector. The offset is updated to ensure proper
+     * alignment for each subsequent buffer.
      *
-     * @param record_batch The record batch to serialize.
-     * @param compression_type The compression algorithm to use (e.g., LZ4_FRAME, ZSTD).
-     * @return A vector of FlatBuffer Buffer objects describing the offset and
-     *         size of each buffer within the compressed body.
+     * @param arrow_proxy The Arrow proxy containing the buffers to be compressed.
+     * @param flatbuf_compressed_buffers A vector to store the resulting compressed buffer metadata.
+     * @param offset The current offset in the buffer layout, which will be updated by the function.
+     * @param compression_type The compression algorithm to use.
      */
-    [[nodiscard]] SPARROW_IPC_API std::vector<org::apache::arrow::flatbuf::Buffer>
-    generate_compressed_buffers(const sparrow::record_batch& record_batch, const CompressionType compression_type);
+    void fill_compressed_buffers(
+        const sparrow::arrow_proxy& arrow_proxy,
+        std::vector<org::apache::arrow::flatbuf::Buffer>& flatbuf_compressed_buffers,
+        int64_t& offset,
+        const CompressionType compression_type
+    );
+
+    /**
+     * @brief Retrieves metadata describing the layout of compressed buffers within a record batch.
+     *
+     * This function processes a record batch to determine the metadata (offset and size)
+     * for each of its buffers, assuming they are compressed using the specified algorithm.
+     * This metadata accounts for each compressed buffer being prefixed by its 8-byte
+     * uncompressed size and padded to ensure 8-byte alignment.
+     *
+     * @param record_batch The record batch whose buffers' compressed metadata is to be retrieved.
+     * @param compression_type The compression algorithm that would be applied (e.g., LZ4_FRAME, ZSTD).
+     * @return A vector of FlatBuffer Buffer objects, each describing the offset and
+     *         size of a corresponding compressed buffer within a larger message body.
+     */
+    [[nodiscard]] std::vector<org::apache::arrow::flatbuf::Buffer>
+    get_compressed_buffers(const sparrow::record_batch& record_batch, const CompressionType compression_type);
 
     /**
      * @brief Calculates the total size of the body section for an Arrow array.
@@ -234,7 +290,7 @@ namespace sparrow_ipc
      * @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, std::optional<CompressionType> compression = std::nullopt);
+    [[nodiscard]] int64_t calculate_body_size(const sparrow::arrow_proxy& arrow_proxy, std::optional<CompressionType> compression = std::nullopt);
 
     /**
      * @brief Calculates the total body size of a record batch by summing the body sizes of all its columns.
@@ -247,7 +303,7 @@ namespace sparrow_ipc
      * @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, std::optional<CompressionType> compression = std::nullopt);
+    [[nodiscard]] int64_t calculate_body_size(const sparrow::record_batch& record_batch, std::optional<CompressionType> compression = std::nullopt);
 
     /**
      * @brief Creates a FlatBuffer message containing a serialized Apache Arrow RecordBatch.

src/flatbuffer_utils.cpp

@@ -2,7 +2,6 @@
 
 #include "compression_impl.hpp"
 #include "sparrow_ipc/flatbuffer_utils.hpp"
-#include "sparrow_ipc/utils.hpp"
 
 namespace sparrow_ipc
 {
@@ -537,49 +536,38 @@ namespace sparrow_ipc
         int64_t& offset
     )
     {
-        const auto& buffers = arrow_proxy.buffers();
-        for (const auto& buffer : buffers)
-        {
-            int64_t size = static_cast<int64_t>(buffer.size());
-            flatbuf_buffers.emplace_back(offset, size);
-            offset += utils::align_to_8(size);
-        }
-        for (const auto& child : arrow_proxy.children())
-        {
-            fill_buffers(child, flatbuf_buffers, offset);
-        }
+        details::fill_buffers_impl(arrow_proxy, flatbuf_buffers, offset, [](const auto& buffer) {
+            return static_cast<int64_t>(buffer.size());
+        });
     }
 
     std::vector<org::apache::arrow::flatbuf::Buffer> get_buffers(const sparrow::record_batch& record_batch)
     {
-        std::vector<org::apache::arrow::flatbuf::Buffer> buffers;
-        std::int64_t offset = 0;
-        for (const auto& column : record_batch.columns())
-        {
-            const auto& arrow_proxy = sparrow::detail::array_access::get_arrow_proxy(column);
-            fill_buffers(arrow_proxy, buffers, offset);
-        }
-        return buffers;
+        return details::get_buffers_impl(record_batch, [](const sparrow::arrow_proxy& proxy, std::vector<org::apache::arrow::flatbuf::Buffer>& buffers, int64_t& offset) {
+            fill_buffers(proxy, buffers, offset);
+        });
     }
 
-    std::vector<org::apache::arrow::flatbuf::Buffer>
-    generate_compressed_buffers(const sparrow::record_batch& record_batch, const CompressionType compression_type)
+    void fill_compressed_buffers(
+        const sparrow::arrow_proxy& arrow_proxy,
+        std::vector<org::apache::arrow::flatbuf::Buffer>& flatbuf_compressed_buffers,
+        int64_t& offset,
+        const CompressionType compression_type
+    )
     {
-        std::vector<org::apache::arrow::flatbuf::Buffer> compressed_buffers;
-        int64_t current_offset = 0;
+        details::fill_buffers_impl(
+            arrow_proxy, flatbuf_compressed_buffers, offset, [&](const auto& buffer) {
+                return compress(compression_type, std::span<const uint8_t>(buffer.data(), buffer.size()))
+                    .size();
+            });
+    }
 
-        for (const auto& column : record_batch.columns())
-        {
-            const auto& arrow_proxy = sparrow::detail::array_access::get_arrow_proxy(column);
-            for (const auto& buffer : arrow_proxy.buffers())
-            {
-                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());
-                compressed_buffers.emplace_back(current_offset, aligned_chunk_size);
-                current_offset += aligned_chunk_size;
-            }
-        }
-        return compressed_buffers;
+    std::vector<org::apache::arrow::flatbuf::Buffer>
+    get_compressed_buffers(const sparrow::record_batch& record_batch, const CompressionType compression_type)
+    {
+        return details::get_buffers_impl(record_batch, [&](const sparrow::arrow_proxy& proxy, std::vector<org::apache::arrow::flatbuf::Buffer>& buffers, int64_t& offset) {
+            fill_compressed_buffers(proxy, buffers, offset, compression_type);
+        });
     }
 
     int64_t calculate_body_size(const sparrow::arrow_proxy& arrow_proxy, std::optional<CompressionType> compression)
@@ -628,7 +616,7 @@ namespace sparrow_ipc
         std::optional<std::vector<org::apache::arrow::flatbuf::Buffer>> compressed_buffers;
         if (compression)
         {
-            compressed_buffers = generate_compressed_buffers(record_batch, compression.value());
+            compressed_buffers = get_compressed_buffers(record_batch, compression.value());
             compression_offset = org::apache::arrow::flatbuf::CreateBodyCompression(record_batch_builder, details::to_fb_compression_type(compression.value()), org::apache::arrow::flatbuf::BodyCompressionMethod::BUFFER);
         }
         const auto& buffers = compressed_buffers ? *compressed_buffers : get_buffers(record_batch);

tests/test_flatbuffer_utils.cpp

@@ -178,40 +178,75 @@ namespace sparrow_ipc
             }
         }
 
+        void test_fill_buffers_variant(
+            const std::function<void(const sparrow::arrow_proxy&, std::vector<org::apache::arrow::flatbuf::Buffer>&, int64_t&)>& fill_func)
+        {
+            auto array = sp::primitive_array<int32_t>({1, 2, 3, 4, 5});
+            auto proxy = sp::detail::array_access::get_arrow_proxy(array);
+
+            std::vector<org::apache::arrow::flatbuf::Buffer> buffers;
+            int64_t offset = 0;
+            fill_func(proxy, buffers, offset);
+
+            CHECK_GT(buffers.size(), 0);
+            CHECK_GT(offset, 0);
+
+            // Verify offsets are aligned
+            for (const auto& buffer : buffers)
+            {
+                CHECK_EQ(buffer.offset() % 8, 0);
+            }
+        }
+
         TEST_CASE("fill_buffers")
         {
             SUBCASE("Simple primitive array")
             {
-                auto array = sp::primitive_array<int32_t>({1, 2, 3, 4, 5});
-                auto proxy = sp::detail::array_access::get_arrow_proxy(array);
-
-                std::vector<org::apache::arrow::flatbuf::Buffer> buffers;
-                int64_t offset = 0;
-                fill_buffers(proxy, buffers, offset);
+                test_fill_buffers_variant([](const sparrow::arrow_proxy& proxy, std::vector<org::apache::arrow::flatbuf::Buffer>& buffers, int64_t& offset) {
+                    fill_buffers(proxy, buffers, offset);
+                });
+            }
+        }
 
-                CHECK_GT(buffers.size(), 0);
-                CHECK_GT(offset, 0);
+        TEST_CASE("fill_compressed_buffers")
+        {
+            SUBCASE("Simple primitive array")
+            {
+                test_fill_buffers_variant([](const sparrow::arrow_proxy& proxy, std::vector<org::apache::arrow::flatbuf::Buffer>& buffers, int64_t& offset) {
+                    fill_compressed_buffers(proxy, buffers, offset, CompressionType::LZ4_FRAME);
+                });
+            }
+        }
 
-                // Verify offsets are aligned
-                for (const auto& buffer : buffers)
-                {
-                    CHECK_EQ(buffer.offset() % 8, 0);
-                }
+        void test_get_buffers_variant(const std::function<std::vector<org::apache::arrow::flatbuf::Buffer>(const sparrow::record_batch&)>& get_func)
+        {
+            auto record_batch = create_test_record_batch();
+            auto buffers = get_func(record_batch);
+            CHECK_GT(buffers.size(), 0);
+            // Verify all offsets are properly calculated and aligned
+            for (size_t i = 1; i < buffers.size(); ++i)
+            {
+                CHECK_GE(buffers[i].offset(), buffers[i - 1].offset() + buffers[i - 1].length());
             }
         }
 
         TEST_CASE("get_buffers")
         {
             SUBCASE("Record batch with multiple columns")
             {
-                auto record_batch = create_test_record_batch();
-                auto buffers = get_buffers(record_batch);
-                CHECK_GT(buffers.size(), 0);
-                // Verify all offsets are properly calculated and aligned
-                for (size_t i = 1; i < buffers.size(); ++i)
-                {
-                    CHECK_GE(buffers[i].offset(), buffers[i - 1].offset() + buffers[i - 1].length());
-                }
+                test_get_buffers_variant([](const sparrow::record_batch& record_batch) {
+                    return get_buffers(record_batch);
+                });
+            }
+        }
+
+        TEST_CASE("get_compressed_buffers")
+        {
+            SUBCASE("Record batch with multiple columns")
+            {
+                test_get_buffers_variant([](const sparrow::record_batch& record_batch) {
+                    return get_compressed_buffers(record_batch, CompressionType::LZ4_FRAME);
+                });
             }
         }
 

tests/test_serialize_utils.cpp

@@ -69,12 +69,18 @@ namespace sparrow_ipc
                 sparrow_ipc::memory_output_stream stream_compressed(body_compressed);
                 sparrow_ipc::any_output_stream astream_compressed(stream_compressed);
                 fill_body(proxy, astream_compressed, CompressionType::LZ4_FRAME);
+                CHECK_GT(body_compressed.size(), 0);
+                // Body size should be aligned
+                CHECK_EQ(body_compressed.size() % 8, 0);
 
                 // 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_GT(body_uncompressed.size(), 0);
+                // Body size should be aligned
+                CHECK_EQ(body_uncompressed.size() % 8, 0);
                 // Check that compressed size is smaller than uncompressed size
                 CHECK_LT(body_compressed.size(), body_uncompressed.size());
             }