feat: Sampler interface and ability to add custom samplers

docs/api.md

@@ -14,6 +14,7 @@
 
     Loader
     Loader.__iter__
+    ChunkSampler
 ```
 
 (io-helpers)=
@@ -28,6 +29,15 @@
     DatasetCollection
 ```
 
+(abc)=
+## abc
+```{eval-rst}
+.. autosummary::
+   :toctree: generated/
+
+    abc.Sampler
+```
+
 (types)=
 ## types
 
@@ -36,4 +46,5 @@
    :toctree: generated/
 
     types.LoaderOutput
+    types.LoadRequest
 ```

pyproject.toml

@@ -178,6 +178,10 @@ omit = [
   "**/test_*.py",
 ]
 
+[[tool.mypy.overrides]]
+module = [ "anndata.*", "cupyx.*", "cupy.*" ]
+ignore_missing_imports = true
+
 [tool.cruft]
 skip = [
   "tests",

src/annbatch/__init__.py

@@ -2,10 +2,18 @@
 
 from importlib.metadata import version
 
-from . import types
+from . import abc, types
 from .io import DatasetCollection, write_sharded
 from .loader import Loader
+from .samplers._chunk_sampler import ChunkSampler
 
 __version__ = version("annbatch")
 
-__all__ = ["Loader", "write_sharded", "DatasetCollection", "types"]
+__all__ = [
+    "Loader",
+    "DatasetCollection",
+    "types",
+    "write_sharded",
+    "ChunkSampler",
+    "abc",
+]

src/annbatch/abc/__init__.py

@@ -0,0 +1,5 @@
+from .sampler import Sampler
+
+__all__ = [
+    "Sampler",
+]

src/annbatch/abc/sampler.py

@@ -0,0 +1,70 @@
+"""Sampler classes for efficient chunk-based data access."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from annbatch.types import LoadRequest
+
+
+class Sampler(ABC):
+    """Base sampler class.
+
+    Samplers control how data is batched and loaded from the underlying datasets.
+    """
+
+    def sample(self, n_obs: int) -> Iterator[LoadRequest]:
+        """Sample load requests given the total number of observations.
+
+        Parameters
+        ----------
+        n_obs
+            The total number of observations available.
+
+        Yields
+        ------
+        LoadRequest
+            Load requests for batching data.
+        """
+        self.validate(n_obs)
+        yield from self._sample(n_obs)
+
+    @abstractmethod
+    def validate(self, n_obs: int) -> None:
+        """Validate the sampler configuration against the given n_obs.
+
+        This method is called at the start of each `sample()` call.
+        Override this method to add custom validation for sampler parameters.
+
+        Parameters
+        ----------
+        n_obs
+            The total number of observations in the loader.
+
+        Raises
+        ------
+        ValueError
+            If the sampler configuration is invalid for the given n_obs.
+        """
+
+    @abstractmethod
+    def _sample(self, n_obs: int) -> Iterator[LoadRequest]:
+        """Implementation of the sample method.
+
+        This method is called by the sample method to perform the actual sampling after
+        validation has passed.
+
+        Parameters
+        ----------
+        n_obs
+            The total number of observations available.
+
+        Yields
+        ------
+        LoadRequest
+            Load requests for batching data.
+        """