wip

Author: Alex-PLACET

include/sparrow_ipc/serialize.hpp

@@ -7,59 +7,111 @@
 #include <sparrow/record_batch.hpp>
 
 #include "Message_generated.h"
+#include "sparrow_ipc/config/config.hpp"
+#include "sparrow_ipc/utils.hpp"
 
 namespace sparrow_ipc
 {
     template <std::ranges::input_range R>
         requires std::same_as<std::ranges::range_value_t<R>, sparrow::record_batch>
-    std::vector<uint8_t> serialize(const R& record_batches);
+    [[nodiscard]] std::vector<uint8_t> serialize(const R& record_batches);
 
-    std::vector<uint8_t> serialize_schema_message(const ArrowSchema& arrow_schema);
+    [[nodiscard]] SPARROW_IPC_API std::vector<uint8_t>
+    serialize_schema_message(const sparrow::record_batch& record_batch);
 
-    std::vector<uint8_t> serialize_record_batch(const sparrow::record_batch& record_batch);
+    [[nodiscard]] SPARROW_IPC_API std::vector<uint8_t>
+    serialize_record_batch(const sparrow::record_batch& record_batch);
 
     template <std::ranges::input_range R>
         requires std::same_as<std::ranges::range_value_t<R>, sparrow::record_batch>
-    std::vector<uint8_t> serialize_record_batches(const R& record_batches);
-
-    flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<org::apache::arrow::flatbuf::KeyValue>>>
-    create_metadata(flatbuffers::FlatBufferBuilder& builder, const ArrowSchema& arrow_schema);
-
-    ::flatbuffers::Offset<org::apache::arrow::flatbuf::Field>
+    [[nodiscard]] std::vector<uint8_t> serialize_record_batches(const R& record_batches)
+    {
+        std::vector<uint8_t> output;
+        for (const auto& record_batch : record_batches)
+        {
+            const auto rb_serialized = serialize_record_batch(record_batch);
+            output.insert(
+                output.end(),
+                std::make_move_iterator(rb_serialized.begin()),
+                std::make_move_iterator(rb_serialized.end())
+            );
+        }
+        return output;
+    }
+
+    [[nodiscard]] SPARROW_IPC_API
+        flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<org::apache::arrow::flatbuf::KeyValue>>>
+        create_metadata(flatbuffers::FlatBufferBuilder& builder, const ArrowSchema& arrow_schema);
+
+    [[nodiscard]] SPARROW_IPC_API ::flatbuffers::Offset<org::apache::arrow::flatbuf::Field>
     create_field(flatbuffers::FlatBufferBuilder& builder, const ArrowSchema& arrow_schema);
 
-    ::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<org::apache::arrow::flatbuf::Field>>>
+    [[nodiscard]] SPARROW_IPC_API ::flatbuffers::Offset<
+        ::flatbuffers::Vector<::flatbuffers::Offset<org::apache::arrow::flatbuf::Field>>>
+    create_children(flatbuffers::FlatBufferBuilder& builder, sparrow::record_batch::column_range columns);
+
+    [[nodiscard]] SPARROW_IPC_API ::flatbuffers::Offset<
+        ::flatbuffers::Vector<::flatbuffers::Offset<org::apache::arrow::flatbuf::Field>>>
     create_children(flatbuffers::FlatBufferBuilder& builder, const ArrowSchema& arrow_schema);
 
-    flatbuffers::FlatBufferBuilder get_schema_message_builder(const ArrowSchema& arrow_schema);
-    std::vector<uint8_t> serialize_schema_message(const ArrowSchema& arrow_schema);
+    [[nodiscard]] SPARROW_IPC_API flatbuffers::FlatBufferBuilder
+    get_schema_message_builder(const sparrow::record_batch& record_batch);
+
+    [[nodiscard]] SPARROW_IPC_API std::vector<uint8_t>
+    serialize_schema_message(const sparrow::record_batch& record_batch);
 
-    void fill_fieldnodes(
+    SPARROW_IPC_API void fill_fieldnodes(
         const sparrow::arrow_proxy& arrow_proxy,
         std::vector<org::apache::arrow::flatbuf::FieldNode>& nodes
     );
 
-    std::vector<org::apache::arrow::flatbuf::FieldNode>
+    [[nodiscard]] SPARROW_IPC_API std::vector<org::apache::arrow::flatbuf::FieldNode>
     create_fieldnodes(const sparrow::record_batch& record_batch);
 
-    void fill_buffers(
+    SPARROW_IPC_API void fill_buffers(
         const sparrow::arrow_proxy& arrow_proxy,
         std::vector<org::apache::arrow::flatbuf::Buffer>& flatbuf_buffers,
         int64_t& offset
     );
 
-    std::vector<org::apache::arrow::flatbuf::Buffer> get_buffers(const sparrow::record_batch& record_batch);
+    [[nodiscard]] SPARROW_IPC_API std::vector<org::apache::arrow::flatbuf::Buffer>
+    get_buffers(const sparrow::record_batch& record_batch);
 
-    void fill_body(const sparrow::arrow_proxy& arrow_proxy, std::vector<uint8_t>& body);
-    std::vector<uint8_t> generate_body(const sparrow::record_batch& record_batch);
-    int64_t calculate_body_size(const sparrow::arrow_proxy& arrow_proxy);
-    int64_t calculate_body_size(const sparrow::record_batch& record_batch);
+    SPARROW_IPC_API void fill_body(const sparrow::arrow_proxy& arrow_proxy, std::vector<uint8_t>& body);
+    [[nodiscard]] SPARROW_IPC_API std::vector<uint8_t> generate_body(const sparrow::record_batch& record_batch);
+    [[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::record_batch& record_batch);
 
-    flatbuffers::FlatBufferBuilder get_record_batch_message_builder(
+    [[nodiscard]] SPARROW_IPC_API flatbuffers::FlatBufferBuilder get_record_batch_message_builder(
         const sparrow::record_batch& record_batch,
         const std::vector<org::apache::arrow::flatbuf::FieldNode>& nodes,
         const std::vector<org::apache::arrow::flatbuf::Buffer>& buffers
     );
 
-    std::vector<uint8_t> serialize_record_batch(const sparrow::record_batch& record_batch);
+    [[nodiscard]] SPARROW_IPC_API std::vector<uint8_t>
+    serialize_record_batch(const sparrow::record_batch& record_batch);
+
+    template <std::ranges::input_range R>
+        requires std::same_as<std::ranges::range_value_t<R>, sparrow::record_batch>
+    std::vector<uint8_t> serialize(const R& record_batches)
+    {
+        if (record_batches.empty())
+        {
+            return {};
+        }
+        if (utils::check_record_batches_consistency(record_batches))
+        {
+            throw std::invalid_argument(
+                "All record batches must have the same schema to be serialized together."
+            );
+        }
+        std::vector<uint8_t> serialized_schema = serialize_schema_message(record_batches[0]);
+        std::vector<uint8_t> serialized_record_batches = serialize_record_batches(record_batches);
+        serialized_schema.insert(
+            serialized_schema.end(),
+            std::make_move_iterator(serialized_record_batches.begin()),
+            std::make_move_iterator(serialized_record_batches.end())
+        );
+        return serialized_schema;
+    }
 }

include/sparrow_ipc/utils.hpp

@@ -22,7 +22,35 @@ namespace sparrow_ipc::utils
 
     template <std::ranges::input_range R>
         requires std::same_as<std::ranges::range_value_t<R>, sparrow::record_batch>
-    SPARROW_IPC_API bool check_record_batches_consistency(const R& record_batches);
+    bool check_record_batches_consistency(const R& record_batches)
+    {
+        if (record_batches.empty())
+        {
+            return true;
+        }
+        const sparrow::record_batch& first_rb = record_batches[0];
+        for (const sparrow::record_batch& rb : record_batches)
+        {
+            if (rb.nb_columns() != first_rb.nb_columns())
+            {
+                return false;
+            }
+            if (rb.nb_rows() != first_rb.nb_rows())
+            {
+                return false;
+            }
+            for (size_t col_idx = 0; col_idx < rb.nb_columns(); ++col_idx)
+            {
+                const sparrow::array& arr = rb.get_column(col_idx);
+                const sparrow::array& first_arr = first_rb.get_column(col_idx);
+                if (arr.data_type() != first_arr.data_type())
+                {
+                    return false;
+                }
+            }
+        }
+        return true;
+    }
 
-    size_t calculate_output_serialized_size(const sparrow::record_batch& record_batch);
+    // size_t calculate_output_serialized_size(const sparrow::record_batch& record_batch);
 }

src/serialize.cpp

@@ -56,25 +56,38 @@ namespace sparrow_ipc
     {
         std::vector<flatbuffers::Offset<org::apache::arrow::flatbuf::Field>> children_vec;
         children_vec.reserve(arrow_schema.n_children);
-        for (int i = 0; i < arrow_schema.n_children; ++i)
+        for (size_t i = 0; i < arrow_schema.n_children; ++i)
         {
             if (arrow_schema.children[i] == nullptr)
             {
-                throw std::invalid_argument("ArrowSchema has null child at index " + std::to_string(i));
+                throw std::invalid_argument("ArrowSchema has null child pointer");
             }
-            flatbuffers::Offset<org::apache::arrow::flatbuf::Field> field = create_field(
-                builder,
-                *(arrow_schema.children[i])
-            );
+            const auto& child = *arrow_schema.children[i];
+            flatbuffers::Offset<org::apache::arrow::flatbuf::Field> field = create_field(builder, child);
             children_vec.emplace_back(field);
         }
         return children_vec.empty() ? 0 : builder.CreateVector(children_vec);
     }
 
-    flatbuffers::FlatBufferBuilder get_schema_message_builder(const ArrowSchema& arrow_schema)
+    ::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset<org::apache::arrow::flatbuf::Field>>>
+    create_children(flatbuffers::FlatBufferBuilder& builder, sparrow::record_batch::column_range columns)
+    {
+        std::vector<flatbuffers::Offset<org::apache::arrow::flatbuf::Field>> children_vec;
+        children_vec.reserve(columns.size());
+        for (const auto& column : columns)
+        {
+            const auto& arrow_schema = sparrow::detail::array_access::get_arrow_proxy(column).schema();
+            flatbuffers::Offset<org::apache::arrow::flatbuf::Field> field = create_field(builder, arrow_schema);
+            children_vec.emplace_back(field);
+        }
+        return children_vec.empty() ? 0 : builder.CreateVector(children_vec);
+    }
+
+    flatbuffers::FlatBufferBuilder get_schema_message_builder(const sparrow::record_batch& record_batch)
     {
         flatbuffers::FlatBufferBuilder schema_builder;
-        const auto fields_vec = create_children(schema_builder, arrow_schema);
+        record_batch.columns();
+        const auto fields_vec = create_children(schema_builder, record_batch.columns());
         const auto schema_offset = org::apache::arrow::flatbuf::CreateSchema(
             schema_builder,
             org::apache::arrow::flatbuf::Endianness::Little,  // TODO: make configurable
@@ -92,12 +105,11 @@ namespace sparrow_ipc
         return schema_builder;
     }
 
-    std::vector<uint8_t> serialize_schema_message(const ArrowSchema& arrow_schema)
+    std::vector<uint8_t> serialize_schema_message(const sparrow::record_batch& record_batch)
     {
         std::vector<uint8_t> schema_buffer;
-
         schema_buffer.insert(schema_buffer.end(), continuation.begin(), continuation.end());
-        flatbuffers::FlatBufferBuilder schema_builder = get_schema_message_builder(arrow_schema);
+        flatbuffers::FlatBufferBuilder schema_builder = get_schema_message_builder(record_batch);
         const flatbuffers::uoffset_t schema_len = schema_builder.GetSize();
         schema_buffer.reserve(schema_buffer.size() + sizeof(uint32_t) + schema_len);
         // Write the 4-byte length prefix after the continuation bytes
@@ -162,8 +174,7 @@ namespace sparrow_ipc
         }
         for (const auto& child : arrow_proxy.children())
         {
-            const auto& child_arrow_proxy = sparrow::detail::array_access::get_arrow_proxy(child);
-            fill_buffers(child_arrow_proxy, flatbuf_buffers, offset);
+            fill_buffers(child, flatbuf_buffers, offset);
         }
     }
 
@@ -190,8 +201,7 @@ namespace sparrow_ipc
         }
         for (const auto& child : arrow_proxy.children())
         {
-            const auto& child_arrow_proxy = sparrow::detail::array_access::get_arrow_proxy(child);
-            fill_body(child_arrow_proxy, body);
+            fill_body(child, body);
         }
     }
 
@@ -215,8 +225,7 @@ namespace sparrow_ipc
         }
         for (const auto& child : arrow_proxy.children())
         {
-            const auto& child_arrow_proxy = sparrow::detail::array_access::get_arrow_proxy(child);
-            total_size += calculate_body_size(child_arrow_proxy);
+            total_size += calculate_body_size(child);
         }
         return total_size;
     }
