Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[python] Remove the no-longer-needed wrapper layers [WIP] #3209

Draft
wants to merge 24 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions apis/python/src/tiledbsoma/_collection.py
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,7 @@ def create(
timestamp=(0, timestamp_ms),
)
handle = wrapper.open(uri, "w", context, tiledb_timestamp)

return cls(
handle,
_dont_call_this_use_create_or_open_instead="tiledbsoma-internal-code",
Expand Down Expand Up @@ -519,6 +520,7 @@ class Collection( # type: ignore[misc] # __eq__ false positive
__slots__ = ()

_wrapper_type = _tdb_handles.CollectionWrapper
_clib_handle_type = clib.SOMACollection


@typeguard_ignore
Expand Down
6 changes: 3 additions & 3 deletions apis/python/src/tiledbsoma/_common_nd_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@ def shape(self) -> Tuple[int, ...]:
Lifecycle:
Maturing.
"""
return cast(Tuple[int, ...], tuple(self._handle.shape))
return cast(Tuple[int, ...], tuple(self._clib_handle.shape))

@property
def maxshape(self) -> Tuple[int, ...]:
Expand All @@ -104,7 +104,7 @@ def maxshape(self) -> Tuple[int, ...]:
Lifecycle:
Maturing.
"""
return cast(Tuple[int, ...], tuple(self._handle.maxshape))
return cast(Tuple[int, ...], tuple(self._clib_handle.maxshape))

@property
def tiledbsoma_has_upgraded_shape(self) -> bool:
Expand All @@ -115,7 +115,7 @@ def tiledbsoma_has_upgraded_shape(self) -> bool:
Lifecycle:
Maturing.
"""
return self._handle.tiledbsoma_has_upgraded_shape
return cast(bool, self._clib_handle.tiledbsoma_has_upgraded_shape)

@classmethod
def _dim_capacity_and_extent(
Expand Down
79 changes: 57 additions & 22 deletions apis/python/src/tiledbsoma/_dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,33 @@ class DataFrame(SOMAArray, somacore.DataFrame):
"""

_wrapper_type = DataFrameWrapper
_clib_handle_type = clib.SOMADataFrame

"""XXX comment me."""

@classmethod
def open(
cls,
uri: str,
mode: options.OpenMode = "r",
*,
tiledb_timestamp: Optional[OpenTimestamp] = None,
context: Optional[SOMATileDBContext] = None,
platform_config: Optional[options.PlatformConfig] = None,
clib_type: Optional[str] = None,
) -> Self:
"""Opens this specific type of SOMA object."""

retval = super().open(
uri,
mode,
tiledb_timestamp=tiledb_timestamp,
context=context,
platform_config=platform_config,
clib_type="SOMAArray",
)

return retval

@classmethod
def create(
Expand Down Expand Up @@ -330,11 +357,21 @@ def create(
raise map_exception_for_create(e, uri) from None

handle = cls._wrapper_type.open(uri, "w", context, tiledb_timestamp)
return cls(
clib_handle = cls._clib_handle_type.open(
uri,
clib.OpenMode.write,
context.native_context,
timestamp=(0, timestamp_ms),
)

retval = cls(
handle,
clib_handle=clib_handle,
_dont_call_this_use_create_or_open_instead="tiledbsoma-internal-code",
)

return retval

def keys(self) -> Tuple[str, ...]:
"""Returns the names of the columns when read back as a dataframe.

Expand Down Expand Up @@ -387,8 +424,8 @@ def count(self) -> int:
Maturing.
"""
self._check_open_read()
# if is it in read open mode, then it is a DataFrameWrapper
return cast(DataFrameWrapper, self._handle).count
# XXX WHY
return cast(int, self._clib_handle.count)

@property
def _maybe_soma_joinid_shape(self) -> Optional[int]:
Expand All @@ -400,7 +437,7 @@ def _maybe_soma_joinid_shape(self) -> Optional[int]:
Lifecycle:
Experimental.
"""
return self._handle.maybe_soma_joinid_shape
return cast(Optional[int], self._clib_handle.maybe_soma_joinid_shape)

@property
def _maybe_soma_joinid_maxshape(self) -> Optional[int]:
Expand All @@ -411,7 +448,7 @@ def _maybe_soma_joinid_maxshape(self) -> Optional[int]:
Lifecycle:
Experimental.
"""
return self._handle.maybe_soma_joinid_maxshape
return cast(Optional[int], self._clib_handle.maybe_soma_joinid_maxshape)

@property
def tiledbsoma_has_upgraded_domain(self) -> bool:
Expand All @@ -422,7 +459,7 @@ def tiledbsoma_has_upgraded_domain(self) -> bool:
Lifecycle:
Maturing.
"""
return self._handle.tiledbsoma_has_upgraded_domain
return cast(bool, self._clib_handle.tiledbsoma_has_upgraded_domain)

def resize_soma_joinid_shape(
self, newshape: int, check_only: bool = False
Expand All @@ -440,10 +477,10 @@ def resize_soma_joinid_shape(
if check_only:
return cast(
StatusAndReason,
self._handle._handle.can_resize_soma_joinid_shape(newshape),
self._clib_handle.can_resize_soma_joinid_shape(newshape),
)
else:
self._handle._handle.resize_soma_joinid_shape(newshape)
self._clib_handle.resize_soma_joinid_shape(newshape)
return (True, "")

def upgrade_soma_joinid_shape(
Expand All @@ -459,10 +496,10 @@ def upgrade_soma_joinid_shape(
if check_only:
return cast(
StatusAndReason,
self._handle._handle.can_upgrade_soma_joinid_shape(newshape),
self._clib_handle.can_upgrade_soma_joinid_shape(newshape),
)
else:
self._handle._handle.upgrade_soma_joinid_shape(newshape)
self._clib_handle.upgrade_soma_joinid_shape(newshape)
return (True, "")

def __len__(self) -> int:
Expand Down Expand Up @@ -536,25 +573,23 @@ def read(
_util.check_unpartitioned(partitions)
self._check_open_read()

handle = self._handle._handle

context = handle.context()
context = self._clib_handle.context()
if platform_config is not None:
config = context.tiledb_config.copy()
config.update(platform_config)
context = clib.SOMAContext(config)

sr = clib.SOMADataFrame.open(
uri=handle.uri,
uri=self._clib_handle.uri,
mode=clib.OpenMode.read,
context=context,
column_names=column_names or [],
result_order=_util.to_clib_result_order(result_order),
timestamp=handle.timestamp and (0, handle.timestamp),
timestamp=self._clib_handle.timestamp and (0, self._clib_handle.timestamp),
)

if value_filter is not None:
sr.set_condition(QueryCondition(value_filter), handle.schema)
sr.set_condition(QueryCondition(value_filter), self._clib_handle.schema)

self._set_reader_coords(sr, coords)

Expand Down Expand Up @@ -608,13 +643,11 @@ def write(
write_options = TileDBWriteOptions.from_platform_config(platform_config)
sort_coords = write_options.sort_coords

clib_dataframe = self._handle._handle

for batch in values.to_batches():
clib_dataframe.write(batch, sort_coords or False)
self._clib_handle.write(batch, sort_coords or False)

if write_options.consolidate_and_vacuum:
clib_dataframe.consolidate_and_vacuum()
self._clib_handle.consolidate_and_vacuum()

return self

Expand Down Expand Up @@ -667,11 +700,13 @@ def _set_reader_coord(
if coord.stop is None:
# There's no way to specify "to infinity" for strings.
# We have to get the nonempty domain and use that as the end.
ned = self._handle.non_empty_domain()
ned = self._clib_handle.non_empty_domain()
_, stop = ned[dim_idx]
else:
stop = coord.stop
sr.set_dim_ranges_string_or_bytes(dim.name, [(start, stop)])
# Use str(...) in case this is an Arrow string type, to satisfy
# the type-checker
sr.set_dim_ranges_string_or_bytes(dim.name, [(str(start), str(stop))])
return True

# Note: slice(None, None) matches the is_slice_of part, unless we also check the dim-type
Expand Down
34 changes: 19 additions & 15 deletions apis/python/src/tiledbsoma/_dense_nd_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,7 @@ class DenseNDArray(NDArray, somacore.DenseNDArray):
__slots__ = ()

_wrapper_type = DenseNDArrayWrapper
_clib_handle_type = clib.SOMADenseNDArray

@classmethod
def create(
Expand Down Expand Up @@ -162,8 +163,15 @@ def create(
raise map_exception_for_create(e, uri) from None

handle = cls._wrapper_type.open(uri, "w", context, tiledb_timestamp)
clib_handle = cls._clib_handle_type.open(
uri,
clib.OpenMode.write,
context.native_context,
timestamp=(0, timestamp_ms),
)
return cls(
handle,
clib_handle=clib_handle,
_dont_call_this_use_create_or_open_instead="tiledbsoma-internal-code",
)

Expand Down Expand Up @@ -216,35 +224,33 @@ def read(
#
# The only exception is if the array has been created but no data have been written at
# all, in which case the best we can do is use the schema shape.
handle: clib.SOMADenseNDArray = self._handle._handle

ned = []
for dim_name in handle.dimension_names:
for dim_name in self._clib_handle.dimension_names:
dtype = np.dtype(self.schema.field(dim_name).type.to_pandas_dtype())
slot = handle.non_empty_domain_slot_opt(dim_name, dtype)
slot = self._clib_handle.non_empty_domain_slot_opt(dim_name, dtype)
if slot is None:
use_shape = True
break
ned.append(slot[1] + 1)
else:
use_shape = False

data_shape = tuple(handle.shape if use_shape else ned)
data_shape = tuple(self._clib_handle.shape if use_shape else ned)
target_shape = dense_indices_to_shape(coords, data_shape, result_order)

context = handle.context()
context = self._clib_handle.context()
if platform_config is not None:
config = context.tiledb_config.copy()
config.update(platform_config)
context = clib.SOMAContext(config)

sr = clib.SOMADenseNDArray.open(
uri=handle.uri,
uri=self._clib_handle.uri,
mode=clib.OpenMode.read,
context=context,
column_names=[],
result_order=_util.to_clib_result_order(result_order),
timestamp=handle.timestamp and (0, handle.timestamp),
timestamp=self._clib_handle.timestamp and (0, self._clib_handle.timestamp),
)

self._set_reader_coords(sr, coords)
Expand Down Expand Up @@ -304,8 +310,6 @@ def write(
"""
_util.check_type("values", values, (pa.Tensor,))

clib_dense_array = self._handle._handle

# Compute the coordinates for the dense array.
new_coords: List[Union[int, Slice[int], None]] = []
for c in coords:
Expand All @@ -325,21 +329,21 @@ def write(
if not input.flags.contiguous:
input = np.ascontiguousarray(input)
order = clib.ResultOrder.rowmajor
clib_dense_array.reset(result_order=order)
self._set_reader_coords(clib_dense_array, new_coords)
clib_dense_array.write(input)
self._clib_handle.reset(result_order=order)
self._set_reader_coords(self._clib_handle, new_coords)
self._clib_handle.write(input)

tiledb_write_options = TileDBWriteOptions.from_platform_config(platform_config)
if tiledb_write_options.consolidate_and_vacuum:
clib_dense_array.consolidate_and_vacuum()
self._clib_handle.consolidate_and_vacuum()
return self

def resize(self, newshape: Sequence[Union[int, None]]) -> None:
"""Supported for ``SparseNDArray``; scheduled for implementation for
``DenseNDArray`` in TileDB-SOMA 1.15
"""
if clib.embedded_version_triple() >= (2, 27, 0):
self._handle.resize(newshape)
self._clib_handle.resize(newshape)
else:
raise NotImplementedError("Not implemented for libtiledbsoma < 2.27.0")

Expand Down
2 changes: 2 additions & 0 deletions apis/python/src/tiledbsoma/_experiment.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from typing_extensions import Self

from . import _tdb_handles
from . import pytiledbsoma as clib
from ._collection import Collection, CollectionBase
from ._dataframe import DataFrame
from ._indexer import IntIndexer
Expand Down Expand Up @@ -69,6 +70,7 @@ class Experiment( # type: ignore[misc] # __eq__ false positive

__slots__ = ()
_wrapper_type = _tdb_handles.ExperimentWrapper
_clib_handle_type = clib.SOMAExperiment

_subclass_constrained_soma_types = {
"obs": ("SOMADataFrame",),
Expand Down
9 changes: 3 additions & 6 deletions apis/python/src/tiledbsoma/_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
"""

from typing import (
Callable,
Dict,
Optional,
Type,
Expand Down Expand Up @@ -123,7 +122,7 @@ def open(
"""
context = _validate_soma_tiledb_context(context)
obj: SOMAObject[_Wrapper] = _open_internal( # type: ignore[valid-type]
_tdb_handles.open, uri, mode, context, tiledb_timestamp
uri, mode, context, tiledb_timestamp
)
try:
if soma_type:
Expand All @@ -144,16 +143,14 @@ def open(


def _open_internal(
opener: Callable[
[str, options.OpenMode, SOMATileDBContext, Optional[OpenTimestamp]], _Wrapper
],
uri: str,
mode: options.OpenMode,
context: SOMATileDBContext,
timestamp: Optional[OpenTimestamp],
) -> SOMAObject[_Wrapper]:
"""Lower-level open function for internal use only."""
handle = opener(uri, mode, context, timestamp)
# XXX temp cast
handle = cast(_Wrapper, _tdb_handles.open(uri, mode, context, timestamp))
try:
return reify_handle(handle)
except Exception:
Expand Down
1 change: 1 addition & 0 deletions apis/python/src/tiledbsoma/_geometry_dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ class GeometryDataFrame(somacore.GeometryDataFrame):
"""

__slots__ = ()
# XXX _clib_handle_type = clib.XXX

# Lifecycle

Expand Down
Loading