Store object references when setting column data until write is submitted (#4350)

Author: XanthosXanthopoulos

apis/python/src/tiledbsoma/_dense_nd_array.py

@@ -366,8 +366,8 @@ def write(
         mq = ManagedQuery(self, platform_config)
         mq._handle.set_layout(order)
         mq.set_coords(new_coords)
-        mq._handle.set_column_data("soma_data", input)
-        mq._handle.submit_write()
+        mq.set_column_data("soma_data", input)
+        mq.submit_write()
         mq._handle.finalize()
 
         tiledb_write_options = TileDBWriteOptions.from_platform_config(platform_config)

apis/python/src/tiledbsoma/_managed_query.py

@@ -30,6 +30,7 @@ class ManagedQuery:
     _array: SOMAArray
     _platform_config: options.PlatformConfig | None = None
     _handle: clib.ManagedQuery = attrs.field(init=False)
+    _ref_store: list[object] = attrs.field(init=False, default=[])
 
     def __attrs_post_init__(self) -> None:
         array_handle = self._array._handle
@@ -218,3 +219,25 @@ def set_geometry_coord(
             column.set_dim_points_double_array(self._handle, [coord])
         else:
             raise ValueError(f"Unsupported spatial coordinate type. Expected slice or float, found {type(coord)}")
+
+    def set_column_data(self, dim_name: str, data: np.typing.NDArray) -> None:
+        # store a reference to the data being written
+        # libtiledbsoma will try not to copy any data to temporary buffers when writing data but the user is free
+        # to pass temporary data objects to write. Submitting the write buffer to TileDB can happen after the
+        # lifespan of the temporary object so ManagedQuery need to preserve a reference to each data object until
+        # the query is submitted
+        self._ref_store.append(data)
+
+        self._handle.set_column_data(dim_name, data)
+
+    def submit_write(self) -> None:
+        self._handle.submit_write()
+
+        # clear stored data objects
+        self._ref_store.clear()
+
+    def submit_and_finalize(self) -> None:
+        self._handle.submit_and_finalize()
+
+        # clear stored data objects
+        self._ref_store.clear()

apis/python/src/tiledbsoma/_sparse_nd_array.py

@@ -389,23 +389,23 @@ def write(
             mq._handle.set_layout(layout)
 
             for i, c in enumerate(coords.T):
-                mq._handle.set_column_data(
+                mq.set_column_data(
                     f"soma_dim_{i}",
                     np.array(
                         c,
                         dtype=self.schema.field(f"soma_dim_{i}").type.to_pandas_dtype(),
                     ),
                 )
-            mq._handle.set_column_data(
+            mq.set_column_data(
                 "soma_data",
                 np.array(data, dtype=self.schema.field("soma_data").type.to_pandas_dtype()),
             )
 
             if layout == clib.ResultOrder.unordered:
-                mq._handle.submit_write()
+                mq.submit_write()
                 mq._handle.finalize()
             else:
-                mq._handle.submit_and_finalize()
+                mq.submit_and_finalize()
 
             if write_options.consolidate_and_vacuum:
                 # Consolidate non-bulk data
@@ -425,23 +425,23 @@ def write(
             mq._handle.set_layout(layout)
 
             for i, c in enumerate([sp.row, sp.col]):
-                mq._handle.set_column_data(
+                mq.set_column_data(
                     f"soma_dim_{i}",
                     np.array(
                         c,
                         dtype=self.schema.field(f"soma_dim_{i}").type.to_pandas_dtype(),
                     ),
                 )
-            mq._handle.set_column_data(
+            mq.set_column_data(
                 "soma_data",
                 np.array(sp.data, dtype=self.schema.field("soma_data").type.to_pandas_dtype()),
             )
 
             if layout == clib.ResultOrder.unordered:
-                mq._handle.submit_write()
+                mq.submit_write()
                 mq._handle.finalize()
             else:
-                mq._handle.submit_and_finalize()
+                mq.submit_and_finalize()
 
             if write_options.consolidate_and_vacuum:
                 # Consolidate non-bulk data

apis/python/src/tiledbsoma/managed_query.cc

@@ -155,7 +155,7 @@ void load_managed_query(py::module& m) {
                 py::gil_scoped_release release;
                 try {
                     mq.setup_write_column(
-                        name, data.size(), (const void*)data_info.ptr, static_cast<uint64_t*>(nullptr), nullptr, true);
+                        name, data.size(), (const void*)data_info.ptr, static_cast<uint64_t*>(nullptr), nullptr, false);
                 } catch (const std::exception& e) {
                     TPY_ERROR_LOC(e.what());
                 }