Fix writes when using slice of arrow table

single-cell-data · Mar 27, 2024 · 67a0b54 · 67a0b54
1 parent 2bb57f1
commit 67a0b54
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 8 deletions.
diff --git a/apis/python/src/tiledbsoma/_dataframe.py b/apis/python/src/tiledbsoma/_dataframe.py
@@ -455,6 +455,8 @@ def write(
         _util.check_type("values", values, (pa.Table,))
 
         target_schema = []
+        print("self.schema")
+        print(self.schema)
         for input_field in values.schema:
             target_field = self.schema.field(input_field.name)
 
@@ -465,11 +467,9 @@ def write(
                 target_schema.append(target_field.with_type(pa.uint8()))
             else:
                 target_schema.append(target_field)
+        # print(values)
         values = values.cast(pa.schema(target_schema, values.schema.metadata))
-
-        print("HELLLLOOOOOOOOOOOOOOOOO")
-        print()
-
+
         for batch in values.to_batches():
             self._handle.write(batch)
 

diff --git a/apis/python/src/tiledbsoma/io/ingest.py b/apis/python/src/tiledbsoma/io/ingest.py
@@ -1219,7 +1219,6 @@ def _write_arrow_table(
         )
         handle.write(arrow_table)
 
-
 def _write_dataframe(
     df_uri: str,
     df: pd.DataFrame,

diff --git a/apis/python/src/tiledbsoma/soma_array.cc b/apis/python/src/tiledbsoma/soma_array.cc
@@ -91,9 +91,25 @@ void write(SOMAArray& array, py::handle py_batch) {
             }
         }
 
+        auto np = py::module::import("numpy");
+        auto table_offset = arr_->offset;
+        auto data_size = tiledb::impl::type_size(ArrowAdapter::to_tiledb_format(sch_->format));
+
+        if(offsets){
+            offsets += table_offset;
+        }
+        if(validities){
+            validities += table_offset;
+        }
+
         array.set_column_data(
-            sch_->name, arr_->length, data, offsets, validities);
+            sch_->name,
+            arr_->length,
+            (char*)data + table_offset * data_size,
+            offsets,
+            nullptr);
     }
+
     try {
         array.write();
     } catch (const std::exception& e) {

diff --git a/apis/python/tests/test_io.py b/apis/python/tests/test_io.py
@@ -171,3 +171,4 @@ def test_write_arrow_table(tmp_path, num_rows, cap_nbytes):
         with soma.DataFrame.open(uri) as sdf:
             pdf = sdf.read().concat().to_pandas()
             assert list(pdf["foo"]) == pydict["foo"]
+            assert list(pdf["bar"]) == pydict["bar"]
diff --git a/libtiledbsoma/src/soma/soma_array.cc b/libtiledbsoma/src/soma/soma_array.cc
@@ -404,11 +404,10 @@ void SOMAArray::write() {
     if (mq_->query_type() != TILEDB_WRITE) {
         throw TileDBSOMAError("[SOMAArray] array must be opened in write mode");
     }
-
     mq_->submit_write();
 
     mq_->reset();
-    array_buffer_ = nullptr;
+    // array_buffer_ = nullptr;
 }
 
 uint64_t SOMAArray::nnz() {