Fix conflicts and rework serialization with compression

Author: Hind-M

CMakeLists.txt

@@ -109,7 +109,6 @@ set(SPARROW_IPC_HEADERS
     ${SPARROW_IPC_INCLUDE_DIR}/sparrow_ipc/config/config.hpp
     ${SPARROW_IPC_INCLUDE_DIR}/sparrow_ipc/config/sparrow_ipc_version.hpp
     ${SPARROW_IPC_INCLUDE_DIR}/sparrow_ipc/compression.hpp
-    ${SPARROW_IPC_INCLUDE_DIR}/sparrow_ipc/deserialize_variable_size_binary_array.hpp
     ${SPARROW_IPC_INCLUDE_DIR}/sparrow_ipc/deserialize_fixedsizebinary_array.hpp
     ${SPARROW_IPC_INCLUDE_DIR}/sparrow_ipc/deserialize_primitive_array.hpp
     ${SPARROW_IPC_INCLUDE_DIR}/sparrow_ipc/deserialize_utils.hpp

include/sparrow_ipc/chunk_memory_serializer.hpp

@@ -1,7 +1,16 @@
 #pragma once
 
+#include <numeric>
+#include <optional>
+#include <ranges>
+#include <stdexcept>
+#include <vector>
+
 #include <sparrow/record_batch.hpp>
 
+#include "Message_generated.h"
+
+#include "sparrow_ipc/any_output_stream.hpp"
 #include "sparrow_ipc/chunk_memory_output_stream.hpp"
 #include "sparrow_ipc/config/config.hpp"
 #include "sparrow_ipc/memory_output_stream.hpp"
@@ -33,8 +42,9 @@ namespace sparrow_ipc
          * @brief Constructs a chunk serializer with a reference to a chunked memory output stream.
          *
          * @param stream Reference to a chunked memory output stream that will receive the serialized chunks
+         * @param compression Optional: The compression type to use for record batch bodies.
          */
-        chunk_serializer(chunked_memory_output_stream<std::vector<std::vector<uint8_t>>>& stream);
+        chunk_serializer(chunked_memory_output_stream<std::vector<std::vector<uint8_t>>>& stream, std::optional<org::apache::arrow::flatbuf::CompressionType> compression = std::nullopt);
 
         /**
          * @brief Writes a single record batch to the chunked stream.
@@ -120,6 +130,7 @@ namespace sparrow_ipc
         std::vector<sparrow::data_type> m_dtypes;
         chunked_memory_output_stream<std::vector<std::vector<uint8_t>>>* m_pstream;
         bool m_ended{false};
+        std::optional<org::apache::arrow::flatbuf::CompressionType> m_compression;
     };
 
     // Implementation
@@ -133,7 +144,21 @@ namespace sparrow_ipc
             throw std::runtime_error("Cannot append record batches to a serializer that has been ended");
         }
 
-        m_pstream->reserve((m_schema_received ? 0 : 1) + m_pstream->size() + record_batches.size());
+        const auto reserve_function = [&record_batches, this]()
+        {
+            return std::accumulate(
+                        record_batches.begin(),
+                        record_batches.end(),
+                        m_pstream->size(),
+                        [this](size_t acc, const sparrow::record_batch& rb)
+                        {
+                            return acc + calculate_record_batch_message_size(rb, m_compression);
+                        }
+                    )
+                    + (m_schema_received ? 0 : calculate_schema_message_size(*record_batches.begin()));
+        };
+
+        m_pstream->reserve(reserve_function);
 
         if (!m_schema_received)
         {
@@ -148,10 +173,14 @@ namespace sparrow_ipc
 
         for (const auto& rb : record_batches)
         {
+            if (get_column_dtypes(rb) != m_dtypes)
+            {
+                throw std::invalid_argument("Record batch schema does not match serializer schema");
+            }
             std::vector<uint8_t> buffer;
             memory_output_stream stream(buffer);
             any_output_stream astream(stream);
-            serialize_record_batch(rb, astream);
+            serialize_record_batch(rb, astream, m_compression);
             m_pstream->write(std::move(buffer));
         }
     }
@@ -169,4 +198,4 @@ namespace sparrow_ipc
         write(record_batches);
         return *this;
     }
-}
+}

include/sparrow_ipc/flatbuffer_utils.hpp

@@ -213,15 +213,20 @@ namespace sparrow_ipc
      * format that conforms to the Arrow IPC specification.
      *
      * @param record_batch The source record batch containing the data to be serialized
-     *
+     * @param compression Optional: The compression algorithm to be used for the message body
+     * @param body_size Optional: An override for the total size of the message body
+     *                  If not provided, the size is calculated from the uncompressed buffers
+     *                  This is required when using compression
+     * @param compressed_buffers Optional: A pointer to a vector of buffer metadata.
+     *                           If provided, this metadata is used instead of generating it from the
+     *                           uncompressed record batch. This is required when using compression.
      * @return A FlatBufferBuilder containing the complete serialized message ready for
      *         transmission or storage. The builder is finished and ready to be accessed
      *         via GetBufferPointer() and GetSize().
      *
      * @note The returned message uses Arrow IPC format version V5
-     * @note Compression and variadic buffer counts are not currently implemented (set to 0)
-     * @note The body size is automatically calculated based on the record batch contents
+     * @note Variadic buffer counts is not currently implemented (set to 0)
      */
     [[nodiscard]] flatbuffers::FlatBufferBuilder
-    get_record_batch_message_builder(const sparrow::record_batch& record_batch);
-}
+    get_record_batch_message_builder(const sparrow::record_batch& record_batch, std::optional<org::apache::arrow::flatbuf::CompressionType> compression = std::nullopt, std::optional<std::int64_t> body_size = std::nullopt, const std::vector<org::apache::arrow::flatbuf::Buffer>* compressed_buffers = nullptr);
+}

