Factorize

Author: Hind-M

include/sparrow_ipc/arrow_interface/arrow_array.hpp

@@ -1,43 +1,86 @@
 #pragma once
 
-#include <vector>
+#include <utility>
 
 #include <sparrow/c_interface.hpp>
+#include <sparrow/utils/contracts.hpp>
 
 #include "sparrow_ipc/config/config.hpp"
+#include "sparrow_ipc/arrow_interface/arrow_array/private_data.hpp"
 
 namespace sparrow_ipc
 {
-    [[nodiscard]] SPARROW_IPC_API ArrowArray make_owning_arrow_array(
-        int64_t length,
-        int64_t null_count,
-        int64_t offset,
-        std::vector<std::vector<std::uint8_t>>&& buffers,
-        size_t children_count,
-        ArrowArray** children,
-        ArrowArray* dictionary
-    );
+    SPARROW_IPC_API void release_arrow_array_children_and_dictionary(ArrowArray* array);
+
+    template <ArrowPrivateData T>
+    void arrow_array_release(ArrowArray* array)
+    {
+        SPARROW_ASSERT_TRUE(array != nullptr)
+        SPARROW_ASSERT_TRUE(array->release == std::addressof(arrow_array_release<T>))
+
+        SPARROW_ASSERT_TRUE(array->private_data != nullptr);
 
-    [[nodiscard]] SPARROW_IPC_API ArrowArray make_non_owning_arrow_array(
+        delete static_cast<T*>(array->private_data);
+        array->private_data = nullptr;
+        array->buffers = nullptr;  // The buffers were deleted with the private data
+
+        release_arrow_array_children_and_dictionary(array);
+        array->release = nullptr;
+    }
+
+    template <ArrowPrivateData T, typename Arg>
+    SPARROW_IPC_API void fill_arrow_array(
+        ArrowArray& array,
         int64_t length,
         int64_t null_count,
         int64_t offset,
-        std::vector<std::uint8_t*>&& buffers,
         size_t children_count,
         ArrowArray** children,
-        ArrowArray* dictionary
-    );
+        ArrowArray* dictionary,
+        Arg&& private_data_arg
+    )
+    {
+        SPARROW_ASSERT_TRUE(length >= 0);
+        SPARROW_ASSERT_TRUE(null_count >= -1);
+        SPARROW_ASSERT_TRUE(offset >= 0);
 
-    SPARROW_IPC_API void release_non_owning_arrow_array(ArrowArray* array);
+        array.length = length;
+        array.null_count = null_count;
+        array.offset = offset;
+        array.n_children = static_cast<int64_t>(children_count);
+        array.children = children;
+        array.dictionary = dictionary;
 
-    SPARROW_IPC_API void fill_non_owning_arrow_array(
-        ArrowArray& array,
+        auto private_data = new T(std::forward<Arg>(private_data_arg));
+        array.private_data = private_data;
+        array.n_buffers = private_data->n_buffers();
+        array.buffers = private_data->buffers_ptrs();
+
+        array.release = &arrow_array_release<T>;
+    }
+
+    template <ArrowPrivateData T, typename Arg>
+    [[nodiscard]] SPARROW_IPC_API ArrowArray make_arrow_array(
         int64_t length,
         int64_t null_count,
         int64_t offset,
-        std::vector<std::uint8_t*>&& buffers,
         size_t children_count,
         ArrowArray** children,
-        ArrowArray* dictionary
-    );
-}
+        ArrowArray* dictionary,
+        Arg&& private_data_arg
+    )
+    {
+        ArrowArray array{};
+        fill_arrow_array<T>(
+            array,
+            length,
+            null_count,
+            offset,
+            children_count,
+            children,
+            dictionary,
+            std::forward<Arg>(private_data_arg)
+        );
+        return array;
+    }
+}

include/sparrow_ipc/arrow_interface/arrow_array/private_data.hpp

@@ -1,12 +1,19 @@
 #pragma once
