Add owning_arrow_array_private_data

Hind-M · Hind-M · commit a1639007c4bb · 2025-10-13T11:46:07.000+02:00
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -103,6 +103,7 @@ 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/decompressed_buffers.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
@@ -124,6 +125,7 @@ set(SPARROW_IPC_SRC
     ${SPARROW_IPC_SOURCE_DIR}/arrow_interface/arrow_schema.cpp
     ${SPARROW_IPC_SOURCE_DIR}/arrow_interface/arrow_schema/private_data.cpp
     ${SPARROW_IPC_SOURCE_DIR}/compression.cpp
+    #${SPARROW_IPC_SOURCE_DIR}/decompressed_buffers.cpp
     ${SPARROW_IPC_SOURCE_DIR}/deserialize_fixedsizebinary_array.cpp
     ${SPARROW_IPC_SOURCE_DIR}/deserialize_utils.cpp
     ${SPARROW_IPC_SOURCE_DIR}/deserialize.cpp
diff --git a/include/sparrow_ipc/arrow_interface/arrow_array.hpp b/include/sparrow_ipc/arrow_interface/arrow_array.hpp
@@ -1,4 +1,3 @@
-
 #pragma once
 
 #include <vector>
@@ -9,6 +8,16 @@
 
 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
+    );
+
     [[nodiscard]] SPARROW_IPC_API ArrowArray make_non_owning_arrow_array(
         int64_t length,
         int64_t null_count,
diff --git a/include/sparrow_ipc/arrow_interface/arrow_array/private_data.hpp b/include/sparrow_ipc/arrow_interface/arrow_array/private_data.hpp
@@ -7,6 +7,36 @@
 
 namespace sparrow_ipc
 {
+    class owning_arrow_array_private_data
+    {
+    public:
+
+        explicit owning_arrow_array_private_data(std::vector<std::vector<std::uint8_t>>&& buffers)
+            : m_buffers(std::move(buffers))
+        {
+            m_buffer_pointers.reserve(m_buffers.size());
+            for (const auto& buffer : m_buffers)
+            {
+                m_buffer_pointers.push_back(buffer.data());
+            }
+        }
+
+        const void** buffers_ptrs() noexcept
+        {
+            return m_buffer_pointers.data();
+        }
+
+        std::size_t n_buffers() const noexcept
+        {
+            return m_buffers.size();
+        }
+
+    private:
+
+        std::vector<std::vector<std::uint8_t>> m_buffers;
+        std::vector<const void*> m_buffer_pointers;
+    };
+
     class non_owning_arrow_array_private_data
     {
     public:
diff --git a/include/sparrow_ipc/deserialize_primitive_array.hpp b/include/sparrow_ipc/deserialize_primitive_array.hpp
@@ -14,25 +14,6 @@
 
 namespace sparrow_ipc
 {
-    namespace
-    {
-        struct DecompressedBuffers
-        {
-            std::vector<uint8_t> validity_buffer;
-            std::vector<uint8_t> data_buffer;
-        };
-
-        void release_decompressed_buffers(ArrowArray* array)
-        {
-            if (array->private_data)
-            {
-                delete static_cast<DecompressedBuffers*>(array->private_data);
-                array->private_data = nullptr;
-            }
-            array->release = nullptr;
-        }
-    }
-
     template <typename T>
     [[nodiscard]] sparrow::primitive_array<T> deserialize_non_owning_primitive_array(
         const org::apache::arrow::flatbuf::RecordBatch& record_batch,
@@ -43,19 +24,15 @@ namespace sparrow_ipc
     )
     {
         const auto compression = record_batch.compression();
-        DecompressedBuffers* decompressed_buffers_owner = nullptr;
 
         // 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)
         {
-            if (!decompressed_buffers_owner)
-            {
-                decompressed_buffers_owner = new DecompressedBuffers();
-            }
-            decompressed_buffers_owner->validity_buffer = decompress(compression->codec(), validity_buffer_span);
-            validity_buffer_span = decompressed_buffers_owner->validity_buffer;
+            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{
@@ -73,12 +50,8 @@ namespace sparrow_ipc
         auto data_buffer_span = body.subspan(primitive_buffer_metadata->offset(), primitive_buffer_metadata->length());
         if (compression)
         {
-            if (!decompressed_buffers_owner)
-            {
-                decompressed_buffers_owner = new DecompressedBuffers();
-            }
-            decompressed_buffers_owner->data_buffer = decompress(compression->codec(), data_buffer_span);
-            data_buffer_span = decompressed_buffers_owner->data_buffer;
+            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());
 
@@ -95,24 +68,35 @@ namespace sparrow_ipc
             nullptr
         );
 
-        std::vector<std::uint8_t*> buffers = {bitmap_ptr, primitives_ptr};
-        ArrowArray array = make_non_owning_arrow_array(
-            record_batch.length(),
-            null_count,
-            0,
-            std::move(buffers),
-            0,
-            nullptr,
-            nullptr
-        );
-
-        if (decompressed_buffers_owner)
+        ArrowArray array;
+        if (compression)
+        {
+            array = make_owning_arrow_array(
+                record_batch.length(),
+                null_count,
+                0,
+                std::move(decompressed_buffers),
+                0,
+                nullptr,
+                nullptr
+            );
+        }
+        else
         {
-            array.private_data = decompressed_buffers_owner;
-            array.release = release_decompressed_buffers;
+            std::vector<std::uint8_t*> buffers = {bitmap_ptr, primitives_ptr};
+            array = make_non_owning_arrow_array(
+                record_batch.length(),
+                null_count,
+                0,
+                std::move(buffers),
+                0,
+                nullptr,
+                nullptr
+            );
         }
 
         sparrow::arrow_proxy ap{std::move(array), std::move(schema)};
         return sparrow::primitive_array<T>{std::move(ap)};
     }
 }
+
diff --git a/src/arrow_interface/arrow_array.cpp b/src/arrow_interface/arrow_array.cpp
@@ -10,6 +10,64 @@
 
 namespace sparrow_ipc
 {
+    void release_owning_arrow_array(ArrowArray* array)
+    {
+        SPARROW_ASSERT_FALSE(array == nullptr)
+        SPARROW_ASSERT_TRUE(array->release == std::addressof(release_owning_arrow_array))
+
+        if (array->private_data)
+        {
+            delete static_cast<owning_arrow_array_private_data*>(array->private_data);
+            array->private_data = nullptr;
+        }
+
+        for (int64_t i = 0; i < array->n_children; ++i)
+        {
+            if (array->children[i] && array->children[i]->release)
+            {
+                array->children[i]->release(array->children[i]);
+            }
+        }
+        delete[] array->children;
+        array->children = nullptr;
+
+        if (array->dictionary && array->dictionary->release)
+        {
+            array->dictionary->release(array->dictionary);
+        }
+        delete array->dictionary;
+        array->dictionary = nullptr;
+
+        array->release = nullptr;
+    }
+
+    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
+    )
+    {
+        ArrowArray array{};
+        array.length = length;
+        array.null_count = null_count;
+        array.offset = offset;
+
+        auto private_data = new owning_arrow_array_private_data(std::move(buffers));
+        array.private_data = private_data;
+        array.n_buffers = private_data->n_buffers();
+        array.buffers = private_data->buffers_ptrs();
+
+        array.n_children = static_cast<int64_t>(children_count);
+        array.children = children;
+        array.dictionary = dictionary;
+        array.release = release_owning_arrow_array;
+        return array;
+    }
+
     void release_non_owning_arrow_array(ArrowArray* array)
     {
         SPARROW_ASSERT_FALSE(array == nullptr)
