Revert "[python] Add new SOMAContext class and deprecate SOMATileDBContext (#4363)"

apis/python/HISTORY.md

@@ -9,20 +9,15 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 ### Added
 
 - \[[#4299](https://github.com/single-cell-data/TileDB-SOMA/pull/4299)\] Use a global memory budget for read operations instead of a per column memory budget. The global memory budget allocates splits the budget per column depending on the type and characteristics of each column. Global memory budget is disabled by default under a feature flag and can be enabled by setting `soma.read.use_memory_pool`.
-- \[[#4363](https://github.com/single-cell-data/TileDB-SOMA/pull/4363)\] Add new `SOMAContext` class that replaces the `SOMATileDBContext` class.
-- `PlatformConfig` and `PlatformSchemaConfig` now support `to_dict()` and have a readable `__repr__` for diagnostic logging.
 
 ### Changed
 
 - \[[#4421](https://github.com/single-cell-data/TileDB-SOMA/pull/4421)\] Max domain for integral dimensions are changed for int32, int64 and uint64 columns and are now $(-2^{31} + 1, 2^{31} - 2)$, $(-2^{63} + 1, 2^{63} - 2)$ and $(0, 2^{63} - 2)$ respectively.
 - \[[#4299](https://github.com/single-cell-data/TileDB-SOMA/pull/4299)\] `ManagedQuery` reuses the same buffers for each incomplete read and allocates dedicated buffers when converting to Arrow.
 - \[[#4294](https://github.com/single-cell-data/TileDB-SOMA/pull/4294)\] Use vcpkg instead of custom superbuild, for installing libtiledbsoma and its dependencies. Changes in CI workflows to use the prebuilt libtiledbsoma.
-- \[[#4363](https://github.com/single-cell-data/TileDB-SOMA/pull/4363)\] (BREAKING) The `context` property of a `SOMAObject` now returns a `SOMAContext` instead of a `SOMATileDBContext`.
 
 ### Deprecated
 
-- \[[#4636](https://github.com/single-cell-data/TileDB-SOMA/pull/4363)\] Deprecate `SOMATileDBContext` class in favor of the new `SOMAContext` class.
-
 ### Removed
 
 - \[[#4431](https://github.com/single-cell-data/TileDB-SOMA/pull/4431)\] Remove deprecated support for allowing a dimension in `shape` to be `None` in the `SparseNDArray`.

apis/python/src/tiledbsoma/__init__.py

@@ -179,7 +179,6 @@
 from ._point_cloud_dataframe import PointCloudDataFrame
 from ._query import AxisColumnNames, ExperimentAxisQuery
 from ._scene import Scene
-from ._soma_context import SOMAContext
 from ._sparse_nd_array import SparseNDArray, SparseNDArrayRead
 from .options import SOMATileDBContext, TileDBCreateOptions, TileDBDeleteOptions, TileDBWriteOptions
 from .pytiledbsoma import (
@@ -219,7 +218,6 @@
     "NotCreateableError",
     "PointCloudDataFrame",
     "ResultOrder",
-    "SOMAContext",
     "SOMAError",
     "SOMATileDBContext",
     "ScaleTransform",

apis/python/src/tiledbsoma/_collection.py

@@ -20,7 +20,6 @@
 from ._dataframe import DataFrame
 from ._dense_nd_array import DenseNDArray
 from ._funcs import typeguard_ignore
-from ._soma_context import SOMAContext
 from ._soma_group import SOMAGroup
 from ._soma_object import SOMAObject
 from ._sparse_nd_array import SparseNDArray
@@ -51,7 +50,7 @@ def create(
         uri: str,
         *,
         platform_config: PlatformConfig | None = None,  # noqa: ARG003
-        context: SOMAContext | SOMATileDBContext | None = None,
+        context: SOMATileDBContext | None = None,
         tiledb_timestamp: OpenTimestamp | None = None,
     ) -> Self:
         """Creates and opens a new SOMA collection in storage.
@@ -68,8 +67,9 @@ def create(
                 located in the ``{'tiledb': {'create': ...}}`` key,
                 or as a :class:`~tiledbsoma.TileDBCreateOptions` object.
                 (Currently unused for collections.)
-            context: If provided, the :class:`SOMAContext` to use when creating and opening this collection. Otherwise,
-                the default context will be used and possibly initialized.
+            context:
+                If provided, the :class:`SOMATileDBContext` to use when creating and
+                opening this collection.
             tiledb_timestamp:
                 If specified, overrides the default timestamp
                 used to open this object. If unset, uses the timestamp provided by

apis/python/src/tiledbsoma/_common_nd_array.py

@@ -14,9 +14,10 @@
 
 from ._core_options import PlatformConfig
 from ._soma_array import SOMAArray
-from ._soma_context import SOMAContext
 from ._types import OpenTimestamp, StatusAndReason
-from .options._soma_tiledb_context import SOMATileDBContext
+from .options._soma_tiledb_context import (
+    SOMATileDBContext,
+)
 from .options._tiledb_create_write_options import TileDBCreateOptions
 
 
@@ -33,7 +34,7 @@ def create(
         type: pa.DataType,
         shape: Sequence[int],
         platform_config: PlatformConfig | None = None,
-        context: SOMAContext | SOMATileDBContext | None = None,
+        context: SOMATileDBContext | None = None,
         tiledb_timestamp: OpenTimestamp | None = None,
     ) -> Self:
         """Creates a SOMA ``NDArray`` at the given URI.
@@ -51,8 +52,6 @@ def create(
                 This may be provided as settings in a dictionary, with options
                 located in the ``{'tiledb': {'create': ...}}`` key,
                 or as a :class:`~tiledbsoma.TileDBCreateOptions` object.
-            context: If provided, the :class:`SOMAContext` to use when creating and opening this collection. If not,
-                provide the default context will be used and possibly initialized.
             tiledb_timestamp:
                 If specified, overrides the default timestamp
                 used to open this object. If unset, uses the timestamp provided by

apis/python/src/tiledbsoma/_dask/load.py

@@ -30,7 +30,7 @@
 from tiledbsoma._indexer import IntIndexer
 from tiledbsoma._sparse_nd_array import SparseNDArray
 from tiledbsoma._types import OpenTimestamp
-from tiledbsoma.options._soma_tiledb_context import ConfigDict
+from tiledbsoma.options._soma_tiledb_context import ConfigDict, ConfigVal
 
 if TYPE_CHECKING:
     try:
@@ -108,7 +108,7 @@ def load_daskarray(
     *,
     coords: SparseNDCoords | None = None,
     chunk_size: ChunkSize | None = None,
-    tiledb_config: dict[str, str | float] | None = None,
+    tiledb_config: dict[str, ConfigVal] | None = None,
     format: Format = "csr",
     result_order: ResultOrderStr = ResultOrder.AUTO,
     platform_config: PlatformConfig | None = None,
@@ -142,10 +142,10 @@ def load_daskarray(
 
     if isinstance(layer, SparseNDArray):
         if tiledb_config:
-            tiledb_config = {**layer.context.config}
+            tiledb_config = {**layer.context.tiledb_config}
             tiledb_config.update(**tiledb_config)
         else:
-            tiledb_config = layer.context.config  # type: ignore[assignment]
+            tiledb_config = layer.context.tiledb_config
 
         chunk_joinids = da.from_array(arr, chunks=(1, 1))
         meta = sp.csr_matrix((0, 0), dtype=dtype) if format == "csr" else sp.csc_matrix((0, 0), dtype=dtype)

apis/python/src/tiledbsoma/_dask/util.py

@@ -17,7 +17,8 @@
 from typing_extensions import TypeAlias
 
 from tiledbsoma._core_options import SparseNDCoord, SparseNDCoords
-from tiledbsoma._soma_context import SOMAContext
+from tiledbsoma.options import SOMATileDBContext
+from tiledbsoma.options._soma_tiledb_context import ConfigVal
 
 ChunkSize = Union[int, tuple[Union[int, None], int]]
 JoinIDs: TypeAlias = NDArray[int64]
@@ -33,7 +34,7 @@ class SOMADaskConfig(TypedDict, total=False):
     """
 
     chunk_size: ChunkSize
-    tiledb_config: dict[str, Union[str, float]]  # noqa: UP007
+    tiledb_config: dict[str, ConfigVal]
 
 
 def chunk_ids_sizes(joinids: JoinIDs, chunk_size: int, dim_size: int) -> tuple[list[JoinIDs], list[int]]:  # noqa: ARG001
@@ -52,7 +53,7 @@ def chunk_ids_sizes(joinids: JoinIDs, chunk_size: int, dim_size: int) -> tuple[l
     return chunk_joinids, chunk_sizes
 
 
-def make_context(tiledb_config: dict[str, Any]) -> SOMAContext:
+def make_context(tiledb_config: dict[str, Any]) -> SOMATileDBContext:
     """Create and cache ``SOMATileDBContext``s within Dask worker processes.
 
     This wrapper just flattens the ``tiledb_config`` ``dict`` into hashable ``tuple``s (for use with ``cache``).
@@ -61,12 +62,12 @@ def make_context(tiledb_config: dict[str, Any]) -> SOMAContext:
 
 
 @cache
-def _make_context(tiledb_configs: tuple[tuple[str, Any], ...]) -> SOMAContext:
+def _make_context(tiledb_configs: tuple[tuple[str, Any], ...]) -> SOMATileDBContext:
     """Create and cache ``SOMATileDBContext``s within Dask worker processes.
 
     ``tiledb_config`` is conceptually a ``dict``, but flattened into hashable ``tuple``s here, for use with ``cache``.
     """
-    return SOMAContext.create(config=dict(tiledb_configs))
+    return SOMATileDBContext(tiledb_config=dict(tiledb_configs))
 
 
 def coord_to_joinids(coord: SparseNDCoord, n: int) -> JoinIDs:

apis/python/src/tiledbsoma/_dataframe.py

@@ -23,7 +23,6 @@
 from ._query_condition import QueryCondition
 from ._read_iters import TableReadIter
 from ._soma_array import SOMAArray
-from ._soma_context import SOMAContext
 from ._types import (
     NPFInfo,
     NPFloating,
@@ -32,8 +31,8 @@
     OpenTimestamp,
     StatusAndReason,
 )
-from ._util import tiledb_timestamp_to_ms
-from .options import SOMATileDBContext, _update_context_and_timestamp
+from .options import SOMATileDBContext
+from .options._soma_tiledb_context import _validate_soma_tiledb_context
 from .options._tiledb_create_write_options import TileDBCreateOptions, TileDBDeleteOptions, TileDBWriteOptions
 from .options._util import build_clib_platform_config
 
@@ -143,7 +142,7 @@ def create(
         domain: Domain,
         index_column_names: Sequence[str] = (SOMA_JOINID,),
         platform_config: PlatformConfig | None = None,
-        context: SOMAContext | SOMATileDBContext | None = None,
+        context: SOMATileDBContext | None = None,
         tiledb_timestamp: OpenTimestamp | None = None,
     ) -> DataFrame:
         """Creates the data structure on disk/S3/cloud.
@@ -200,6 +199,7 @@ def create(
             Maturing.
         """
         _util.check_type("schema", schema, (pa.Schema,))
+        context = _validate_soma_tiledb_context(context)
 
         soma_domain: list[tuple[Any, Any] | None] = []
         if domain is None:
@@ -213,27 +213,26 @@ def create(
             soma_domain.append(_cast_domain_to_cpp_type(slot_soma_domain, schema, index_column_name))
 
         plt_cfg = build_clib_platform_config(platform_config)
-        context, tiledb_timestamp = _update_context_and_timestamp(context, tiledb_timestamp)
-        timestamp_ms = tiledb_timestamp_to_ms(tiledb_timestamp)
+        timestamp_ms = context._open_timestamp_ms(tiledb_timestamp)
         try:
             clib.SOMADataFrame.create(
                 uri,
                 schema=schema,
                 index_column_names=index_column_names,
                 index_column_domains=soma_domain,
-                ctx=context._handle,
+                ctx=context.native_context,
                 platform_config=plt_cfg,
                 timestamp=(0, timestamp_ms),
             )
         except SOMAError as e:
             raise map_exception_for_create(e, uri) from None
 
-        timestamp_ms = tiledb_timestamp_to_ms(tiledb_timestamp)
         try:
+            timestamp_ms = context._open_timestamp_ms(tiledb_timestamp)
             handle = clib.SOMADataFrame.open(
                 uri,
                 mode=clib.OpenMode.soma_write,
-                context=context._handle,
+                context=context.native_context,
                 timestamp=(0, timestamp_ms),
             )
 

apis/python/src/tiledbsoma/_dense_nd_array.py

@@ -21,7 +21,6 @@
 from ._exception import SOMAError
 from ._managed_query import ManagedQuery
 from ._read_iters import TableReadIter
-from ._soma_context import SOMAContext
 from ._types import OpenTimestamp, Slice
 from ._util import dense_indices_to_shape
 from .options._soma_tiledb_context import SOMATileDBContext
@@ -92,7 +91,7 @@ def create(
         type: pa.DataType,
         shape: Sequence[int | None],
         platform_config: PlatformConfig | None = None,
-        context: SOMAContext | SOMATileDBContext | None = None,
+        context: SOMATileDBContext | None = None,
         tiledb_timestamp: OpenTimestamp | None = None,
     ) -> Self:
         """Creates a SOMA ``DenseNDArray`` at the given URI.

apis/python/src/tiledbsoma/_experiment.py

@@ -26,9 +26,9 @@
 from ._point_cloud_dataframe import PointCloudDataFrame
 from ._query import ExperimentAxisQuery
 from ._scene import Scene
-from ._soma_context import SOMAContext
 from ._soma_object import SOMAObject
 from ._sparse_nd_array import SparseNDArray
+from .options import SOMATileDBContext
 from .options._tiledb_create_write_options import TileDBDeleteOptions
 
 
@@ -375,7 +375,9 @@ def _create_var_axis_candidates(exp: Experiment, ms_name: str) -> list[_ArrayDel
     return candidates
 
 
-def _query_joinids(uri: str, coords: SparseDFCoords, value_filter: str | None, context: SOMAContext) -> pa.Int64Array:
+def _query_joinids(
+    uri: str, coords: SparseDFCoords, value_filter: str | None, context: SOMATileDBContext
+) -> pa.Int64Array:
     with DataFrame.open(uri, context=context) as df:
         return (
             df

apis/python/src/tiledbsoma/_factory.py

@@ -31,11 +31,10 @@
 from ._core_options import OpenMode
 from ._exception import DoesNotExistError, SOMAError, is_does_not_exist_error
 from ._funcs import typeguard_ignore
-from ._soma_context import SOMAContext
 from ._soma_object import SOMAObject
 from ._types import OpenTimestamp
-from ._util import tiledb_timestamp_to_ms
-from .options import SOMATileDBContext, _update_context_and_timestamp
+from .options import SOMATileDBContext
+from .options._soma_tiledb_context import _validate_soma_tiledb_context
 
 _Obj = TypeVar("_Obj", bound="SOMAObject")
 
@@ -46,7 +45,7 @@ def open(
     mode: OpenMode = ...,
     *,
     soma_type: str | None = None,
-    context: SOMAContext | SOMATileDBContext | None = None,
+    context: SOMATileDBContext | None = None,
     tiledb_timestamp: OpenTimestamp | None = None,
 ) -> SOMAObject: ...
 
@@ -57,7 +56,7 @@ def open(
     mode: OpenMode,
     *,
     soma_type: type[_Obj],
-    context: SOMAContext | SOMATileDBContext | None = None,
+    context: SOMATileDBContext | None = None,
     tiledb_timestamp: OpenTimestamp | None = None,
 ) -> _Obj: ...
 
@@ -68,7 +67,7 @@ def open(
     mode: OpenMode = "r",
     *,
     soma_type: type[SOMAObject] | str | None = None,
-    context: SOMAContext | SOMATileDBContext | None = None,
+    context: SOMATileDBContext | None = None,
     tiledb_timestamp: OpenTimestamp | None = None,
 ) -> SOMAObject:
     """Opens a TileDB SOMA object.
@@ -87,8 +86,7 @@ def open(
             If the stored SOMA object is not of the correct type, an error will be
             raised.
         context:
-            If provided, the :class:`SOMAContext` to use when creating and opening this collection. If not,
-            provide the default context will be used and possibly initialized.
+            If set, the :class:`SOMATileDBContext` data to use.
         tiledb_timestamp:
             If specified, overrides the default timestamp
             used to open this object. If unset, uses the timestamp provided by
@@ -116,7 +114,7 @@ def open(
         Maturing.
     """
     if soma_type is None:
-        context, tiledb_timestamp = _update_context_and_timestamp(context, tiledb_timestamp)
+        context = _validate_soma_tiledb_context(context)
         return _open_soma_object(uri, mode, context, tiledb_timestamp)
 
     if isinstance(soma_type, str):
@@ -138,19 +136,19 @@ def open(
 def _open_soma_object(
     uri: str,
     mode: OpenMode,
-    context: SOMAContext,
+    context: SOMATileDBContext,
     tiledb_timestamp: OpenTimestamp | None,
     clib_type: str | None = None,
 ) -> SOMAObject:
     """Picks out the appropriate SOMA class for a handle and wraps it."""
     if clib_type is None or clib_type.lower() in ["somaarray", "somagroup"]:
-        timestamp_ms = tiledb_timestamp_to_ms(tiledb_timestamp)
+        timestamp_ms = context._open_timestamp_ms(tiledb_timestamp)
         open_mode = _tdb_handles._open_mode_to_clib_mode(mode)
         try:
             handle = clib.SOMAObject.open(
                 uri=uri,
                 mode=open_mode,
-                context=context._handle,
+                context=context.native_context,
                 timestamp=(0, timestamp_ms),
                 clib_type=clib_type,
             )