-
+#include <concepts>
 #include <cstdint>
 #include <vector>
 
 #include "sparrow_ipc/config/config.hpp"
 
 namespace sparrow_ipc
 {
+    template <typename T>
+    concept ArrowPrivateData = requires(T& t)
+    {
+        { t.buffers_ptrs() } -> std::same_as<const void**>;
+        { t.n_buffers() } -> std::convertible_to<std::size_t>;
+    };
+
     class owning_arrow_array_private_data
     {
     public:
@@ -21,12 +28,12 @@ namespace sparrow_ipc
             }
         }
 
-        const void** buffers_ptrs() noexcept
+        [[nodiscard]] SPARROW_IPC_API const void** buffers_ptrs() noexcept
         {
             return m_buffer_pointers.data();
         }
 
-        std::size_t n_buffers() const noexcept
+        [[nodiscard]] SPARROW_IPC_API std::size_t n_buffers() const noexcept
         {
             return m_buffers.size();
         }
@@ -42,14 +49,22 @@ namespace sparrow_ipc
     public:
 
         explicit constexpr non_owning_arrow_array_private_data(std::vector<std::uint8_t*>&& buffers_pointers)
-            : m_buffers_pointers(std::move(buffers_pointers))
+            : m_buffer_pointers(std::move(buffers_pointers))
         {
         }
 
-        [[nodiscard]] SPARROW_IPC_API const void** buffers_ptrs() noexcept;
+        [[nodiscard]] SPARROW_IPC_API const void** buffers_ptrs() noexcept
+        {
+            return const_cast<const void**>(reinterpret_cast<void**>(m_buffer_pointers.data()));
+        }
+
+        [[nodiscard]] SPARROW_IPC_API std::size_t n_buffers() const noexcept
+        {
+            return m_buffer_pointers.size();
+        }
 
     private:
 
-        std::vector<std::uint8_t*> m_buffers_pointers;
+        std::vector<std::uint8_t*> m_buffer_pointers;
     };
 }

include/sparrow_ipc/deserialize_primitive_array.hpp

@@ -9,7 +9,6 @@
 #include "Message_generated.h"
 #include "sparrow_ipc/arrow_interface/arrow_array.hpp"
 #include "sparrow_ipc/arrow_interface/arrow_schema.hpp"
-#include "sparrow_ipc/compression.hpp"
 #include "sparrow_ipc/deserialize_utils.hpp"
 
 namespace sparrow_ipc
@@ -23,38 +22,6 @@ namespace sparrow_ipc
         size_t& buffer_index
     )
     {
-        const auto compression = record_batch.compression();
-
-        // Validity buffer
-        const auto validity_buffer_metadata = record_batch.buffers()->Get(buffer_index++);
-        auto validity_buffer_span = body.subspan(validity_buffer_metadata->offset(), validity_buffer_metadata->length());
-        std::vector<std::vector<std::uint8_t>> decompressed_buffers;
-        if (compression)
-        {
-            decompressed_buffers.emplace_back(decompress(compression->codec(), validity_buffer_span));
-            validity_buffer_span = decompressed_buffers.back();
-        }
-        auto bitmap_ptr = const_cast<uint8_t*>(validity_buffer_span.data());
-        const sparrow::dynamic_bitset_view<const std::uint8_t> bitmap_view{
-            bitmap_ptr,
-            static_cast<size_t>(record_batch.length())};
-        auto null_count = bitmap_view.null_count();
-        if (validity_buffer_metadata->length() == 0)
-        {
-            bitmap_ptr = nullptr;
-            null_count = 0;
-        }
-
-        // Data buffer
-        const auto primitive_buffer_metadata = record_batch.buffers()->Get(buffer_index++);
-        auto data_buffer_span = body.subspan(primitive_buffer_metadata->offset(), primitive_buffer_metadata->length());
-        if (compression)
-        {
-            decompressed_buffers.emplace_back(decompress(compression->codec(), data_buffer_span));
-            data_buffer_span = decompressed_buffers.back();
-        }
-        auto primitives_ptr = const_cast<uint8_t*>(data_buffer_span.data());
-
         const std::string_view format = data_type_to_format(
             sparrow::detail::get_data_type_from_array<sparrow::primitive_array<T>>::get()
         );
@@ -68,35 +35,54 @@ namespace sparrow_ipc
             nullptr
         );
 