include/sparrow_ipc/serialize.hpp

@@ -25,9 +25,8 @@ namespace sparrow_ipc
      * @tparam R Container type that holds record batches (must support empty(), operator[], begin(), end())
      * @param record_batches Collection of record batches to serialize. All batches must have identical
      * schemas.
-     * @param compression The compression type to use when serializing
-     *
      * @param stream The output stream where the serialized data will be written.
+     * @param compression The compression type to use when serializing.
      *
      * @throws std::invalid_argument If record batches have inconsistent schemas or if the collection
      *                               contains batches that cannot be serialized together.
@@ -70,13 +69,14 @@ namespace sparrow_ipc
      *
      * @param record_batch The sparrow record batch to serialize
      * @param stream The output stream where the serialized record batch will be written
+     * @param compression The compression type to use when serializing.
      *
      * @note The output follows Arrow IPC message format with proper alignment and
      *       includes both metadata and data portions of the record batch
      */
 
     SPARROW_IPC_API void
-    serialize_record_batch(const sparrow::record_batch& record_batch, any_output_stream& stream);
+    serialize_record_batch(const sparrow::record_batch& record_batch, any_output_stream& stream, std::optional<org::apache::arrow::flatbuf::CompressionType> compression);
 
     /**
      * @brief Serializes a schema message for a record batch into a byte buffer.

include/sparrow_ipc/serialize_utils.hpp

@@ -8,7 +8,6 @@
 #include "Message_generated.h"
 #include "sparrow_ipc/any_output_stream.hpp"
 #include "sparrow_ipc/config/config.hpp"
-#include "sparrow_ipc/compression.hpp"
 #include "sparrow_ipc/utils.hpp"
 
 namespace sparrow_ipc
@@ -40,8 +39,8 @@ namespace sparrow_ipc
      * consists of a metadata section followed by a body section containing the actual data.
      *
      * @param record_batch The sparrow record batch to be serialized
-     * @param compression The compression type to use when serializing
      * @param stream The output stream where the serialized record batch will be written
+     * @param compression The compression type to use when serializing
      */
     SPARROW_IPC_API void
     serialize_record_batch(const sparrow::record_batch& record_batch, any_output_stream& stream, std::optional<org::apache::arrow::flatbuf::CompressionType> compression);
