data.py

"""Definitions of data storage interfaces for SOMA implementations.

SOMA users should ordinarily not need to import this module directly; relevant
members will be exported to the ``somacore`` namespace.

Default values are provided here as a reference for implementors.
"""

from __future__ import annotations

from typing import Any, Sequence, TypeVar, Union

import pyarrow as pa
from typing_extensions import Protocol, Self, runtime_checkable

from . import base
from . import options
from .types import StatusAndReason

_RO_AUTO = options.ResultOrder.AUTO

AxisDomain = Union[None, tuple[Any, Any], list[Any]]
Domain = Sequence[AxisDomain]


@runtime_checkable
class DataFrame(base.SOMAObject, Protocol):
    """A multi-column table with a user-defined schema.

    Lifecycle: maturing
    """

    # Lifecycle

    @classmethod
    def create(
        cls,
        uri: str,
        *,
        schema: pa.Schema,
        domain: Sequence[tuple[Any, Any] | None],
        index_column_names: Sequence[str] = (options.SOMA_JOINID,),
        platform_config: options.PlatformConfig | None = None,
        context: Any | None = None,
    ) -> Self:
        """Creates a new ``DataFrame`` at the given URI.

        The schema of the created dataframe will include a column named
        ``soma_joinid`` of type ``pyarrow.int64``, with negative values
        disallowed.  If a ``soma_joinid`` column is present in the provided
        schema, it must be of the correct type.  If no ``soma_joinid`` column
        is provided, one will be added.  It may be used as an indexed column.

        Args:
            uri: The URI where the dataframe will be created.
            schema: Arrow schema defining the per-column schema. This schema
                must define all columns, including columns to be named as index
                columns.  If the schema includes types unsupported by the SOMA
                implementation, an error will be raised.
            domain: A sequence of tuples specifying the domain of each index
                column. Each tuple must be a pair consisting of the minimum
                and maximum values storable in the index column. This sequence
                must have the same length as ``index_column_names``. Use
                ``None`` for string index columns when the implementation does
                not support string domains.
            index_column_names: A list of column names to use as user-defined
                index columns (e.g., ``['cell_type', 'tissue_type']``).
                All named columns must exist in the schema, and at least one
                index column name is required.
            platform_config: platform-specific configuration; keys are SOMA
                implementation names.
            context: Other implementation-specific configuration.

        Returns:
            The newly created dataframe, opened for writing.

        Lifecycle: maturing
        """
        ...

    # Data operations

    def read(
        self,
        coords: options.SparseDFCoords = (),
        column_names: Sequence[str] | None = None,
        *,
        batch_size: options.BatchSize = options.BatchSize(),
        partitions: options.ReadPartitions | None = None,
        result_order: options.ResultOrderStr = _RO_AUTO,
        value_filter: str | None = None,
        platform_config: options.PlatformConfig | None = None,
    ) -> "ReadIter[pa.Table]":
        """Reads a user-defined slice of data into Arrow tables.

        Args:
            coords: for each index dimension, which rows to read.
                Defaults to ``()``, meaning no constraint -- all IDs.
            column_names: the named columns to read and return.
                Defaults to ``None``, meaning no constraint -- all column names.
            partitions: If present, specifies that this is part of
                a partitioned read, and which part of the data to include.
            result_order: the order to return results, specified as a
                :class:`~options.ResultOrder` or its string value.
            value_filter: an optional value filter to apply to the results.
                The default of ``None`` represents no filter. Value filter
                syntax is implementation-defined; see the documentation
                for the particular SOMA implementation for details.

        Returns:
            A :class:`ReadIter` of :class:`pa.Table`s.

        **Indexing:**

        Indexing is performed on a per-column basis for each indexed column.
        To specify dimensions:

        - A sequence of coordinates is accepted, one per indexed dimension.
        - The sequence length must be less than or equal to the number of
          indexed dimensions.
        - If the sequence is shorter than the number of indexed coordinates,
          then no constraint (i.e. ``None``) is used for the remaining
          indexed dimensions.
        - Specifying an empty sequence (e.g. ``()``, the default) represents
          no constraints over any dimension, returning the entire dataset.

        Each dimension may be indexed as follows:

        - ``None`` or ``slice(None)`` places no constraint on the dimension.
        - Coordinates can be specified as a scalar value, a Python sequence
          (``list``, ``tuple``, etc.), a NumPy ndarray, an Arrow array, or
          similar objects (as defined by ``SparseDFCoords``).
        - Slices specify a closed range: ``slice(2, 4)`` includes both 2 and 4.
          Slice *steps* may not be used: ``slice(10, 20, 2)`` is invalid.
          ``slice(None)`` places no constraint on the dimension. Half-specified
          slices like ``slice(None, 99)`` and ``slice(5, None)`` specify
          all indices up to and including the value, and all indices
          starting from and including the value.
        - Negative values in indices and slices are treated as raw domain values
          and not as indices relative to the end, unlike traditional Python
          sequence indexing. For instance, ``slice(-10, 3)`` indicates the range
          from −10 to 3 on the given dimension.

        Lifecycle: maturing
        """
        ...

    def change_domain(
        self,
        newdomain: Domain,
        check_only: bool = False,
    ) -> StatusAndReason:
        """Allows you to enlarge the domain of a SOMA :class:`DataFrame`, when
        the ``DataFrame`` already has a domain.

        The argument must be a tuple of pairs of low/high values for the desired
        domain, one pair per index column. For string index columns, you must
        offer the low/high pair as `("", "")`, or as ``None`` for string columns.  If
        ``check_only`` is ``True``, returns whether the operation would succeed if
        attempted, and a reason why it would not.

        For example, suppose the dataframe's sole index-column name is
        ``"soma_joinid"`` (which is the default at create).  If the dataframe's
        ``.maxdomain`` is ``((0, 999999),)`` and its ``.domain`` is ``((0,
        2899),)``, this means that ``soma_joinid`` values between 0 and 2899 can
        be read or written; any attempt to read or write ``soma_joinid`` values
        outside this range will result in an error. If you then apply
        ``.change_domain([(0, 5700)])``, then ``.domain`` will
        report ``((0, 5699),)``, and now ``soma_joinid`` values in the range 0
        to 5699 can now be written to the dataframe.

        If you use non-default ``index_column_names`` in the dataframe's
        ``create`` then you need to specify the (low, high) pairs for each
        index column. For example, if the dataframe's ``index_column_names``
        is ``["soma_joinid", "cell_type"]``, then you can upgrade domain using
        ``[(0, 5699), ("", "")]``.

        Lastly, it is an error to try to set the ``domain`` to be smaller than
        ``maxdomain`` along any index column.  The ``maxdomain`` of a dataframe is
        set at creation time, and cannot be extended afterward.

        Lifecycle: maturing
        """
        ...

    def write(
        self,
        values: pa.RecordBatch | pa.Table,
        *,
        platform_config: options.PlatformConfig | None = None,
    ) -> Self:
        """Writes the data from an Arrow table to the persistent object.

        As duplicate index values are not allowed, index values already present
        in the object are overwritten and new index values are added.

        Args:
            values: An Arrow table containing all columns, including
                the index columns. The schema for the values must match
                the schema for the ``DataFrame``.

        Returns:
            ``self``, to enable method chaining.

        Lifecycle: maturing
        """
        ...

    # Metadata operations

    @property
    def schema(self) -> pa.Schema:
        """The schema of the data in this dataframe.

        Lifecycle: maturing
        """
        ...

    @property
    def index_column_names(self) -> tuple[str, ...]:
        """The names of the index (dimension) columns.

        Lifecycle: maturing
        """
        ...

    @property
    def domain(self) -> tuple[tuple[Any, Any], ...]:
        """The allowable range of values in each index column.

        Returns:
            A tuple of minimum and maximum values, inclusive,
            storable on each index column of the dataframe.

        Lifecycle: maturing
        """
        ...