+        const auto compression = record_batch.compression();
+        std::vector<std::vector<std::uint8_t>> decompressed_buffers;
+
+        auto validity_buffer_span = utils::get_and_decompress_buffer(record_batch, body, buffer_index, compression, decompressed_buffers);
+
+        uint8_t* bitmap_ptr = nullptr;
+        int64_t null_count = 0;
+
+        if (validity_buffer_span.size() > 0)
+        {
+            bitmap_ptr = const_cast<uint8_t*>(validity_buffer_span.data());
+            const sparrow::dynamic_bitset_view<const std::uint8_t> bitmap_view{
+                bitmap_ptr,
+                static_cast<size_t>(record_batch.length())};
+            null_count = bitmap_view.null_count();
+        }
+
+        auto data_buffer_span = utils::get_and_decompress_buffer(record_batch, body, buffer_index, compression, decompressed_buffers);
+
         ArrowArray array;
         if (compression)
         {
-            array = make_owning_arrow_array(
+            array = make_arrow_array<owning_arrow_array_private_data>(
                 record_batch.length(),
                 null_count,
                 0,
-                std::move(decompressed_buffers),
                 0,
                 nullptr,
-                nullptr
+                nullptr,
+                std::move(decompressed_buffers)
             );
         }
         else
         {
+            auto primitives_ptr = const_cast<uint8_t*>(data_buffer_span.data());
             std::vector<std::uint8_t*> buffers = {bitmap_ptr, primitives_ptr};
-            array = make_non_owning_arrow_array(
+            array = make_arrow_array<non_owning_arrow_array_private_data>(
                 record_batch.length(),
                 null_count,
                 0,
-                std::move(buffers),
                 0,
                 nullptr,
-                nullptr
+                nullptr,
+                std::move(buffers)
             );
         }
 
         sparrow::arrow_proxy ap{std::move(array), std::move(schema)};
         return sparrow::primitive_array<T>{std::move(ap)};
     }
 }
-

include/sparrow_ipc/deserialize_utils.hpp

@@ -2,12 +2,12 @@
 
 #include <span>
 #include <utility>
+#include <vector>
 
 #include <sparrow/buffer/dynamic_bitset/dynamic_bitset_view.hpp>
 #include <sparrow/u8_buffer.hpp>
 
 #include "Message_generated.h"
-#include "Schema_generated.h"
 
 namespace sparrow_ipc::utils
 {
@@ -33,4 +33,29 @@ namespace sparrow_ipc::utils
         std::span<const uint8_t> body,
         size_t index
     );
-}
+
+    /**
+     * @brief Extracts a buffer from a RecordBatch and decompresses it if necessary.
+     *
+     * This function retrieves a buffer span from the specified index, increments the index,
+     * and applies decompression if specified. If the buffer is decompressed, the new
+     * data is stored in `decompressed_storage` and the returned span will point to this new data.
+     *
+     * @param record_batch The Arrow RecordBatch containing buffer metadata.
+     * @param body The raw buffer data as a byte span.
+     * @param buffer_index The index of the buffer to retrieve. This value is incremented by the function.
+     * @param compression The compression algorithm to use. If nullptr, no decompression is performed.
+     * @param decompressed_storage A vector that will be used to store the data of any decompressed buffers.
+     *
+     * @return A span viewing the resulting buffer data. This will be a view of the original
+     *         `body` if no decompression occurs, or a view of the newly added buffer in
+     *         `decompressed_storage` if decompression occurs.
+     */
+    [[nodiscard]] std::span<const uint8_t> get_and_decompress_buffer(
+        const org::apache::arrow::flatbuf::RecordBatch& record_batch,
+        std::span<const uint8_t> body,
+        size_t& buffer_index,
+        const org::apache::arrow::flatbuf::BodyCompression* compression,
+        std::vector<std::vector<uint8_t>>& decompressed_storage
+    );
+}