@@ -74,10 +73,11 @@ namespace sparrow_ipc
      * - Body data with 8-byte alignment between buffers
      *
      * @param record_batch The record batch to be measured
+     * @param compression The compression type to use when serializing
      * @return The total size in bytes that the serialized record batch would occupy
      */
     [[nodiscard]] SPARROW_IPC_API std::size_t
-    calculate_record_batch_message_size(const sparrow::record_batch& record_batch);
+    calculate_record_batch_message_size(const sparrow::record_batch& record_batch, std::optional<org::apache::arrow::flatbuf::CompressionType> compression = std::nullopt);
 
     /**
      * @brief Calculates the total serialized size for a collection of record batches.
@@ -87,12 +87,13 @@ namespace sparrow_ipc
      *
      * @tparam R Range type containing sparrow::record_batch objects
      * @param record_batches Collection of record batches to be measured
+     * @param compression The compression type to use when serializing
      * @return The total size in bytes for the complete serialized output
      * @throws std::invalid_argument if record batches have inconsistent schemas
      */
     template <std::ranges::input_range R>
         requires std::same_as<std::ranges::range_value_t<R>, sparrow::record_batch>
-    [[nodiscard]] std::size_t calculate_total_serialized_size(const R& record_batches)
+    [[nodiscard]] std::size_t calculate_total_serialized_size(const R& record_batches, std::optional<org::apache::arrow::flatbuf::CompressionType> compression = std::nullopt)
     {
         if (record_batches.empty())
         {
@@ -111,7 +112,7 @@ namespace sparrow_ipc
         // Calculate record batch message sizes
         for (const auto& record_batch : record_batches)
         {
-            total_size += calculate_record_batch_message_size(record_batch);
+            total_size += calculate_record_batch_message_size(record_batch, compression);
         }
 
         return total_size;

include/sparrow_ipc/serializer.hpp

@@ -41,8 +41,8 @@ namespace sparrow_ipc
          *               The serializer stores a pointer to this stream for later use.
          */
         template <writable_stream TStream>
-        serializer(TStream& stream)
-            : m_stream(stream)
+        serializer(TStream& stream, std::optional<org::apache::arrow::flatbuf::CompressionType> compression = std::nullopt)
+            : m_stream(stream), m_compression(compression)
         {
         }
 
@@ -94,7 +94,7 @@ namespace sparrow_ipc
                            m_stream.size(),
                            [this](size_t acc, const sparrow::record_batch& rb)
                            {
-                               return acc + calculate_record_batch_message_size(rb);
+                               return acc + calculate_record_batch_message_size(rb, m_compression);
                            }
                        )
                        + (m_schema_received ? 0 : calculate_schema_message_size(*record_batches.begin()));
@@ -115,7 +115,7 @@ namespace sparrow_ipc
                 {
                     throw std::invalid_argument("Record batch schema does not match serializer schema");
                 }
-                serialize_record_batch(rb, m_stream);
+                serialize_record_batch(rb, m_stream, m_compression);
             }
         }
 
@@ -206,6 +206,7 @@ namespace sparrow_ipc
         std::vector<sparrow::data_type> m_dtypes;
         any_output_stream m_stream;
         bool m_ended{false};
+        std::optional<org::apache::arrow::flatbuf::CompressionType> m_compression;
     };
 
     inline serializer& end_stream(serializer& serializer)

src/chunk_memory_serializer.cpp