@runtime_checkable
class NDArray(base.SOMAObject, Protocol):
    """Common behaviors of N-dimensional arrays of a single primitive type."""

    # Lifecycle

    @classmethod
    def create(
        cls,
        uri: str,
        *,
        type: pa.DataType,
        shape: Sequence[int],
        platform_config: options.PlatformConfig | None = None,
        context: Any | None = None,
    ) -> Self:
        """Creates a new ND array of the current type at the given URI.

        Args:
            uri: The URI where the array will be created.
            type: The Arrow type to store in the array.
                If the type is unsupported, an error will be raised.
            shape: The initial maximum capacity of each dimension, specified
                as one element per dimension, e.g. ``(100, 10)`` for a 2D array.
                All lengths must be in the positive int64 range. The shape can
                be increased after creation using :meth:`resize`.

        Returns:
            The newly created array, opened for writing.

        Lifecycle: maturing
        """
        ...

    def resize(self, newshape: Sequence[int | None], check_only: bool = False) -> StatusAndReason:
        """Increases the shape of the array as specified.

        Lifecycle: maturing
        """
        ...

    # Metadata operations

    @property
    def shape(self) -> tuple[int, ...]:
        """The maximum capacity (domain) of each dimension of this array.

        Lifecycle: maturing
        """
        ...

    @property
    def ndim(self) -> int:
        """The number of dimensions in this array.

        Lifecycle: maturing
        """
        ...

    @property
    def schema(self) -> pa.Schema:
        """The schema of the data in this array.

        Lifecycle: maturing
        """
        ...


