[c++/python] Revert 3764 for global-order writes (#3886)

johnkerl · web-flow · commit 5e5212233c24 · 2025-04-01T15:12:10.000-04:00
diff --git a/apis/python/src/tiledbsoma/_dataframe.py b/apis/python/src/tiledbsoma/_dataframe.py
@@ -29,7 +29,7 @@
 from . import pytiledbsoma as clib
 from ._constants import SOMA_GEOMETRY, SOMA_JOINID
 from ._exception import SOMAError, map_exception_for_create
-from ._read_iters import TableReadIter
+from ._read_iters import ManagedQuery, TableReadIter
 from ._soma_array import SOMAArray
 from ._tdb_handles import DataFrameWrapper
 from ._types import (
@@ -801,7 +801,12 @@ def write(
                 "TileDBWriteOptions instead of TileDBCreateOptions"
             )
         write_options = TileDBWriteOptions.from_platform_config(platform_config)
-        self._write_table(values, write_options.sort_coords)
+        sort_coords = write_options.sort_coords
+
+        for batch in values.to_batches():
+            mq = ManagedQuery(self)
+            mq._handle.set_array_data(batch)
+            mq._handle.submit_write(sort_coords or False)
 
         if write_options.consolidate_and_vacuum:
             self._handle._handle.consolidate_and_vacuum()
diff --git a/apis/python/src/tiledbsoma/_dense_nd_array.py b/apis/python/src/tiledbsoma/_dense_nd_array.py
@@ -315,7 +315,6 @@ def write(
         _util._set_coords(mq, new_coords)
         mq._handle.set_column_data("soma_data", input)
         mq._handle.submit_write()
-        mq._handle.finalize()
 
         tiledb_write_options = TileDBWriteOptions.from_platform_config(platform_config)
         if tiledb_write_options.consolidate_and_vacuum:
diff --git a/apis/python/src/tiledbsoma/_geometry_dataframe.py b/apis/python/src/tiledbsoma/_geometry_dataframe.py
@@ -38,7 +38,7 @@
     _revise_domain_for_extent,
 )
 from ._exception import SOMAError, map_exception_for_create
-from ._read_iters import TableReadIter
+from ._read_iters import ManagedQuery, TableReadIter
 from ._spatial_dataframe import SpatialDataFrame
 from ._spatial_util import (
     coordinate_space_from_json,
@@ -498,16 +498,24 @@ def write(
         _util.check_type("values", values, (pa.Table,))
 
         write_options: Union[TileDBCreateOptions, TileDBWriteOptions]
+        sort_coords = None
         if isinstance(platform_config, TileDBCreateOptions):
             raise ValueError(
                 "As of TileDB-SOMA 1.13, the write method takes "
                 "TileDBWriteOptions instead of TileDBCreateOptions"
             )
         write_options = TileDBWriteOptions.from_platform_config(platform_config)
-        self._write_table(values, write_options.sort_coords)
+        sort_coords = write_options.sort_coords
+
+        clib_dataframe = self._handle._handle
+
+        for batch in values.to_batches():
+            mq = ManagedQuery(self, None)
+            mq._handle.set_array_data(batch)
+            mq._handle.submit_write(sort_coords or False)
 
         if write_options.consolidate_and_vacuum:
-            self._handle._handle.consolidate_and_vacuum()
+            clib_dataframe.consolidate_and_vacuum()
 
         return self
 
diff --git a/apis/python/src/tiledbsoma/_point_cloud_dataframe.py b/apis/python/src/tiledbsoma/_point_cloud_dataframe.py
@@ -30,7 +30,7 @@
     _revise_domain_for_extent,
 )
 from ._exception import SOMAError, map_exception_for_create
-from ._read_iters import TableReadIter
+from ._read_iters import ManagedQuery, TableReadIter
 from ._spatial_dataframe import SpatialDataFrame
 from ._spatial_util import (
     coordinate_space_from_json,
@@ -480,10 +480,16 @@ def write(
             )
         write_options = TileDBWriteOptions.from_platform_config(platform_config)
         sort_coords = write_options.sort_coords
-        self._write_table(values, sort_coords)
+
+        clib_dataframe = self._handle._handle
+
+        for batch in values.to_batches():
+            mq = ManagedQuery(self, None)
+            mq._handle.set_array_data(batch)
+            mq._handle.submit_write(sort_coords or False)
 
         if write_options.consolidate_and_vacuum:
-            self._handle._handle.consolidate_and_vacuum()
+            clib_dataframe.consolidate_and_vacuum()
 
         return self
 
diff --git a/apis/python/src/tiledbsoma/_soma_array.py b/apis/python/src/tiledbsoma/_soma_array.py
@@ -11,7 +11,6 @@
 
 # This package's pybind11 code
 from . import pytiledbsoma as clib  # noqa: E402
-from ._read_iters import ManagedQuery
 from ._soma_object import SOMAObject
 
 
@@ -160,58 +159,3 @@ def _maxdomain(self) -> Tuple[Tuple[Any, Any], ...]:
           resized up to core (max) domain.
         """
         return self._handle.maxdomain
-
-    def _write_table(self, values: pa.Table, sort_coords: bool) -> None:
-        """Helper function that sets the correct result order for the layout
-        and allows for multiple submissions before calling `submit` and `finalize`
-        for unordered write or `submit_and_finalize` for global order writes.
-
-        Args:
-            values:
-                An `Arrow table <https://arrow.apache.org/docs/python/generated/pyarrow.Table.html>`_
-                containing all columns, including the index columns. The schema
-                for the values must match the schema for the :class:`DataFrame`.
-
-                If a column is of categorical type in the schema and a
-                flattened/non-categorical column is presented for data on write,
-                a ``ValueError`` is raised.  If a column is of non-categorical
-                type in the schema and a categorical column is presented for data
-                on write, the data are written as an array of category values,
-                and the category-type information is not saved.
-            sort_coords:
-                Whether the coordinates need to be sorted (True) or are already
-                sorted in global order (False). In the PlatformConfig, this is
-                is to True by default.
-        """
-        batches = values.to_batches()
-        if not batches:
-            return
-
-        layout = (
-            clib.ResultOrder.unordered if sort_coords else clib.ResultOrder.globalorder
-        )
-
-        if layout == clib.ResultOrder.unordered:
-            for batch in batches:
-                # Create new ManagedQuery per each batch
-                mq = ManagedQuery(self)._handle
-                mq.set_layout(layout)
-                mq.set_array_data(batch)
-
-                # Submit and flush every batch
-                mq.submit_write()
-                mq.finalize()
-
-        else:  # global order
-            # Create a single ManagedQuery at the beginning
-            mq = ManagedQuery(self)._handle
-            mq.set_layout(layout)
-
-            # Submit for each batch but don't finalize
-            for batch in batches[:-1]:
-                mq.set_array_data(batch)
-                mq.submit_write()
-
-            # Only finalize at the last batch
-            mq.set_array_data(batches[-1])
-            mq.submit_and_finalize()
diff --git a/apis/python/src/tiledbsoma/_sparse_nd_array.py b/apis/python/src/tiledbsoma/_sparse_nd_array.py
@@ -308,12 +308,14 @@ def write(
         """
 
         write_options: Union[TileDBCreateOptions, TileDBWriteOptions]
+        sort_coords = None
         if isinstance(platform_config, TileDBCreateOptions):
             raise ValueError(
                 "As of TileDB-SOMA 1.13, the write method takes "
                 "TileDBWriteOptions instead of TileDBCreateOptions"
             )
         write_options = TileDBWriteOptions.from_platform_config(platform_config)
+        sort_coords = write_options.sort_coords
 
         clib_sparse_array = self._handle._handle
 
@@ -322,14 +324,6 @@ def write(
             data, coords = values.to_numpy()
 
             mq = ManagedQuery(self, platform_config)
-
-            layout = (
-                clib.ResultOrder.unordered
-                if write_options.sort_coords
-                else clib.ResultOrder.globalorder
-            )
-            mq._handle.set_layout(layout)
-
             for i, c in enumerate(coords.T):
                 mq._handle.set_column_data(
                     f"soma_dim_{i}",
@@ -344,12 +338,7 @@ def write(
                     data, dtype=self.schema.field("soma_data").type.to_pandas_dtype()
                 ),
             )
-
-            if layout == clib.ResultOrder.unordered:
-                mq._handle.submit_write()
-                mq._handle.finalize()
-            else:
-                mq._handle.submit_and_finalize()
+            mq._handle.submit_write(sort_coords or True)
 
             if write_options.consolidate_and_vacuum:
                 # Consolidate non-bulk data
@@ -366,14 +355,6 @@ def write(
             sp = values.to_scipy().tocoo()
 
             mq = ManagedQuery(self, platform_config)
-
-            layout = (
-                clib.ResultOrder.unordered
-                if write_options.sort_coords
-                else clib.ResultOrder.globalorder
-            )
-            mq._handle.set_layout(layout)
-
             for i, c in enumerate([sp.row, sp.col]):
                 mq._handle.set_column_data(
                     f"soma_dim_{i}",
@@ -388,20 +369,20 @@ def write(
                     sp.data, dtype=self.schema.field("soma_data").type.to_pandas_dtype()
                 ),
             )
-
-            if layout == clib.ResultOrder.unordered:
-                mq._handle.submit_write()
-                mq._handle.finalize()
-            else:
-                mq._handle.submit_and_finalize()
+            mq._handle.submit_write(sort_coords or True)
 
             if write_options.consolidate_and_vacuum:
                 # Consolidate non-bulk data
                 clib_sparse_array.consolidate_and_vacuum()
             return self
 
         if isinstance(values, pa.Table):
-            self._write_table(values, write_options.sort_coords)
+            # Write bulk data
+            for batch in values.to_batches():
+                # clib_sparse_array.write(batch, sort_coords or False)
+                mq = ManagedQuery(self, None)
+                mq._handle.set_array_data(batch)
+                mq._handle.submit_write(sort_coords or False)
 
             if write_options.consolidate_and_vacuum:
                 # Consolidate non-bulk data
diff --git a/apis/python/src/tiledbsoma/managed_query.cc b/apis/python/src/tiledbsoma/managed_query.cc
@@ -156,39 +156,16 @@ void load_managed_query(py::module& m) {
                 }
                 py::gil_scoped_acquire acquire;
             })
-        .def("reset_columns", &ManagedQuery::reset_columns)
-
         .def(
             "submit_write",
-            [](ManagedQuery& mq) {
-                try {
-                    mq.submit_write();
-                } catch (const std::exception& e) {
-                    TPY_ERROR_LOC(e.what());
-                }
-            },
-            py::call_guard<py::gil_scoped_release>())
-
-        .def(
-            "submit_and_finalize",
-            [](ManagedQuery& mq) {
-                try {
-                    mq.submit_and_finalize();
-                } catch (const std::exception& e) {
-                    TPY_ERROR_LOC(e.what());
-                }
-            },
-            py::call_guard<py::gil_scoped_release>())
-
-        .def(
-            "finalize",
-            [](ManagedQuery& mq) {
+            [](ManagedQuery& mq, bool sort_coords) {
                 try {
-                    mq.finalize();
+                    mq.submit_write(sort_coords);
                 } catch (const std::exception& e) {
                     TPY_ERROR_LOC(e.what());
                 }
             },
+            "sort_coords"_a = false,
             py::call_guard<py::gil_scoped_release>())
 
         .def("reset", &ManagedQuery::reset)
diff --git a/apis/python/src/tiledbsoma/pytiledbsoma.cc b/apis/python/src/tiledbsoma/pytiledbsoma.cc
@@ -68,9 +68,7 @@ PYBIND11_MODULE(pytiledbsoma, m) {
     py::enum_<ResultOrder>(m, "ResultOrder")
         .value("automatic", ResultOrder::automatic)
         .value("rowmajor", ResultOrder::rowmajor)
-        .value("colmajor", ResultOrder::colmajor)
-        .value("unordered", ResultOrder::unordered)
-        .value("globalorder", ResultOrder::global);
+        .value("colmajor", ResultOrder::colmajor);
 
     py::enum_<URIType>(m, "URIType")
         .value("automatic", URIType::automatic)
diff --git a/apis/python/tests/test_dataframe.py b/apis/python/tests/test_dataframe.py
@@ -2442,58 +2442,3 @@ def nan_check(expected_nan, dict_vals):
         assert nan_check(signaling_nan, actual_dict_vals)
         assert_array_equal(expected_data1["A"], actual_data[:4])
         assert_array_equal(expected_data2["A"], actual_data[4:])
-
-
-def test_fragments_in_writes(tmp_path):
-    uri = tmp_path.as_posix()
-
-    # --- three dataframes, all with identical schema
-    df_0 = pd.DataFrame(
-        {
-            "soma_joinid": pd.Series([0, 1, 2, 3], dtype=np.int64),
-            "obs": pd.Series(["A", "B", "A", "B"], dtype="str"),
-        }
-    )
-    df_1 = pd.DataFrame(
-        {
-            "soma_joinid": pd.Series([4, 5, 6, 7], dtype=np.int64),
-            "obs": pd.Series(["A", "A", "B", "B"], dtype="str"),
-        }
-    )
-    df_2 = pd.DataFrame(
-        {
-            "soma_joinid": pd.Series([8, 9, 10, 11], dtype=np.int64),
-            "obs": pd.Series(["B", "C", "B", "C"], dtype="str"),
-        }
-    )
-    expected_df = pd.concat([df_0, df_1, df_2], ignore_index=True)
-
-    soma.DataFrame.create(
-        uri,
-        schema=pa.Schema.from_pandas(df_0, preserve_index=False),
-        domain=[[0, 11]],
-    ).close()
-
-    with soma.open(uri, mode="w") as A:
-        # Three-chunk table
-        A.write(
-            pa.concat_tables(
-                [
-                    pa.Table.from_pandas(df_0, preserve_index=False),
-                    pa.Table.from_pandas(df_1, preserve_index=False),
-                    pa.Table.from_pandas(df_2, preserve_index=False),
-                ]
-            ),
-            platform_config=soma.TileDBWriteOptions(**{"sort_coords": False}),
-        )
-
-    # There should be a single fragment even though there are three chunks (and
-    # therefore three submits) in the array because we only finalize once at
-    # the end
-    assert len(list((Path(uri) / "__commits").iterdir())) == 1
-    assert len(list((Path(uri) / "__fragments").iterdir())) == 1
-
-    with soma.open(uri) as A:
-        df = A.read().concat().to_pandas()
-
-    assert df.equals(expected_df)
diff --git a/libtiledbsoma/src/soma/managed_query.cc b/libtiledbsoma/src/soma/managed_query.cc
diff --git a/libtiledbsoma/src/soma/managed_query.h b/libtiledbsoma/src/soma/managed_query.h
diff --git a/libtiledbsoma/test/unit_soma_array.cc b/libtiledbsoma/test/unit_soma_array.cc
