cellarr-array 0.0.3__py3-none-any.whl → 0.2.0__py3-none-any.whl

cellarr_array/__init__.py

@@ -15,7 +15,5 @@ except PackageNotFoundError:  # pragma: no cover
 finally:
     del version, PackageNotFoundError
 
-from .config import CellArrConfig, ConsolidationConfig
-from .DenseCellArray import DenseCellArray
-from .SparseCellArray import SparseCellArray
-from .helpers import create_cellarray, SliceHelper
+from .core import DenseCellArray, SparseCellArray
+from .utils import CellArrConfig, ConsolidationConfig, create_cellarray

cellarr_array/core/__init__.py

@@ -0,0 +1,3 @@
+from .base import CellArray
+from .dense import DenseCellArray
+from .sparse import SparseCellArray

cellarr_array/core/base.py

@@ -0,0 +1,344 @@
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+
+try:
+    from types import EllipsisType
+except ImportError:
+    # TODO: This is required for Python <3.10. Remove once Python 3.9 reaches EOL in October 2025
+    EllipsisType = type(...)
+from typing import Any, List, Literal, Optional, Tuple, Union
+
+import numpy as np
+import tiledb
+from scipy import sparse
+
+from ..utils.config import ConsolidationConfig
+from .helpers import SliceHelper
+
+__author__ = "Jayaram Kancherla"
+__copyright__ = "Jayaram Kancherla"
+__license__ = "MIT"
+
+
+class CellArray(ABC):
+    """Abstract base class for TileDB array operations."""
+
+    def __init__(
+        self,
+        uri: Optional[str] = None,
+        tiledb_array_obj: Optional[tiledb.Array] = None,
+        attr: str = "data",
+        mode: Optional[Literal["r", "w", "d", "m"]] = None,
+        config_or_context: Optional[Union[tiledb.Config, tiledb.Ctx]] = None,
+        validate: bool = True,
+    ):
+        """Initialize the object.
+
+        Args:
+            uri:
+                URI to the array.
+                Required if 'tiledb_array_obj' is not provided.
+
+            tiledb_array_obj:
+                Optional, an already opened ``tiledb.Array`` instance.
+                If provided, 'uri' can be None, and 'config_or_context' is ignored.
+
+            attr:
+                Attribute to access.
+                Defaults to "data".
+
+            mode:
+                Open the array object in read 'r', write 'w', modify
+                'm' mode, or delete 'd' mode.
+
+                Defaults to None for automatic mode switching.
+
+                If 'tiledb_array_obj' is provided, this mode should ideally match
+                the mode of the provided array or be None.
+
+            config_or_context:
+                Optional config or context object. Ignored if 'tiledb_array_obj' is provided,
+                as context will be derived from the object.
+
+                Defaults to None.
+
+            validate:
+                Whether to validate the attributes.
+                Defaults to True.
+        """
+        self._array_passed_in = False
+        self._opened_array_external = None
+        self._ctx = None
+
+        if tiledb_array_obj is not None:
+            if not isinstance(tiledb_array_obj, tiledb.Array):
+                raise ValueError("'tiledb_array_obj' must be a tiledb.Array instance.")
+
+            if not tiledb_array_obj.isopen:
+                # Option 1: Raise error
+                raise ValueError("If 'tiledb_array_obj' is provided, it must be an open tiledb.Array instance.")
+                # Option 2: Try to reopen (less safe as we don't know original intent)
+                # try:
+                #     tiledb_array_obj.reopen()
+                # except tiledb.TileDBError as e:
+                #     raise ValueError(
+                #         f"Provided 'tiledb_array_obj' is closed and could not be reopened: {e}"
+                #     )
+
+            self.uri = tiledb_array_obj.uri
+            self._array_passed_in = True
+            self._opened_array_external = tiledb_array_obj
+
+            # infer mode if possible, or require it matches
+            if mode is not None and tiledb_array_obj.mode != mode:
+                # we could try to reopen with the desired mode
+                raise ValueError(
+                    f"Provided array mode '{tiledb_array_obj.mode}' does not match requested mode '{mode}'.",
+                    "Re-open the external array with the desired mode or pass matching mode.",
+                )
+
+            self._mode = tiledb_array_obj.mode
+            self._ctx = tiledb_array_obj.ctx
+        elif uri is not None:
+            self.uri = uri
+            self._mode = mode
+            self._array_passed_in = False
+            self._opened_array_external = None
+
+            if config_or_context is None:
+                self._ctx = None
+            elif isinstance(config_or_context, tiledb.Config):
+                self._ctx = tiledb.Ctx(config_or_context)
+            elif isinstance(config_or_context, tiledb.Ctx):
+                self._ctx = config_or_context
+            else:
+                raise TypeError("'config_or_context' must be a TileDB Config or Ctx object.")
+        else:
+            raise ValueError("Either 'uri' or 'tiledb_array_obj' must be provided.")
+
+        self._shape = None
+        self._ndim = None
+        self._dim_names = None
+        self._attr_names = None
+        self._nonempty_domain = None
+
+        if validate:
+            self._validate(attr=attr)
+
+        self._attr = attr
+
+    def _validate(self, attr):
+        with self.open_array(mode="r") as A:
+            schema = A.schema
+            if schema.ndim > 2:
+                raise ValueError("Only 1D and 2D arrays are supported.")
+
+            current_attr_names = [schema.attr(i).name for i in range(schema.nattr)]
+            if attr not in current_attr_names:
+                raise ValueError(
+                    f"Attribute '{attr}' does not exist in the array. Available attributes: {current_attr_names}."
+                )
+
+    @property
+    def mode(self) -> Optional[str]:
+        """Get current array mode. If an external array is used, this is its open mode."""
+        if self._array_passed_in and self._opened_array_external is not None:
+            return self._opened_array_external.mode
+        return self._mode
+
+    @mode.setter
+    def mode(self, value: Optional[str]):
+        """Set array mode for subsequent operations if not using an external array.
+
+        This action does not affect an already passed-in external array's mode.
+        """
+        if self._array_passed_in:
+            # To change mode of an external array, user must reopen it and pass it again.
+            current_ext_mode = self._opened_array_external.mode if self._opened_array_external else "unknown"
+            if value != current_ext_mode:
+                raise ValueError(
+                    f"Cannot change mode of an externally managed array (current: {current_ext_mode}). "
+                    "Re-open the external array with the new mode and re-initialize CellArray."
+                )
+        if value is not None and value not in ["r", "w", "m", "d"]:
+            raise ValueError("Mode must be one of: None, 'r', 'w', 'm', 'd'")
+
+        self._mode = value
+
+    @property
+    def dim_names(self) -> List[str]:
+        """Get dimension names of the array."""
+        if self._dim_names is None:
+            with self.open_array(mode="r") as A:
+                self._dim_names = [dim.name for dim in A.schema.domain]
+        return self._dim_names
+
+    @property
+    def attr_names(self) -> List[str]:
+        """Get attribute names of the array."""
+        if self._attr_names is None:
+            with self.open_array(mode="r") as A:
+                self._attr_names = [A.schema.attr(i).name for i in range(A.schema.nattr)]
+        return self._attr_names
+
+    @property
+    def shape(self) -> Tuple[int, ...]:
+        if self._shape is None:
+            with self.open_array(mode="r") as A:
+                self._shape = tuple(int(dim.domain[1] - dim.domain[0] + 1) for dim in A.schema.domain)
+        return self._shape
+
+    @property
+    def nonempty_domain(self) -> Optional[Tuple[Any, ...]]:
+        if self._nonempty_domain is None:
+            with self.open_array(mode="r") as A:
+                # nonempty_domain() can return None if the array is empty.
+                ned = A.nonempty_domain()
+                if ned is None:
+                    self._nonempty_domain = None
+                else:
+                    self._nonempty_domain = tuple(ned) if isinstance(ned[0], tuple) else (ned,)
+        return self._nonempty_domain
+
+    @property
+    def ndim(self) -> int:
+        """Get number of dimensions."""
+        if self._ndim is None:
+            with self.open_array(mode="r") as A:
+                self._ndim = A.schema.ndim
+                # self._ndim = len(self.shape)
+        return self._ndim
+
+    @contextmanager
+    def open_array(self, mode: Optional[str] = None):
+        """Context manager for array operations.
+
+        Uses the externally provided array if available, otherwise opens from URI.
+
+        Args:
+            mode:
+                Desired mode for the operation ('r', 'w', 'm', 'd').
+                If an external array is used, this mode must be compatible with
+                (or same as) the mode the external array was opened with.
+
+                If None, uses the CellArray's default mode.
+        """
+        if self._array_passed_in and self._opened_array_external is not None:
+            if not self._opened_array_external.isopen:
+                # Attempt to reopen if closed. This assumes the user might have closed it
+                # and expects CellArr to reopen it if still possible.
+                try:
+                    self._opened_array_external.reopen()
+                except Exception as e:
+                    raise tiledb.TileDBError(
+                        f"Externally provided array is closed and could not be reopened: {e}"
+                    ) from e
+
+            effective_mode = mode if mode is not None else self._opened_array_external.mode
+
+            current_external_mode = self._opened_array_external.mode
+            if effective_mode == "r" and current_external_mode not in ["r", "w", "m"]:
+                # Read ops ok on write/modify modes
+                pass
+            elif effective_mode in ["w", "d"] and current_external_mode != effective_mode:
+                raise tiledb.TileDBError(
+                    f"Requested operation mode '{effective_mode}' is incompatible with the "
+                    f"externally provided array's mode '{current_external_mode}'. "
+                    "Ensure the external array is opened in a compatible mode."
+                )
+
+            # DO NOT close self._opened_array_external here; its lifecycle is managed by the user.
+            yield self._opened_array_external
+        else:
+            effective_mode = mode if mode is not None else self.mode
+            effective_mode = effective_mode if effective_mode is not None else "r"
+            array = tiledb.open(self.uri, mode=effective_mode, ctx=self._ctx)
+
+            try:
+                yield array
+            finally:
+                array.close()
+
+    def __getitem__(self, key: Union[slice, EllipsisType, Tuple[Union[slice, List[int]], ...], EllipsisType]):
+        """Get item implementation that routes to either direct slicing or multi_index
+        based on the type of indices provided.
+
+        Args:
+            key:
+                Slice or list of indices for each dimension in the array.
+        """
+        if not isinstance(key, tuple):
+            key = (key,)
+
+        if len(key) > self.ndim:
+            raise IndexError(f"Invalid number of dimensions: got {len(key)}, expected {self.ndim}")
+
+        # Normalize all indices
+        normalized_key = tuple(SliceHelper.normalize_index(idx, self.shape[i]) for i, idx in enumerate(key))
+
+        num_ellipsis = sum(isinstance(i, EllipsisType) for i in normalized_key)
+        if num_ellipsis > 1:
+            raise IndexError(f"Found more than 1 Ellipsis (...) in key: {normalized_key}")
+
+        # Check if we can use direct slicing
+        use_direct = all(isinstance(idx, (slice, EllipsisType)) for idx in normalized_key)
+
+        if use_direct:
+            return self._direct_slice(normalized_key)
+        else:
+            if num_ellipsis > 0:
+                raise IndexError(f"tiledb does not support ellipsis in multi-index access: {normalized_key}")
+            return self._multi_index(normalized_key)
+
+    @abstractmethod
+    def _direct_slice(self, key: Tuple[Union[slice, EllipsisType], ...]) -> np.ndarray:
+        """Implementation for direct slicing."""
+        pass
+
+    @abstractmethod
+    def _multi_index(self, key: Tuple[Union[slice, List[int]], ...]) -> np.ndarray:
+        """Implementation for multi-index access."""
+        pass
+
+    def vacuum(self) -> None:
+        """Remove deleted fragments from the array."""
+        tiledb.vacuum(self.uri)
+
+    def consolidate(self, config: Optional[ConsolidationConfig] = None) -> None:
+        """Consolidate array fragments.
+
+        Args:
+            config:
+                Optional consolidation configuration.
+        """
+        if config is None:
+            config = ConsolidationConfig()
+
+        consolidation_cfg = tiledb.Config()
+
+        consolidation_cfg["sm.consolidation.steps"] = config.steps
+        consolidation_cfg["sm.consolidation.step_min_frags"] = config.step_min_frags
+        consolidation_cfg["sm.consolidation.step_max_frags"] = config.step_max_frags
+        consolidation_cfg["sm.consolidation.buffer_size"] = config.buffer_size
+        consolidation_cfg["sm.mem.total_budget"] = config.total_budget
+
+        tiledb.consolidate(self.uri, config=consolidation_cfg)
+
+        if config.vacuum_after:
+            self.vacuum()
+
+    @abstractmethod
+    def write_batch(self, data: Union[np.ndarray, sparse.spmatrix], start_row: int, **kwargs) -> None:
+        """Write a batch of data to the array starting at the specified row.
+
+        Args:
+            data:
+                Data to write (numpy array for dense, scipy sparse matrix for sparse).
+
+            start_row:
+                Starting row index for writing.
+
+            **kwargs:
+                Additional arguments for write operation.
+        """
+        pass

