Reorganize code

Author: Hind-M

include/sparrow_ipc/flatbuffer_utils.hpp

@@ -207,6 +207,48 @@ namespace sparrow_ipc
     [[nodiscard]] std::vector<org::apache::arrow::flatbuf::Buffer>
     get_buffers(const sparrow::record_batch& record_batch);
 
+    /**
+     * @brief Generates the compressed message body and buffer metadata for a record batch.
+     *
+     * 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.
+     *
+     * @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.
+     */
+    [[nodiscard]] SPARROW_IPC_API std::vector<org::apache::arrow::flatbuf::Buffer>
+    generate_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.
+     *
+     * This function recursively computes the total size needed for all buffers
+     * in an Arrow array structure, including buffers from child arrays. Each
+     * 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, 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.
+     *
+     * This function iterates through all columns in the given record batch and accumulates
+     * the body size of each column's underlying Arrow array proxy. The body size represents
+     * 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, std::optional<CompressionType> compression = std::nullopt);
+
     /**
      * @brief Creates a FlatBuffer message containing a serialized Apache Arrow RecordBatch.
      *

include/sparrow_ipc/serialize_utils.hpp

@@ -119,22 +119,6 @@ namespace sparrow_ipc
         return total_size;
     }
 
-    /**
-     * @brief Generates the compressed message body and buffer metadata for a record batch.
-     *
-     * 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.
-     *
-     * @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.
-     */
-    [[nodiscard]] SPARROW_IPC_API std::vector<org::apache::arrow::flatbuf::Buffer>
-    generate_compressed_buffers(const sparrow::record_batch& record_batch, const CompressionType compression_type);
-
     /**
      * @brief Fills the body vector with serialized data from an arrow proxy and its children.
      *
@@ -166,31 +150,5 @@ namespace sparrow_ipc
      */
     SPARROW_IPC_API void generate_body(const sparrow::record_batch& record_batch, any_output_stream& stream, std::optional<CompressionType> compression = std::nullopt);
 
-    /**
-     * @brief Calculates the total size of the body section for an Arrow array.
-     *
-     * This function recursively computes the total size needed for all buffers
-     * in an Arrow array structure, including buffers from child arrays. Each
-     * 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, 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.
-     *
-     * This function iterates through all columns in the given record batch and accumulates
-     * the body size of each column's underlying Arrow array proxy. The body size represents
-     * 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, std::optional<CompressionType> compression = std::nullopt);
-
     SPARROW_IPC_API std::vector<sparrow::data_type> get_column_dtypes(const sparrow::record_batch& rb);
 }

src/flatbuffer_utils.cpp

@@ -2,7 +2,6 @@
 
 #include "compression_impl.hpp"
 #include "sparrow_ipc/flatbuffer_utils.hpp"
-#include "sparrow_ipc/serialize_utils.hpp"
 #include "sparrow_ipc/utils.hpp"
 
 namespace sparrow_ipc
@@ -563,6 +562,65 @@ namespace sparrow_ipc
         return buffers;
     }
 
+    std::vector<org::apache::arrow::flatbuf::Buffer>
+    generate_compressed_buffers(const sparrow::record_batch& record_batch, const CompressionType compression_type)
+    {
+        std::vector<org::apache::arrow::flatbuf::Buffer> compressed_buffers;
+        int64_t current_offset = 0;
+
+        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;
+    }
+
+    int64_t calculate_body_size(const sparrow::arrow_proxy& arrow_proxy, std::optional<CompressionType> compression)
+    {
+        int64_t total_size = 0;
+        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
+        {
+            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, compression);
+        }
+        return total_size;
+    }
+
+    int64_t calculate_body_size(const sparrow::record_batch& record_batch, std::optional<CompressionType> compression)
+    {
+        return std::accumulate(
+            record_batch.columns().begin(),
+            record_batch.columns().end(),
+            int64_t{0},
+            [&](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, compression);
+            }
+        );
+    }
+
     flatbuffers::FlatBufferBuilder get_record_batch_message_builder(const sparrow::record_batch& record_batch, std::optional<CompressionType> compression)
     {
         flatbuffers::FlatBufferBuilder record_batch_builder;

src/serialize_utils.cpp

@@ -34,45 +34,6 @@ namespace sparrow_ipc
         });
     }
 
-    int64_t calculate_body_size(const sparrow::arrow_proxy& arrow_proxy, std::optional<CompressionType> compression)
-    {
-        int64_t total_size = 0;
-        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
-        {
-            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, compression);
-        }
-        return total_size;
-    }
-
-    int64_t calculate_body_size(const sparrow::record_batch& record_batch, std::optional<CompressionType> compression)
-    {
-        return std::accumulate(
-            record_batch.columns().begin(),
-            record_batch.columns().end(),
-            int64_t{0},
-            [&](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, compression);
-            }
-        );
-    }
-
     std::size_t calculate_schema_message_size(const sparrow::record_batch& record_batch)
     {
         // Build the schema message to get its exact size
@@ -108,26 +69,6 @@ namespace sparrow_ipc
         return metadata_size + actual_body_size;
     }
 
-    std::vector<org::apache::arrow::flatbuf::Buffer>
-    generate_compressed_buffers(const sparrow::record_batch& record_batch, const CompressionType compression_type)
-    {
-        std::vector<org::apache::arrow::flatbuf::Buffer> compressed_buffers;
-        int64_t current_offset = 0;
-
-        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<sparrow::data_type> get_column_dtypes(const sparrow::record_batch& rb)
     {
         std::vector<sparrow::data_type> dtypes;

tests/test_flatbuffer_utils.cpp

@@ -532,4 +532,4 @@ namespace sparrow_ipc
             }
         }
     }
-}
+}

tests/test_serialize_utils.cpp

@@ -4,6 +4,7 @@
 #include <sparrow.hpp>
 
 #include "sparrow_ipc/any_output_stream.hpp"
+#include "sparrow_ipc/flatbuffer_utils.hpp"
 #include "sparrow_ipc/magic_values.hpp"
 #include "sparrow_ipc/memory_output_stream.hpp"
 #include "sparrow_ipc/serialize_utils.hpp"