@runtime_checkable
class DenseNDArray(NDArray, Protocol):
    """An N-dimensional array stored densely.

    Lifecycle: maturing
    """

    def read(
        self,
        coords: options.DenseNDCoords = (),
        *,
        partitions: options.ReadPartitions | None = None,
        result_order: options.ResultOrderStr = _RO_AUTO,
        platform_config: options.PlatformConfig | None = None,
    ) -> pa.Tensor:
        """Reads the specified subarray as a Tensor.

        Coordinates must specify a contiguous subarray, and the number of
        coordinates must be less than or equal to the number of dimensions.
        For example, if the array is 10×20, acceptable values of ``coords``
        include ``()``, ``(3, 4)``, ``[slice(5, 10)]``, and
        ``[slice(5, 10), slice(6, 12)]``.

        Args:
            coords: A per-dimension sequence of coordinates defining
                the range to read.
            partitions: If present, specifies that this is part of
                a partitioned read, and which part of the data to include.
            result_order: the order to return results, specified as a
                :class:`~options.ResultOrder` or its string value.
            platform_config: platform-specific configuration; keys are SOMA
                implementation names.

        Returns:
            The data over the requested range as a tensor.

        **Indexing:**

        Indexing is performed on a per-dimension basis.

        - A sequence of coordinates is accepted, one per dimension.
        - The sequence length must be less than or equal to
          the number of dimensions.
        - If the sequence is shorter than the number of dimensions, the
          remaining dimensions are unconstrained.
        - Specifying an empty sequence (e.g. ``()``, the default) represents
          no constraints over any dimension, returning the entire dataset.

        Each dimension may be indexed by value or slice:

        - Slices specify a closed range: ``slice(2, 4)`` includes 2, 3, and 4.
          Slice *steps* may not be used: ``slice(10, 20, 2)`` is invalid.
          ``slice(None)`` places no constraint on the dimension. Half-specified
          slices like ``slice(None, 99)`` and ``slice(5, None)`` specify
          all indices up to and including the value, and all indices
          starting from and including the value.
        - Negative indexing is not supported.

        Lifecycle: maturing
        """
        ...

    def write(
        self,
        coords: options.DenseNDCoords,
        values: pa.Tensor,
        *,
        platform_config: options.PlatformConfig | None = None,
    ) -> Self:
        """Writes an Arrow tensor to a subarray of the persistent object.

        The subarray written is defined by ``coords`` and ``values``. This will
        overwrite existing values in the array.

        Args:
            coords: A per-dimension tuple of scalars or slices
                defining the bounds of the subarray to be written.
                See :meth:`read` for details about indexing.
            values: The values to be written to the subarray.  Must have
                the same shape as ``coords``, and matching type to the array.

        Returns:
            ``self``, to enable method chaining.

        Lifecycle: maturing
        """
        ...


