diffract-core 0.2.1__py3-none-any.whl

diffract/core/storage/hybrid_manager.py

@@ -0,0 +1,218 @@
+"""Hybrid storage manager combining light and heavy backends.
+
+Routes small/structured data to light storage and large arrays to heavy storage.
+Uses a sentinel value ("__HEAVY__") in light storage to indicate the data lives
+in the heavy backend.
+
+Thread-safety: inherits the concurrency models of the underlying managers.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import types
+from typing import Any, Self
+
+import numpy as np
+
+from .interface import DEFAULT_TABLE, UID, IStorageManager
+
+
+class HybridStorageManager(IStorageManager):
+    """Storage manager that routes data between light and heavy backends.
+
+    Small values and metadata go to light storage for fast random access. Large
+    data (exceeding threshold) are stored in heavy storage with a sentinel value
+    in light storage to mark the indirection.
+
+    This provides a unified interface while allowing optimization for different
+    data access patterns and storage backends.
+
+    Example:
+        >>> # Local SQLite + Cloud Zarr
+        >>> light = SQLiteStorageManager("meta.db")
+        >>> heavy = ZarrStorageManager("s3://bucket/data")
+        >>> storage = HybridStorageManager(
+        ...     light, heavy, array_threshold=128 * 1024 * 1024
+        ... )
+        >>> with storage:
+        ...     storage.set_field("param_001", "weights", large_array)
+
+    Args:
+        light_storage: Storage backend for small/structured data and metadata.
+        heavy_storage: Storage backend for large data arrays.
+        array_threshold: Byte threshold above which data goes to heavy storage.
+    """
+
+    _HEAVY_SENTINEL = "__HEAVY__"
+
+    def __init__(
+        self,
+        light_storage: IStorageManager,
+        heavy_storage: IStorageManager,
+        array_threshold: int = 128 * 1024 * 1024,
+    ) -> None:
+        """Initialize hybrid storage manager.
+
+        Args:
+            light_storage: Storage for small data and metadata.
+            heavy_storage: Blob storage for large data arrays.
+            array_threshold: Byte threshold for routing to heavy storage.
+        """
+        self.light = light_storage
+        self.heavy = heavy_storage
+        self.array_threshold = array_threshold
+
+    def _should_use_heavy(self, value: Any) -> bool:
+        """Determine if value should be stored in heavy storage."""
+        if isinstance(value, np.ndarray):
+            return value.nbytes >= self.array_threshold
+        return False
+
+    def has_field(
+        self, obj_uid: UID, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> bool:
+        """Check if a field exists in storage."""
+        return self.light.has_field(obj_uid, field_name, table=table)
+
+    def get_field(
+        self, obj_uid: UID, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> Any:
+        """Retrieve field value, following heavy storage indirection if needed.
+
+        Raises:
+            KeyError: If the field doesn't exist.
+        """
+        value = self.light.get_field(obj_uid, field_name, table=table)
+        if isinstance(value, str) and value == self._HEAVY_SENTINEL:
+            return self.heavy.get_field(obj_uid, field_name, table=table)
+        return value
+
+    def set_field(
+        self, obj_uid: UID, field_name: str, value: Any, *, table: str = DEFAULT_TABLE
+    ) -> None:
+        """Store field value, routing data according to routing rules.
+
+        Large or heavy data goes to heavy storage with a sentinel in light storage.
+        Smaller values go directly to light storage. If a field moves between
+        backends, the old copy is erased.
+        """
+        if self._should_use_heavy(value):
+            self.light.set_field(obj_uid, field_name, self._HEAVY_SENTINEL, table=table)
+            self.heavy.set_field(obj_uid, field_name, value, table=table)
+        else:
+            self.light.set_field(obj_uid, field_name, value, table=table)
+            self.heavy.erase_field(obj_uid, field_name, table=table)
+
+    def erase_field(
+        self, obj_uid: UID, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> None:
+        """Remove a field from storage."""
+        self.light.erase_field(obj_uid, field_name, table=table)
+        self.heavy.erase_field(obj_uid, field_name, table=table)
+
+    def erase_obj(self, obj_uid: UID, *, table: str = DEFAULT_TABLE) -> None:
+        """Remove an object and all its fields from storage."""
+        self.light.erase_obj(obj_uid, table=table)
+        self.heavy.erase_obj(obj_uid, table=table)
+
+    def list_fields(
+        self, obj_uid: UID = None, *, table: str = DEFAULT_TABLE
+    ) -> list[str]:
+        """List all fields for an object."""
+        return self.light.list_fields(obj_uid, table=table)
+
+    def list_objs(self, *, table: str = DEFAULT_TABLE) -> list[str]:
+        """List all objects."""
+        return self.light.list_objs(table=table)
+
+    def list_objs_has_field(
+        self, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> list[UID]:
+        """List objects having a field."""
+        return self.light.list_objs_has_field(field_name, table=table)
+
+    def erase_field_for_all(
+        self, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> None:
+        """Remove a field from all objects.
+
+        Delegates to underlying backends which handle batch optimization.
+        """
+        self.light.erase_field_for_all(field_name, table=table)
+        self.heavy.erase_field_for_all(field_name, table=table)
+
+    def get_field_metadata(
+        self, obj_uid: UID, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> dict[str, Any] | None:
+        """Return stored metadata from whichever backend holds the value."""
+        try:
+            value = self.light.get_field(obj_uid, field_name, table=table)
+        except KeyError:
+            return None
+
+        if isinstance(value, str) and value == self._HEAVY_SENTINEL:
+            return self.heavy.get_field_metadata(obj_uid, field_name, table=table)
+        return self.light.get_field_metadata(obj_uid, field_name, table=table)
+
+    def clear(self, *, table: str | None = None) -> None:
+        """Clear data from both backends.
+
+        Args:
+            table: If provided, clear only this table. If None, clear all data.
+        """
+        self.light.clear(table=table)
+        self.heavy.clear(table=table)
+
+    def connect(self) -> None:
+        """Initialize connections for backends that need it."""
+        if hasattr(self.light, "connect"):
+            self.light.connect()
+        if hasattr(self.heavy, "connect"):
+            self.heavy.connect()
+
+    def close(self) -> None:
+        """Close both backends. Raises ExceptionGroup on errors."""
+        exceptions: list[Exception] = []
+
+        for manager in [self.light, self.heavy]:
+            try:
+                if hasattr(manager, "close"):
+                    manager.close()
+            except Exception as exc:  # noqa: BLE001
+                exceptions.append(exc)
+
+        if exceptions:
+            raise ExceptionGroup("Error occurred during closing", exceptions)
+
+    def __enter__(self) -> Self:
+        """Enter batch context for both backends."""
+        if hasattr(self.light, "__enter__"):
+            self.light.__enter__()
+        if hasattr(self.heavy, "__enter__"):
+            self.heavy.__enter__()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: types.TracebackType | None,
+    ) -> None:
+        """Exit batch context for both backends. Raises ExceptionGroup on errors."""
+        exceptions: list[Exception] = []
+
+        for manager in [self.light, self.heavy]:
+            try:
+                if hasattr(manager, "__exit__"):
+                    manager.__exit__(exc_type, exc_val, exc_tb)
+            except Exception as exc:  # noqa: BLE001
+                exceptions.append(exc)
+
+        if exceptions:
+            raise ExceptionGroup("Error occurred during exiting contexts", exceptions)
+
+    def __del__(self) -> None:
+        """Best-effort cleanup; suppresses exceptions during GC."""
+        with contextlib.suppress(Exception):
+            self.close()

diffract/core/storage/interface.py

@@ -0,0 +1,222 @@
+"""Storage manager interface and protocols.
+
+This module defines the core protocol and type aliases that all storage
+manager implementations must follow. It provides a framework-agnostic
+contract for persistent data storage operations.
+
+The interface supports hierarchical data organization where objects are
+identified by UID, organized into tables, and can contain multiple named
+fields with arbitrary values. All storage operations are designed to be
+atomic and durable.
+
+Example:
+    >>> storage_manager: IStorageManager = get_storage_manager()
+    >>> storage_manager.set_field(
+    ...     "obj123", "result", computed_array, table="parameters"
+    ... )
+    >>> value = storage_manager.get_field("obj123", "result", table="parameters")
+    >>> all_objects = storage_manager.list_objs(table="parameters")
+"""
+
+from __future__ import annotations
+
+import types
+from typing import Any, Protocol, Self, runtime_checkable
+
+type UID = str
+
+DEFAULT_TABLE = "default"
+"""Default table name for storage operations when not specified."""
+
+
+@runtime_checkable
+class IStorageManager(Protocol):
+    """Protocol defining the storage manager interface.
+
+    Provides a unified interface for persistent storage operations across
+    different backends. All storage manager implementations must implement
+    these methods to ensure consistent behavior and data durability.
+
+    The storage is organized as a three-level structure: tables contain objects
+    identified by UID, and objects contain named fields with arbitrary values.
+    Storage operations should be atomic and provide durability guarantees.
+    """
+
+    def __enter__(self) -> Self:
+        """Enter batch operation context for optimized writes."""
+        ...
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: types.TracebackType | None,
+    ) -> None:
+        """Exit batch context and flush pending operations."""
+        ...
+
+    def has_field(
+        self, obj_uid: UID, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> bool:
+        """Check if an object has a specific field in storage.
+
+        Args:
+            obj_uid: Unique identifier for the stored object.
+            field_name: Name of the field to check.
+            table: Table name for logical data separation.
+
+        Returns:
+            True if the field exists in storage, False otherwise.
+        """
+        ...
+
+    def get_field(
+        self, obj_uid: UID, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> Any:
+        """Retrieve a field value from storage.
+
+        Args:
+            obj_uid: Unique identifier for the stored object.
+            field_name: Name of the field to retrieve.
+            table: Table name for logical data separation.
+
+        Returns:
+            The stored value.
+
+        Raises:
+            KeyError: If the field doesn't exist.
+            OSError: If storage read operation fails.
+            ValueError: If stored data is corrupted.
+        """
+        ...
+
+    def list_fields(
+        self, obj_uid: UID = None, *, table: str = DEFAULT_TABLE
+    ) -> list[str]:
+        """List all field names for a specific object.
+
+        Args:
+            obj_uid: Unique identifier for the stored object.
+            table: Table name for logical data separation.
+
+        Returns:
+            List of field names that exist for the specified object.
+
+        Raises:
+            OSError: If storage read operation fails.
+        """
+        ...
+
+    def list_objs(self, *, table: str = DEFAULT_TABLE) -> list[str]:
+        """List all object UIDs in storage.
+
+        Args:
+            table: Table name for logical data separation.
+
+        Returns:
+            List of all unique object identifiers in storage.
+
+        Raises:
+            OSError: If storage read operation fails.
+        """
+        ...
+
+    def list_objs_has_field(
+        self, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> list[UID]:
+        """List all object UIDs that have a specific field.
+
+        Args:
+            field_name: Name of the field to search for.
+            table: Table name for logical data separation.
+
+        Returns:
+            List of UIDs for objects that have the specified field.
+
+        Raises:
+            OSError: If storage read operation fails.
+        """
+        ...
+
+    def set_field(
+        self, obj_uid: UID, field_name: str, value: Any, *, table: str = DEFAULT_TABLE
+    ) -> None:
+        """Store a field value in persistent storage.
+
+        Args:
+            obj_uid: Unique identifier for the stored object.
+            field_name: Name of the field to store.
+            value: Value to store (must be serializable by the backend).
+            table: Table name for logical data separation.
+
+        Raises:
+            OSError: If storage write operation fails.
+        """
+        ...
+
+    def erase_field(
+        self, obj_uid: UID, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> None:
+        """Remove a specific field from storage.
+
+        Args:
+            obj_uid: Unique identifier for the stored object.
+            field_name: Name of the field to remove.
+            table: Table name for logical data separation.
+
+        Raises:
+            OSError: If storage delete operation fails.
+        """
+        ...
+
+    def erase_obj(self, obj_uid: UID, *, table: str = DEFAULT_TABLE) -> None:
+        """Remove an entire object and all its fields from storage.
+
+        Args:
+            obj_uid: Unique identifier for the object to remove.
+            table: Table name for logical data separation.
+
+        Raises:
+            OSError: If storage delete operation fails.
+        """
+        ...
+
+    def erase_field_for_all(
+        self, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> None:
+        """Remove a field from all stored objects.
+
+        Args:
+            field_name: Name of the field to remove from all objects.
+            table: Table name for logical data separation.
+
+        Raises:
+            OSError: If storage delete operation fails.
+        """
+        ...
+
+    def clear(self, *, table: str | None = None) -> None:
+        """Remove stored data.
+
+        Args:
+            table: If provided, clear only this table. If None, clear all data.
+
+        Raises:
+            OSError: If storage clear operation fails.
+        """
+        ...
+
+    def get_field_metadata(
+        self, obj_uid: UID, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> dict[str, Any] | None:
+        """Return stored metadata (dtype/shape/kind) if available.
+
+        Args:
+            obj_uid: Unique identifier for the stored object.
+            field_name: Name of the field.
+            table: Table name for logical data separation.
+
+        Returns:
+            Metadata dictionary or None if not available.
+        """
+        ...

diffract/core/storage/metadata.py

@@ -0,0 +1,89 @@
+"""Helpers for capturing lightweight value metadata (dtype, shape, kind)."""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+from typing import Any, Literal
+
+import numpy as np
+
+Kind = Literal["scalar", "vector", "matrix", "ndarray", "object"]
+
+
+_SCALAR_NDIM = 0
+_VECTOR_NDIM = 1
+_MATRIX_NDIM = 2
+
+
+@dataclass(slots=True)
+class ValueMetadata:
+    """Normalized description of a stored value."""
+
+    kind: Kind
+    dtype: str | None
+    shape: tuple[int, ...] | None
+    ndim: int | None
+    is_numeric: bool
+
+    def to_jsonable(self) -> dict[str, Any]:
+        """Convert metadata to JSON-serializable dictionary."""
+        return asdict(self)
+
+    @classmethod
+    def from_jsonable(cls, data: dict[str, Any]) -> ValueMetadata:
+        """Reconstruct metadata from JSON-serializable dictionary."""
+        return cls(
+            kind=data.get("kind", "object"),
+            dtype=data.get("dtype"),
+            shape=tuple(data["shape"]) if data.get("shape") is not None else None,
+            ndim=data.get("ndim"),
+            is_numeric=bool(data.get("is_numeric", False)),
+        )
+
+
+def _kind_from_ndim(ndim: int) -> Kind:
+    if ndim == _SCALAR_NDIM:
+        return "scalar"
+    if ndim == _VECTOR_NDIM:
+        return "vector"
+    if ndim == _MATRIX_NDIM:
+        return "matrix"
+    return "ndarray"
+
+
+def infer_value_metadata(value: Any) -> ValueMetadata:
+    """Infer basic metadata for a value that will be stored.
+
+    Keeps the schema minimal while still being enough for visualization/signature
+    purposes.
+    """
+    if isinstance(value, np.ndarray):
+        ndim = int(value.ndim)
+        return ValueMetadata(
+            kind=_kind_from_ndim(ndim),
+            dtype=str(value.dtype),
+            shape=tuple(int(x) for x in value.shape),
+            ndim=ndim,
+            is_numeric=bool(np.issubdtype(value.dtype, np.number)),
+        )
+
+    if isinstance(value, (int, float, bool, complex, np.generic)):
+        # Scalars treated as 0-D.
+        dtype_name = type(value).__name__
+        is_numeric = isinstance(value, (int, float, complex, np.generic, bool))
+        return ValueMetadata(
+            kind="scalar",
+            dtype=dtype_name,
+            shape=None,
+            ndim=0,
+            is_numeric=is_numeric,
+        )
+
+    # Fallback: treat as opaque object.
+    return ValueMetadata(
+        kind="object",
+        dtype=type(value).__name__,
+        shape=None,
+        ndim=None,
+        is_numeric=False,
+    )

diffract/core/storage/ram_manager.py

@@ -0,0 +1,159 @@
+"""RAM-based storage manager implementation for diffract.
+
+Provides high-performance, RAM-only storage (no disk persistence).
+
+Example:
+>>> storage = RAMStorageManager()
+>>> storage.set_field("obj1", "data", [1, 2, 3], table="parameters")
+>>> storage.get_field("obj1", "data", table="parameters")
+[1, 2, 3]
+>>> storage.clear()
+"""
+
+from __future__ import annotations
+
+import types
+from typing import Any, Self
+
+from .interface import DEFAULT_TABLE, UID, IStorageManager
+from .metadata import infer_value_metadata
+
+
+class RAMStorageManager(IStorageManager):
+    """RAM-only storage manager with table support."""
+
+    def __init__(self) -> None:
+        # Maps a table and field name to a mapping from object uid to value.
+        self._storage: dict[tuple[str, str], dict[UID, Any]] = {}
+        # Maps a table, field name and object uid to that value's metadata.
+        self._metadata: dict[tuple[str, str, UID], dict[str, Any]] = {}
+
+    def _key(self, table: str, field_name: str) -> tuple[str, str]:
+        """Create storage key from table and field name."""
+        return (table, field_name)
+
+    def has_field(
+        self, obj_uid: UID, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> bool:
+        """Return True if the field exists in memory."""
+        key = self._key(table, field_name)
+        return key in self._storage and obj_uid in self._storage[key]
+
+    def get_field(
+        self, obj_uid: UID, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> Any:
+        """Get a field value from memory."""
+        key = self._key(table, field_name)
+        return self._storage[key][obj_uid]
+
+    def list_fields(
+        self, obj_uid: UID = None, *, table: str = DEFAULT_TABLE
+    ) -> list[str]:
+        """List field names (optionally only those containing the given object)."""
+        if obj_uid is None:
+            return [field_name for (tbl, field_name) in self._storage if tbl == table]
+        return [
+            field_name
+            for (tbl, field_name), objs in self._storage.items()
+            if tbl == table and obj_uid in objs
+        ]
+
+    def list_objs(self, *, table: str = DEFAULT_TABLE) -> list[str]:
+        """List all object UIDs in specified table."""
+        seen: set[str] = set()
+        for (tbl, _), objs in self._storage.items():
+            if tbl == table:
+                seen.update(objs.keys())
+        return list(seen)
+
+    def list_objs_has_field(
+        self, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> list[UID]:
+        """List objects that have the given field."""
+        key = self._key(table, field_name)
+        if key in self._storage:
+            return list(self._storage[key].keys())
+        return []
+
+    def set_field(
+        self, obj_uid: UID, field_name: str, value: Any, *, table: str = DEFAULT_TABLE
+    ) -> None:
+        """Store a field value in memory."""
+        key = self._key(table, field_name)
+        if key not in self._storage:
+            self._storage[key] = {}
+        self._storage[key][obj_uid] = value
+        self._metadata[(table, field_name, obj_uid)] = infer_value_metadata(
+            value
+        ).to_jsonable()
+
+    def erase_field(
+        self, obj_uid: UID, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> None:
+        """Remove a field for an object."""
+        if self.has_field(obj_uid, field_name, table=table):
+            key = self._key(table, field_name)
+            del self._storage[key][obj_uid]
+            self._metadata.pop((table, field_name, obj_uid), None)
+            if not self._storage[key]:
+                del self._storage[key]
+
+    def erase_obj(self, obj_uid: UID, *, table: str = DEFAULT_TABLE) -> None:
+        """Remove an object and all its fields from specified table."""
+        keys_to_remove = []
+        for (tbl, field_name), objs in self._storage.items():
+            if tbl == table and obj_uid in objs:
+                del objs[obj_uid]
+                self._metadata.pop((table, field_name, obj_uid), None)
+                if not objs:
+                    keys_to_remove.append((tbl, field_name))
+
+        for key in keys_to_remove:
+            del self._storage[key]
+
+    def erase_field_for_all(
+        self, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> None:
+        """Remove a field from all objects in specified table."""
+        key = self._key(table, field_name)
+        if key in self._storage:
+            del self._storage[key]
+        to_delete = [k for k in self._metadata if k[0] == table and k[1] == field_name]
+        for k in to_delete:
+            del self._metadata[k]
+
+    def clear(self, *, table: str | None = None) -> None:
+        """Remove data from memory.
+
+        Args:
+            table: If provided, clear only this table. If None, clear all data.
+        """
+        if table is None:
+            self._storage.clear()
+            self._metadata.clear()
+        else:
+            keys_to_remove = [k for k in self._storage if k[0] == table]
+            for k in keys_to_remove:
+                del self._storage[k]
+            meta_to_remove = [k for k in self._metadata if k[0] == table]
+            for k in meta_to_remove:
+                del self._metadata[k]
+
+    def get_field_metadata(
+        self, obj_uid: UID, field_name: str, *, table: str = DEFAULT_TABLE
+    ) -> dict[str, Any] | None:
+        """Return stored metadata for a field if available."""
+        return self._metadata.get((table, field_name, obj_uid))
+
+    def __enter__(self) -> Self:
+        """Enter context; returns self."""
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: types.TracebackType | None,
+    ) -> None:
+        """Exit context; no-op for RAM backend."""
+        return