@@ -270,14 +279,11 @@ namespace sparrow_ipc
     {
         std::vector<org::apache::arrow::flatbuf::FieldNode> nodes = create_fieldnodes(record_batch);
         std::vector<org::apache::arrow::flatbuf::Buffer> flatbuf_buffers = get_buffers(record_batch);
-        flatbuffers::FlatBufferBuilder record_batch_builder;
-        ::flatbuffers::Offset<org::apache::arrow::flatbuf::RecordBatch>
-            record_batch_offset = org::apache::arrow::flatbuf::CreateRecordBatchDirect(
-                record_batch_builder,
-                static_cast<int64_t>(record_batch.nb_rows()),
-                &nodes,
-                &flatbuf_buffers
-            );
+        flatbuffers::FlatBufferBuilder record_batch_builder = get_record_batch_message_builder(
+            record_batch,
+            nodes,
+            flatbuf_buffers
+        );
         std::vector<uint8_t> output;
         output.insert(output.end(), continuation.begin(), continuation.end());
         const flatbuffers::uoffset_t record_batch_len = record_batch_builder.GetSize();
@@ -302,40 +308,4 @@ namespace sparrow_ipc
         return output;
     }
 
-    template <std::ranges::input_range R>
-        requires std::same_as<std::ranges::range_value_t<R>, sparrow::record_batch>
-    std::vector<uint8_t> serialize_record_batches(const R& record_batches)
-    {
-        std::vector<uint8_t> output;
-        for (const auto& record_batch : record_batches)
-        {
-            const auto rb_serialized = serialize_record_batch(record_batch);
-            output.insert(
-                output.end(),
-                std::make_move_iterator(rb_serialized.begin()),
-                std::make_move_iterator(rb_serialized.end())
-            );
-        }
-        return output;
-    }
-
-    template <std::ranges::input_range R>
-        requires std::same_as<std::ranges::range_value_t<R>, sparrow::record_batch>
-    std::vector<uint8_t> serialize(const R& record_batches)
-    {
-        if (check_record_batches_consistency(record_batches))
-        {
-            throw std::invalid_argument(
-                "All record batches must have the same schema to be serialized together."
-            );
-        }
-        std::vector<uint8_t> serialized_schema = serialize_schema_message(record_batches[0].schema());
-        std::vector<uint8_t> serialized_record_batches = serialize_record_batches(record_batches);
-        serialized_schema.insert(
-            serialized_schema.end(),
-            std::make_move_iterator(serialized_record_batches.begin()),
-            std::make_move_iterator(serialized_record_batches.end())
-        );
-        return serialized_schema;
-    }
 }

