Add a multithreaded alternative to to_arrow

Author: XanthosXanthopoulos

apis/python/src/tiledbsoma/common.cc

@@ -195,13 +195,18 @@ py::object _buffer_to_table(std::shared_ptr<ArrayBuffers> buffers) {
     py::list array_list;
     py::list field_list;
 
-    for (auto& name : buffers->names()) {
-        auto [pa_array, pa_schema] = ArrowAdapter::to_arrow(buffers->at<ReadColumnBuffer>(name));
+    auto column_names = buffers->names();
+
+    auto arrays = ArrowAdapter::buffer_to_arrow(buffers);
+
+    for (size_t i = 0; i < column_names.size(); ++i) {
+        auto& [pa_array, pa_schema] = arrays[i];
         auto nullable = (pa_schema->flags & ARROW_FLAG_NULLABLE) != 0;
         auto dtype = pa_dtype_import(py::capsule(pa_schema.get()));
         array_list.append(pa_array_import(py::capsule(pa_array.get()), dtype));
-        field_list.append(pa.attr("field")(name, dtype, nullable));
+        field_list.append(pa.attr("field")(column_names[i], dtype, nullable));
     }
+
     return pa_table_from_arrays(array_list, "schema"_a = pa.attr("schema")(field_list));
 }
 

libtiledbsoma/src/utils/arrow_adapter.cc

@@ -12,7 +12,9 @@
  */
 
 #include <ranges>
+#include <thread>
 
+#include "../soma/array_buffers.h"
 #include "../soma/column_buffer.h"
 #include "arrow_adapter.h"
 #include "logger.h"
@@ -756,6 +758,28 @@ inline void exitIfError(const ArrowErrorCode ec, const std::string& msg) {
         throw TileDBSOMAError(fmt::format("ArrowAdapter: Arrow Error {} ", msg));
 }
 
+std::vector<std::pair<managed_unique_ptr<ArrowArray>, managed_unique_ptr<ArrowSchema>>> ArrowAdapter::buffer_to_arrow(
+    std::shared_ptr<ArrayBuffers> buffer, bool downcast_dict_of_large_var) {
+    std::vector<std::future<std::pair<managed_unique_ptr<ArrowArray>, managed_unique_ptr<ArrowSchema>>>> arrow_columns;
+
+    for (const auto& name : buffer->names()) {
+        arrow_columns.emplace_back(
+            std::async(
+                std::launch::async,
+                ArrowAdapter::to_arrow,
+                buffer->at<ReadColumnBuffer>(name),
+                downcast_dict_of_large_var));
+    }
+
+    std::vector<std::pair<managed_unique_ptr<ArrowArray>, managed_unique_ptr<ArrowSchema>>> columns;
+
+    for (auto& arrow_column : arrow_columns) {
+        columns.push_back(std::move(arrow_column.get()));
+    }
+
+    return columns;
+}
+
 std::pair<managed_unique_ptr<ArrowArray>, managed_unique_ptr<ArrowSchema>> ArrowAdapter::to_arrow(
     std::shared_ptr<ReadColumnBuffer> column, bool downcast_dict_of_large_var) {
     managed_unique_ptr<ArrowSchema> schema = make_managed_unique<ArrowSchema>();

libtiledbsoma/src/utils/arrow_adapter.h

@@ -35,6 +35,7 @@ using namespace tiledb;
 using json = nlohmann::json;
 
 class ReadColumnBuffer;
+class ArrayBuffers;
 class SOMACoordinateSpace;
 
 /**
@@ -106,6 +107,14 @@ class ArrowAdapter {
 
     static bool _isstr(const char* format);
 
+    /**
+     * @brief Convert ArrayBuffer to an list of Arrow arrays.
+     *
+     * @return std::vector<std::pair<std::unique_ptr<ArrowArray>, std::unique_ptr<ArrowSchema>>>
+     */
+    static std::vector<std::pair<managed_unique_ptr<ArrowArray>, managed_unique_ptr<ArrowSchema>>> buffer_to_arrow(
+        std::shared_ptr<ArrayBuffers> buffers, bool downcast_dict_of_large_var = false);
+
     /**
      * @brief Convert ColumnBuffer to an Arrow array.
      *