[c++] Move re-open array from SOMAArray::write to `ManagedQuery::wr…

…ite`
single-cell-data · Feb 4, 2025 · 40f486b · 40f486b
1 parent ce0dc82
commit 40f486b
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 3 deletions.
diff --git a/apis/python/tests/test_dataframe.py b/apis/python/tests/test_dataframe.py
@@ -2124,3 +2124,44 @@ def test_arrow_table_validity_with_slicing(tmp_path):
         assert_array_equal(pdf["mybool"], table["mybool"])
         assert_array_equal(pdf["mydatetime"], table["mydatetime"])
         assert_array_equal(pdf["myenum"], table["myenum"])
+
+
+def test_enum_regression_62887(tmp_path):
+    uri = tmp_path.as_posix()
+
+    schema = pa.schema(
+        [
+            pa.field("soma_joinid", pa.int64(), nullable=False),
+            pa.field("A", pa.dictionary(pa.int8(), pa.int8())),
+        ]
+    )
+
+    tbl = pa.Table.from_pydict(
+        {
+            "soma_joinid": pa.chunked_array([[0, 1, 2, 3, 4, 5, 6, 7], [8, 9]]),
+            "A": pa.chunked_array(
+                [
+                    pa.DictionaryArray.from_arrays(
+                        indices=pa.array([0, 0, 0, 0, 0, 0, 0, 0], type=pa.int8()),
+                        dictionary=pa.array(
+                            [0, 1, 2, 3, 4, 5, 6, 7, 8], type=pa.int8()
+                        ),
+                    ),
+                    pa.DictionaryArray.from_arrays(
+                        indices=pa.array([0, 0], type=pa.int8()),
+                        dictionary=pa.array(
+                            [0, 1, 2, 3, 4, 5, 6, 7, 8], type=pa.int8()
+                        ),
+                    ),
+                ]
+            ),
+        }
+    )
+
+    with soma.DataFrame.create(
+        uri, schema=schema, index_column_names=["soma_joinid"], domain=[(0, 10000000)]
+    ) as A:
+        A.write(tbl)
+
+    with soma.open(uri) as A:
+        assert_array_equal(A.read().concat()["A"], tbl["A"])
diff --git a/libtiledbsoma/src/soma/managed_query.cc b/libtiledbsoma/src/soma/managed_query.cc
@@ -193,6 +193,11 @@ void ManagedQuery::submit_write(bool sort_coords) {
         query_->submit();
         query_->finalize();
     }
+
+    // When we evolve the schema, the ArraySchema needs to be updated to the
+    // latest version so re-open the Array
+    array_->close();
+    array_->open(TILEDB_WRITE);
 }
 
 void ManagedQuery::submit_read() {

diff --git a/libtiledbsoma/src/soma/soma_array.cc b/libtiledbsoma/src/soma/soma_array.cc
@@ -394,9 +394,8 @@ void SOMAArray::write(bool sort_coords) {
     }
     mq_->submit_write(sort_coords);
 
-    // When we evolve the schema, the ArraySchema needs to be updated to the
-    // latest version so re-open the Array
-    arr_ = std::make_shared<Array>(*ctx_->tiledb_ctx(), uri_, TILEDB_WRITE);
+    // ManagedQuery::submit_write re-opens the array so reset the ManagedQuery
+    // with the latest version of the array
     mq_ = std::make_unique<ManagedQuery>(arr_, ctx_->tiledb_ctx(), name_);
 }