@@ -6,8 +6,8 @@
 
 namespace sparrow_ipc
 {
-    chunk_serializer::chunk_serializer(chunked_memory_output_stream<std::vector<std::vector<uint8_t>>>& stream)
-        : m_pstream(&stream)
+    chunk_serializer::chunk_serializer(chunked_memory_output_stream<std::vector<std::vector<uint8_t>>>& stream, std::optional<org::apache::arrow::flatbuf::CompressionType> compression)
+        : m_pstream(&stream), m_compression(compression)
     {
     }
 

src/flatbuffer_utils.cpp

@@ -562,23 +562,28 @@ namespace sparrow_ipc
         return buffers;
     }
 
-    flatbuffers::FlatBufferBuilder get_record_batch_message_builder(const sparrow::record_batch& record_batch)
+    flatbuffers::FlatBufferBuilder get_record_batch_message_builder(const sparrow::record_batch& record_batch, std::optional<org::apache::arrow::flatbuf::CompressionType> compression, std::optional<std::int64_t> body_size_override, const std::vector<org::apache::arrow::flatbuf::Buffer>* compressed_buffers)
     {
         const std::vector<org::apache::arrow::flatbuf::FieldNode> nodes = create_fieldnodes(record_batch);
-        const std::vector<org::apache::arrow::flatbuf::Buffer> buffers = get_buffers(record_batch);
+        const std::vector<org::apache::arrow::flatbuf::Buffer>& buffers = compressed_buffers ? *compressed_buffers : get_buffers(record_batch);
         flatbuffers::FlatBufferBuilder record_batch_builder;
         auto nodes_offset = record_batch_builder.CreateVectorOfStructs(nodes);
         auto buffers_offset = record_batch_builder.CreateVectorOfStructs(buffers);
+        flatbuffers::Offset<org::apache::arrow::flatbuf::BodyCompression> compression_offset = 0;
+        if (compression)
+        {
+            compression_offset = org::apache::arrow::flatbuf::CreateBodyCompression(record_batch_builder, compression.value(), org::apache::arrow::flatbuf::BodyCompressionMethod::BUFFER);
+        }
         const auto record_batch_offset = org::apache::arrow::flatbuf::CreateRecordBatch(
             record_batch_builder,
             static_cast<int64_t>(record_batch.nb_rows()),
             nodes_offset,
             buffers_offset,
-            0,  // TODO: Compression
+            compression_offset,
             0   // TODO :variadic buffer Counts
         );
 
-        const int64_t body_size = calculate_body_size(record_batch);
+        const int64_t body_size = body_size_override.value_or(calculate_body_size(record_batch));
         const auto record_batch_message_offset = org::apache::arrow::flatbuf::CreateMessage(
             record_batch_builder,
             org::apache::arrow::flatbuf::MetadataVersion::V5,

src/serialize.cpp

@@ -1,11 +1,11 @@
-#include "sparrow_ipc/serialize.hpp"
+#include <optional>
 
+#include "sparrow_ipc/serialize.hpp"
 #include "sparrow_ipc/flatbuffer_utils.hpp"
 
 namespace sparrow_ipc
 {
     void common_serialize(
-        const sparrow::record_batch& record_batch,
         const flatbuffers::FlatBufferBuilder& builder,
         any_output_stream& stream
     )
@@ -20,12 +20,23 @@ namespace sparrow_ipc
 
     void serialize_schema_message(const sparrow::record_batch& record_batch, any_output_stream& stream)
     {
-        common_serialize(record_batch, get_schema_message_builder(record_batch), stream);
+        common_serialize(get_schema_message_builder(record_batch), stream);
     }
 
-    void serialize_record_batch(const sparrow::record_batch& record_batch, any_output_stream& stream)
+    void serialize_record_batch(const sparrow::record_batch& record_batch, any_output_stream& stream, std::optional<org::apache::arrow::flatbuf::CompressionType> compression)
     {
-        common_serialize(record_batch, get_record_batch_message_builder(record_batch), stream);
-        generate_body(record_batch, stream);
+        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()));
+        }
+        else
+        {
+            common_serialize(get_record_batch_message_builder(record_batch, compression), stream);
+            generate_body(record_batch, stream);
+        }
     }
-}
+}