src/utils.cpp

@@ -445,39 +445,8 @@ namespace sparrow_ipc
             }
         }
 
-        template <std::ranges::input_range R>
-            requires std::same_as<std::ranges::range_value_t<R>, sparrow::record_batch>
-        bool check_record_batches_consistency(const R& record_batches)
-        {
-            if (record_batches.empty())
-            {
-                return true;
-            }
-            const sparrow::record_batch& first_rb = record_batches[0];
-            for (const sparrow::record_batch& rb : record_batches)
-            {
-                if (rb.nb_columns() != first_rb.nb_columns())
-                {
-                    return false;
-                }
-                if (rb.nb_rows() != first_rb.nb_rows())
-                {
-                    return false;
-                }
-                for (size_t col_idx = 0; col_idx < rb.nb_columns(); ++col_idx)
-                {
-                    const sparrow::array& arr = rb.get_column(col_idx);
-                    const sparrow::array& first_arr = first_rb.get_column(col_idx);
-                    if (arr.data_type() != first_arr.data_type())
-                    {
-                        return false;
-                    }
-                }
-            }
-        }
-
-        size_t calculate_output_serialized_size(const sparrow::record_batch& record_batch)
-        {
-        }
+        // size_t calculate_output_serialized_size(const sparrow::record_batch& record_batch)
+        // {
+        // }
     }
 }