@runtime_checkable
class SparseNDArray(NDArray, Protocol):
    """A N-dimensional array stored sparsely.

    Lifecycle: maturing
    """

    def read(
        self,
        coords: options.SparseNDCoords = (),
        *,
        batch_size: options.BatchSize = options.BatchSize(),
        partitions: options.ReadPartitions | None = None,
        result_order: options.ResultOrderStr = _RO_AUTO,
        platform_config: options.PlatformConfig | None = None,
    ) -> "SparseRead":
        """Reads the specified subarray in batches.

        Values returned are a :class:`SparseRead` object which can be converted
        to any number of formats::

            some_dense_array.read(...).tables()
            # -> an iterator of Arrow Tables

        Args:
            coords: A per-dimension sequence of coordinates defining
                the range to be read.
            batch_size: The size of batches that should be returned from a read.
                See :class:`options.BatchSize` for details.
            partitions: Specifies that this is part of a partitioned read,
                and which partition to include, if present.
            result_order: the order to return results, specified as a
                :class:`~options.ResultOrder` or its string value.
            platform_config: platform-specific configuration; keys are SOMA
                implementation names.

        Returns:
            The data that was requested in a :class:`SparseRead`,
            allowing access in any supported format.

        **Indexing:**

        Indexing is performed on a per-dimension basis.

        - A sequence of coordinates is accepted, one per dimension.
        - The sequence length must be less than or equal to
          the number of dimensions.
        - If the sequence is shorter than the number of dimensions, the
          remaining dimensions are unconstrained.
        - Specifying an empty sequence (e.g. ``()``, the default) represents
          no constraints over any dimension, returning the entire dataset.

        Each dimension may be indexed as follows:

        - ``None`` or ``slice(None)`` places no constraint on the dimension.
        - Coordinates can be specified as a scalar value, a Python sequence
          (``list``, ``tuple``, etc.), a ``ndarray``, an Arrow array, and
          similar objects (as defined by ``SparseNDCoords``).
        - Slices specify a closed range: ``slice(2, 4)`` includes 2, 3, and 4.
          Slice *steps* may not be used: ``slice(10, 20, 2)`` is invalid.
          ``slice(None)`` places no constraint on the dimension. Half-specified
          slices like ``slice(None, 99)`` and ``slice(5, None)`` specify
          all indices up to and including the value, and all indices
          starting from and including the value.
        - Negative indexing is not supported.

        Lifecycle: maturing
        """
        ...

    def write(
        self,
        values: pa.SparseCSCMatrix | pa.SparseCSRMatrix | pa.SparseCOOTensor | pa.Table,
        *,
        platform_config: options.PlatformConfig | None = None,
    ) -> Self:
        """Writes a Tensor to a subarray of the persistent object.

        Args:
            values: The values to write to the array. Supported types are:

                Arrow sparse tensor: the coordinates in the tensor are
                interpreted as the coordinates to write to.  Supports the
                *experimental* types SparseCOOTensor, SparseCSRMatrix, and
                SparseCSCMatrix. There is currently no support for
                SparseCSFTensor or dense Tensor.

                Arrow table: a COO table, with columns named ``soma_dim_0``,
                    ..., ``soma_dim_N`` and ``soma_data``.
            platform_config: platform-specific configuration; keys are SOMA
                implementation names.

        Returns:
            ``self``, to enable method chaining.

        Lifecycle: maturing
        """
        ...

    @property
    def nnz(self) -> int:
        """The number of values stored in the array, including explicit zeros.

        Lifecycle: maturing
        """
        ...


#
# Read types
#

_T = TypeVar("_T", covariant=True)


# Sparse reads are returned as an iterable structure:


@runtime_checkable
class ReadIter(Protocol[_T]):
    """SparseRead result iterator allowing users to flatten the iteration.

    Lifecycle: maturing
    """

    def __next__(self) -> _T: ...

    def __iter__(self) -> _T: ...

    def concat(self) -> _T:
        """Returns all the requested data in a single operation.

        If some data has already been retrieved using ``next``, this will return
        the remaining data, excluding that which as already been returned.

        Lifecycle: maturing
        """
        ...


@runtime_checkable
class SparseRead(Protocol):
    """Intermediate type to choose result format when reading a sparse array.

    A query may not be able to return all of these formats. The concrete result
    may raise a ``NotImplementedError`` or may choose to raise a different
    exception (likely a ``TypeError``) containing more specific information
    about why the given format is not supported.

    Lifecycle: maturing
    """

    def coos(self) -> ReadIter[pa.SparseCOOTensor]: ...

    def dense_tensors(self) -> ReadIter[pa.Tensor]: ...

    def record_batches(self) -> ReadIter[pa.RecordBatch]: ...

    def tables(self) -> ReadIter[pa.Table]: ...