Simplify get_record_batch_message_builder

Author: Hind-M

include/sparrow_ipc/flatbuffer_utils.hpp

@@ -214,12 +214,6 @@ namespace sparrow_ipc
      *
      * @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().
@@ -228,5 +222,5 @@ namespace sparrow_ipc
      * @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, 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);
+    get_record_batch_message_builder(const sparrow::record_batch& record_batch, std::optional<org::apache::arrow::flatbuf::CompressionType> compression = std::nullopt);
 }

src/flatbuffer_utils.cpp

@@ -562,18 +562,20 @@ namespace sparrow_ipc
         return buffers;
     }
 
-    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)
+    flatbuffers::FlatBufferBuilder get_record_batch_message_builder(const sparrow::record_batch& record_batch, std::optional<org::apache::arrow::flatbuf::CompressionType> compression)
     {
-        const std::vector<org::apache::arrow::flatbuf::FieldNode> nodes = create_fieldnodes(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;
+        std::optional<std::vector<org::apache::arrow::flatbuf::Buffer>> compressed_buffers;
         if (compression)
         {
+            compressed_buffers = generate_compressed_buffers(record_batch, compression.value());
             compression_offset = org::apache::arrow::flatbuf::CreateBodyCompression(record_batch_builder, compression.value(), org::apache::arrow::flatbuf::BodyCompressionMethod::BUFFER);
         }
+        const auto& buffers = compressed_buffers ? *compressed_buffers : get_buffers(record_batch);
+        const std::vector<org::apache::arrow::flatbuf::FieldNode> nodes = create_fieldnodes(record_batch);
+        auto nodes_offset = record_batch_builder.CreateVectorOfStructs(nodes);
+        auto buffers_offset = record_batch_builder.CreateVectorOfStructs(buffers);
         const auto record_batch_offset = org::apache::arrow::flatbuf::CreateRecordBatch(
             record_batch_builder,
             static_cast<int64_t>(record_batch.nb_rows()),
@@ -583,7 +585,7 @@ namespace sparrow_ipc
             0   // TODO :variadic buffer Counts
         );
 
-        const int64_t body_size = body_size_override.value_or(calculate_body_size(record_batch, compression));
+        const int64_t body_size = 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

@@ -25,17 +25,7 @@ namespace sparrow_ipc
 
     void serialize_record_batch(const sparrow::record_batch& record_batch, any_output_stream& stream, std::optional<org::apache::arrow::flatbuf::CompressionType> compression)
     {
-        if (compression.has_value())
-        {
-            // TODO Handle this inside get_record_batch_message_builder
-            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, std::nullopt, nullptr), stream);
-        }
+        common_serialize(get_record_batch_message_builder(record_batch, compression), stream);
         generate_body(record_batch, stream, compression);
     }
 }