cellarr_array/{DenseCellArray.py → core/dense.py}

@@ -7,7 +7,7 @@ from typing import List, Tuple, Union
 
 import numpy as np
 
-from .CellArray import CellArray
+from .base import CellArray
 from .helpers import SliceHelper
 
 __author__ = "Jayaram Kancherla"
@@ -92,7 +92,6 @@ class DenseCellArray(CellArray):
         if len(data.shape) != self.ndim:
             raise ValueError(f"Data dimensions {data.shape} don't match array dimensions {self.shape}.")
 
-        # Check bounds
         end_row = start_row + data.shape[0]
         if end_row > self.shape[0]:
             raise ValueError(
@@ -102,7 +101,6 @@ class DenseCellArray(CellArray):
         if self.ndim == 2 and data.shape[1] != self.shape[1]:
             raise ValueError(f"Data columns {data.shape[1]} don't match array columns {self.shape[1]}.")
 
-        # Construct write region
         if self.ndim == 1:
             write_region = slice(start_row, end_row)
         else:  # 2D
@@ -110,4 +108,5 @@ class DenseCellArray(CellArray):
 
         # write_data = {self._attr: data} if len(self.attr_names) > 1 else data
         with self.open_array(mode="w") as array:
+            print("write_region", write_region)
             array[write_region] = data

cellarr_array/{helpers.py → core/helpers.py}

@@ -8,7 +8,7 @@ from typing import List, Optional, Tuple, Union
 import numpy as np
 import tiledb
 
-from .config import CellArrConfig
+from ..utils.config import CellArrConfig
 
 __author__ = "Jayaram Kancherla"
 __copyright__ = "Jayaram Kancherla"
@@ -52,7 +52,7 @@ def create_cellarray(
             Optional list of dimension names.
 
         dim_dtypes:
-            Optional list of dimension dtypes.
+            Optional list of dimension dtypes. Defaults to numpy's uint32.
 
         attr_name:
             Name of the data attribute.
@@ -67,29 +67,28 @@ def create_cellarray(
         ValueError: If dimensions are invalid or inputs are inconsistent.
     """
     config = config or CellArrConfig()
+    tiledb_ctx = tiledb.Config(config.ctx_config) if config.ctx_config else None
 
     if attr_dtype is None:
         attr_dtype = np.float32
     if isinstance(attr_dtype, str):
         attr_dtype = np.dtype(attr_dtype)
 
-    # Require either shape or dim_dtypes
     if shape is None and dim_dtypes is None:
         raise ValueError("Either 'shape' or 'dim_dtypes' must be provided.")
 
     if shape is not None:
         if len(shape) not in (1, 2):
-            raise ValueError("Only 1D and 2D arrays are supported.")
+            raise ValueError("Shape must have 1 or 2 dimensions.")
 
     # Set dimension dtypes, defaults to numpy uint32
     if dim_dtypes is None:
         dim_dtypes = [np.uint32] * len(shape)
     else:
         if len(dim_dtypes) not in (1, 2):
-            raise ValueError("Only 1D and 2D arrays are supported.")
+            raise ValueError("Array must have 1 or 2 dimensions.")
         dim_dtypes = [np.dtype(dt) if isinstance(dt, str) else dt for dt in dim_dtypes]
 
-    # Calculate shape from dtypes if needed
     if shape is None:
         shape = tuple(np.iinfo(dt).max if np.issubdtype(dt, np.integer) else None for dt in dim_dtypes)
     if None in shape:
@@ -97,7 +96,6 @@ def create_cellarray(
             np.iinfo(dt).max if s is None and np.issubdtype(dt, np.integer) else s for s, dt in zip(shape, dim_dtypes)
         )
 
-    # Set dimension names
     if dim_names is None:
         dim_names = [f"dim_{i}" for i in range(len(shape))]
 
@@ -107,37 +105,44 @@ def create_cellarray(
 
     dom = tiledb.Domain(
         *[
-            tiledb.Dim(name=name, domain=(0, s - 1), tile=min(s, config.tile_capacity), dtype=dt)
+            tiledb.Dim(
+                name=name,
+                # supporting empty dimensions
+                domain=(0, 0 if s == 0 else s - 1),
+                tile=min(1 if s == 0 else s // 2, config.tile_capacity // 2),
+                dtype=dt,
+            )
             for name, s, dt in zip(dim_names, shape, dim_dtypes)
         ],
-        ctx=tiledb.Ctx(config.ctx_config),
+        ctx=tiledb_ctx,
     )
-
-    attr = tiledb.Attr(
+    attr_obj = tiledb.Attr(
         name=attr_name,
         dtype=attr_dtype,
         filters=config.attrs_filters.get(attr_name, config.attrs_filters.get("", None)),
+        ctx=tiledb_ctx,
     )
-
     schema = tiledb.ArraySchema(
         domain=dom,
-        attrs=[attr],
+        attrs=[attr_obj],
         cell_order=config.cell_order,
         tile_order=config.tile_order,
         sparse=sparse,
         coords_filters=config.coords_filters,
         offsets_filters=config.offsets_filters,
-        ctx=tiledb.Ctx(config.ctx_config),
+        ctx=tiledb_ctx,
     )
-
-    tiledb.Array.create(uri, schema)
+    tiledb.Array.create(uri, schema, ctx=tiledb_ctx)
 
     # Import here to avoid circular imports
-    from .DenseCellArray import DenseCellArray
-    from .SparseCellArray import SparseCellArray
+    from .dense import DenseCellArray
+    from .sparse import SparseCellArray
 
-    # Return appropriate array type
-    return SparseCellArray(uri, attr=attr_name, mode=mode) if sparse else DenseCellArray(uri, attr=attr_name, mode=mode)
+    return (
+        SparseCellArray(uri=uri, attr=attr_name, mode=mode, config_or_context=tiledb_ctx)
+        if sparse
+        else DenseCellArray(uri=uri, attr=attr_name, mode=mode, config_or_context=tiledb_ctx)
+    )
 
 
 class SliceHelper:
@@ -145,19 +150,27 @@ class SliceHelper:
 
     @staticmethod
     def is_contiguous_indices(indices: List[int]) -> Optional[slice]:
-        """Check if indices can be represented as a contiguous slice."""
         if not indices:
             return None
 
-        diffs = np.diff(indices)
+        sorted_indices = sorted(list(set(indices)))
+        if not sorted_indices:
+            return None
+
+        if len(sorted_indices) == 1:
+            return slice(sorted_indices[0], sorted_indices[0] + 1, None)
+
+        diffs = np.diff(sorted_indices)
         if np.all(diffs == 1):
-            return slice(indices[0], indices[-1] + 1, None)
+            return slice(sorted_indices[0], sorted_indices[-1] + 1, None)
+
         return None
 
     @staticmethod
-    def normalize_index(idx: Union[int, slice, List[int]], dim_size: int) -> Union[slice, List[int], EllipsisType]:
+    def normalize_index(
+        idx: Union[int, range, slice, List[int], EllipsisType], dim_size: int
+    ) -> Union[slice, List[int], EllipsisType]:
         """Normalize index to handle negative indices and ensure consistency."""
-
         if isinstance(idx, EllipsisType):
             return idx
 
@@ -166,36 +179,61 @@ class SliceHelper:
             idx = slice(idx.start, idx.stop, idx.step)
 
         if isinstance(idx, slice):
-            start = idx.start if idx.start is not None else 0
-            stop = idx.stop if idx.stop is not None else dim_size
+            start = idx.start
+            stop = idx.stop
             step = idx.step
 
+            # Resolve None to full dimension slice parts
+            if start is None:
+                start = 0
+
+            if stop is None:
+                stop = dim_size
+
             # Handle negative indices
             if start < 0:
-                start = dim_size + start
-
+                start += dim_size
             if stop < 0:
-                stop = dim_size + stop
+                stop += dim_size
 
-            if start < 0 or start > dim_size:
-                raise IndexError(f"Start index {start} out of bounds for dimension size {dim_size}")
-            if stop < 0 or stop > dim_size:
-                raise IndexError(f"Stop index {stop} out of bounds for dimension size {dim_size}")
+            # slice allows start > dim_size or stop < 0 to result in empty slices.
+            # Note: start == dim_size is OK for empty slice like arr[dim_size:]
+            if start < 0 or (start >= dim_size and dim_size > 0):
+                if not (start == dim_size and (step is None or step > 0)):
+                    if start >= dim_size:
+                        raise IndexError(
+                            f"Start index {idx.start if idx.start is not None else 'None'} results in {start}, which is out of bounds for dimension size {dim_size}."
+                        )
 
-            return slice(start, stop, step)
+            # Clamping slice arguments to dimensions
+            stop = min(stop, dim_size)
+            start = max(0, start)
 
+            return slice(start, stop, step)
         elif isinstance(idx, list):
+            if not idx:
+                return []
+
             norm_idx = [i if i >= 0 else dim_size + i for i in idx]
             if any(i < 0 or i >= dim_size for i in norm_idx):
-                raise IndexError(f"List indices {idx} out of bounds for dimension size {dim_size}")
-            return norm_idx
-
-        else:  # Single integer index
-            norm_idx = idx if idx >= 0 else dim_size + idx
-
-            if norm_idx < 0 or norm_idx >= dim_size:
+                oob_indices = [orig_i for orig_i, norm_i in zip(idx, norm_idx) if not (0 <= norm_i < dim_size)]
+                raise IndexError(
+                    f"List indices {oob_indices} (original values) are out of bounds for dimension size {dim_size}."
+                )
+
+            # TileDB multi_index usually returns data sorted by coordinates
+            return sorted(list(set(norm_idx)))
+        elif isinstance(idx, (int, np.integer)):
+            norm_idx = int(idx)
+            if norm_idx < 0:
+                norm_idx += dim_size
+
+            if not (0 <= norm_idx < dim_size):
                 raise IndexError(f"Index {idx} out of bounds for dimension size {dim_size}")
+
             return slice(norm_idx, norm_idx + 1, None)
+        else:
+            raise TypeError(f"Index type {type(idx)} not supported for normalization.")
 
 
 def create_group(output_path, group_name):