diffract-core 0.2.1__py3-none-any.whl

diffract/core/data/view.py

@@ -0,0 +1,389 @@
+"""Generic data view for batch operations on filtered collections.
+
+This module provides the base DataView class that can be specialized
+for different data types (parameters, relations, etc.).
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+import types
+from collections.abc import Iterator, Sequence
+from typing import TYPE_CHECKING, Self, TypeVar
+
+from tqdm.auto import tqdm
+
+from diffract.core.constants import (
+    PROGRESS_BAR_DELAY_SEC,
+    PROGRESS_BAR_MIN_ITEMS,
+)
+from diffract.core.parallel import ParallelContext, map_maybe_parallel
+from diffract.core.utils.exceptions import format_exception_message
+
+from .interface import EntityIndex, EntityUID, FieldName
+
+if TYPE_CHECKING:
+    from .interface import IDataRepository
+
+TMetadata = TypeVar("TMetadata")
+TProxy = TypeVar("TProxy")
+
+logger = logging.getLogger(__name__)
+
+_ERROR_PREVIEW_LIMIT = 10
+
+
+class DataView[TMetadata, TProxy]:
+    """Generic view over a subset of entities owned by a repository.
+
+    Provides numpy-like view semantics with support for filtering,
+    iteration, prefetching, and batch operations. Subclasses can add
+    domain-specific filtering methods.
+
+    Type Parameters:
+        TMetadata: The metadata type for entities in this view.
+        TProxy: The proxy type for entities in this view.
+    """
+
+    def __init__(
+        self,
+        *,
+        repository: IDataRepository[TMetadata, TProxy],
+        uids: list[EntityUID] | None = None,
+    ) -> None:
+        self._repository = repository
+        self._uids: list[EntityUID] | None = list(uids) if uids is not None else None
+        self._sorted: bool = False
+
+    def _ensure_uids(self) -> list[EntityUID]:
+        """Ensure UIDs are loaded (lazy load from repository if None)."""
+        if self._uids is None:
+            self._uids = self._repository.list_uids()
+            self._sorted = False
+        return self._uids
+
+    def _sort_in_place(self) -> None:
+        """Sort UIDs for deterministic ordering."""
+        uids = self._ensure_uids()
+        uids.sort()
+        self._sorted = True
+
+    def __len__(self) -> int:
+        """Return number of entities in the view."""
+        return len(self._ensure_uids())
+
+    def __iter__(self) -> Iterator[TProxy]:
+        """Iterate entity proxies in view order."""
+        if not self._sorted:
+            self._sort_in_place()
+
+        for uid in self._ensure_uids():
+            yield self._repository.get_proxy(uid)
+
+    def __getitem__(
+        self,
+        index: EntityIndex
+        | slice
+        | EntityUID
+        | Sequence[EntityIndex]
+        | Sequence[EntityUID],
+    ) -> TProxy | Self:
+        """Index into the view by position(s) or uid(s)."""
+        if not self._sorted:
+            self._sort_in_place()
+
+        uids = self._ensure_uids()
+        match index:
+            case int():
+                return self._repository.get_proxy(uids[index])
+            case slice():
+                return self.__class__(
+                    repository=self._repository,
+                    uids=uids[index],
+                )
+            case str():
+                return self._repository.get_proxy(index)
+            case _:
+                if isinstance(index, Sequence) and index:
+                    first = index[0]
+                    if isinstance(first, int):
+                        result_uids = [uids[i] for i in index]
+                    elif isinstance(first, str):
+                        result_uids = list(index)
+                    else:
+                        msg = f"Invalid index type: {type(first)}"
+                        raise ValueError(msg)
+                    return self.__class__(repository=self._repository, uids=result_uids)
+
+                return self.__class__(repository=self._repository, uids=[])
+
+    def __enter__(self) -> Self:
+        """Enter a batch write context for this view's repository storage."""
+        self._repository.storage_manager.__enter__()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: types.TracebackType | None,
+    ) -> None:
+        """Exit the batch write context for this view's repository storage."""
+        self._repository.storage_manager.__exit__(exc_type, exc_val, exc_tb)
+
+    def list_uids(self) -> list[EntityUID]:
+        """Return the list of entity UIDs contained in this view."""
+        if not self._sorted:
+            self._sort_in_place()
+        return list(self._ensure_uids())
+
+    def list_fields_by_uid(
+        self,
+        *,
+        parallel: ParallelContext | None = None,
+    ) -> dict[EntityUID, list[FieldName]]:
+        """Return mapping uid -> available fields for entities in this view."""
+        if not self._sorted:
+            self._sort_in_place()
+
+        uids = self._ensure_uids()
+        if not uids:
+            return {}
+
+        def _safe_list_fields(
+            uid: EntityUID,
+        ) -> tuple[EntityUID, list[FieldName], str | None]:
+            try:
+                return uid, self._repository.get_proxy(uid).list_fields(), None
+            except Exception as e:  # noqa: BLE001
+                return uid, [], format_exception_message(e)
+
+        dict_items = map_maybe_parallel(
+            uids,
+            _safe_list_fields,
+            parallel=parallel,
+        )
+
+        total = len(uids)
+        field_map: dict[EntityUID, list[FieldName]] = {}
+        failures: list[tuple[EntityUID, str]] = []
+        for uid, fields, error_type in tqdm(
+            dict_items,
+            desc="Listing fields...",
+            delay=PROGRESS_BAR_DELAY_SEC,
+            total=total,
+            disable=total < PROGRESS_BAR_MIN_ITEMS,
+        ):
+            field_map[uid] = fields
+            if error_type is not None:
+                failures.append((uid, error_type))
+
+        if failures:
+            preview = ", ".join(
+                f"{uid}({err})" for uid, err in failures[:_ERROR_PREVIEW_LIMIT]
+            )
+            suffix = "..." if len(failures) > _ERROR_PREVIEW_LIMIT else ""
+            logger.warning(
+                "Failed to list fields for %d entities: %s%s",
+                len(failures),
+                preview,
+                suffix,
+            )
+
+        return field_map
+
+    def prefetch_fields(
+        self,
+        *,
+        fields_by_uid: dict[EntityUID, list[FieldName]] | None = None,
+        fields: list[FieldName] | None = None,
+        verify_prefetch: bool = False,
+        parallel: ParallelContext | None = None,
+    ) -> bool:
+        """Prefetch specified fields for entities in this view into cache."""
+        if fields_by_uid is not None and fields is not None:
+            msg = "fields_by_uid and fields are mutually exclusive."
+            raise ValueError(msg)
+
+        if fields is None and fields_by_uid is None:
+            return True
+
+        if fields is not None and not fields:
+            return True
+
+        if fields_by_uid is not None and not fields_by_uid:
+            return True
+
+        if not self._sorted:
+            self._sort_in_place()
+
+        if fields is not None:
+            fields_by_uid = dict.fromkeys(self._ensure_uids(), fields)
+
+        if fields_by_uid is None:
+            return True
+
+        def _prefetch_one(
+            uid: EntityUID,
+        ) -> tuple[EntityUID, bool, str | None]:
+            uid_fields = fields_by_uid.get(uid) if fields_by_uid else []
+            if not uid_fields:
+                return uid, True, None
+            try:
+                proxy = self._repository.get_proxy(uid)
+                for field_name in uid_fields:
+                    if not proxy.try_prefetch_field(field_name):
+                        return uid, False, f"Missing field: {field_name}"
+            except Exception as e:  # noqa: BLE001
+                return uid, False, format_exception_message(e)
+            else:
+                return uid, True, None
+
+        results = map_maybe_parallel(
+            list(fields_by_uid.keys()),
+            _prefetch_one,
+            parallel=parallel,
+        )
+
+        total = len(fields_by_uid)
+        all_success = True
+        failures: list[tuple[EntityUID, str]] = []
+
+        for uid, success, error_msg in tqdm(
+            results,
+            desc="Prefetching fields...",
+            delay=PROGRESS_BAR_DELAY_SEC,
+            total=total,
+            disable=total < PROGRESS_BAR_MIN_ITEMS,
+        ):
+            if not success:
+                all_success = False
+                if error_msg:
+                    failures.append((uid, error_msg))
+
+        if failures:
+            preview = ", ".join(
+                f"{uid}({err})" for uid, err in failures[:_ERROR_PREVIEW_LIMIT]
+            )
+            suffix = "..." if len(failures) > _ERROR_PREVIEW_LIMIT else ""
+            logger.warning(
+                "Failed to prefetch fields for %d entities: %s%s",
+                len(failures),
+                preview,
+                suffix,
+            )
+
+        if verify_prefetch and all_success:
+            for uid in fields_by_uid:
+                proxy = self._repository.get_proxy(uid)
+                for field_name in fields_by_uid[uid]:
+                    if not proxy.is_field_prefetched(field_name):
+                        return False
+
+        return all_success
+
+    def erase_fields(self, *fields: FieldName) -> None:
+        """Remove specified fields from entities in this view."""
+        if not fields:
+            return
+
+        if not self._sorted:
+            self._sort_in_place()
+
+        uids = self._ensure_uids()
+        total = len(uids)
+        for uid in tqdm(
+            uids,
+            desc="Erasing fields...",
+            delay=PROGRESS_BAR_DELAY_SEC,
+            total=total,
+            disable=total < PROGRESS_BAR_MIN_ITEMS,
+        ):
+            proxy = self._repository.get_proxy(uid)
+            for field_name in fields:
+                if proxy.has_field(field_name):
+                    proxy.erase_field(field_name)
+
+    def erase_fields_with_regexp(self, *patterns: str) -> None:
+        """Remove fields matching regex patterns from entities in this view."""
+        if not patterns:
+            return
+
+        if not self._sorted:
+            self._sort_in_place()
+
+        compiled = [re.compile(p) for p in patterns]
+
+        uids = self._ensure_uids()
+        total = len(uids)
+        for uid in tqdm(
+            uids,
+            desc="Erasing fields by pattern...",
+            delay=PROGRESS_BAR_DELAY_SEC,
+            total=total,
+            disable=total < PROGRESS_BAR_MIN_ITEMS,
+        ):
+            proxy = self._repository.get_proxy(uid)
+            for field_name in proxy.list_fields():
+                if any(p.fullmatch(field_name) for p in compiled):
+                    proxy.erase_field(field_name)
+
+    def filter_by_fields(
+        self,
+        *fields: FieldName,
+        inverse_mask: bool = False,
+        parallel: ParallelContext | None = None,
+    ) -> Self:
+        """Filter entities by the presence of one or more fields."""
+        if not self._sorted:
+            self._sort_in_place()
+
+        uids = self._ensure_uids()
+        need_fields = tuple(fields)
+        if not need_fields:
+            return self.__class__(
+                repository=self._repository,
+                uids=uids if not inverse_mask else [],
+            )
+
+        def _filter(uid: EntityUID) -> EntityUID | None:
+            proxy = self._repository.get_proxy(uid)
+            for field in need_fields:
+                if not proxy.has_field(field) ^ inverse_mask:
+                    return None
+            return uid
+
+        filtered_uids = map_maybe_parallel(uids, _filter, parallel=parallel)
+
+        total = len(uids)
+        out = [
+            uid
+            for uid in tqdm(
+                filtered_uids,
+                desc="Filtering by fields...",
+                delay=PROGRESS_BAR_DELAY_SEC,
+                total=total,
+                disable=total < PROGRESS_BAR_MIN_ITEMS,
+            )
+            if uid is not None
+        ]
+
+        return self.__class__(repository=self._repository, uids=out)
+
+    def clear(self, erase: bool = False) -> None:
+        """Clear the view.
+
+        Args:
+            erase: If True, also erase data from storage and cache.
+        """
+        if erase:
+            for uid in self._ensure_uids():
+                proxy = self._repository.get_proxy(uid)
+                for field_name in proxy.list_fields():
+                    proxy.erase_field(field_name)
+        if self._uids is not None:
+            self._uids.clear()
+        else:
+            self._uids = []
+        self._sorted = True

diffract/core/export/__init__.py

@@ -0,0 +1,27 @@
+"""Export and formatting utilities for computation results.
+
+This package provides a unified interface to export results from parameter
+collections into multiple output formats using pluggable formatters.
+
+Features:
+    - Single entry point via ResultExporter
+    - Protocols for formatters and exporters
+    - DI container wiring for default formatter selection
+
+Example:
+    >>> from diffract.core.export import ResultExporter
+    >>> exporter = ResultExporter()
+    >>> out = exporter.export_results(
+    ...     "weights", parameters=collection, formatter=formatter
+    ... )
+"""
+
+from .exporters import ResultExporter
+from .interface import IResultExporter, IResultFormatter, StructuredExportResult
+
+__all__ = [
+    "IResultExporter",
+    "IResultFormatter",
+    "ResultExporter",
+    "StructuredExportResult",
+]

diffract/core/export/containers.py

@@ -0,0 +1,47 @@
+"""Dependency injection container for exporting results.
+
+Provides providers for selecting a formatter and exporting results using
+ResultExporter. The default export format is read from config.default_export_format.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+from dependency_injector import containers, providers
+
+from .exporters import ResultExporter
+from .formatters.registry import get_formatter
+
+
+def _export_with_formatter(
+    *args: Any,
+    exporter: ResultExporter,
+    export_format: str,
+    formatter_factory: Callable[[str], Any],
+    **kwargs: Any,
+) -> Any:
+    """Export results using dynamically selected formatter."""
+    return exporter.export_results(
+        *args,
+        formatter=formatter_factory(export_format),
+        **kwargs,
+    )
+
+
+class ExportContainer(containers.DeclarativeContainer):
+    """Container for export-related dependencies."""
+
+    config = providers.Configuration()
+
+    result_formatter = providers.Factory(get_formatter)
+
+    result_exporter = providers.Singleton(ResultExporter)
+
+    export_results = providers.Factory(
+        _export_with_formatter,
+        exporter=result_exporter,
+        export_format=config.default_export_format,
+        formatter_factory=result_formatter.provider,
+    )

diffract/core/export/exporters.py

@@ -0,0 +1,191 @@
+"""Exporter implementation for collecting and formatting results.
+
+The ResultExporter gathers requested fields from both parameter and aggregate
+repositories, then delegates formatting to a provided IResultFormatter implementation.
+
+Behavior:
+    - Requires at least one field name to export.
+    - Best-effort prefetch; continues if prefetch fails.
+    - Returns formatter-specific output with separate scalars and aggregates.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from diffract.core.parallel import ParallelContext
+from diffract.core.utils.exceptions import format_exception_message
+
+from .interface import AggregateData, IResultExporter, IResultFormatter, ResultData
+
+if TYPE_CHECKING:
+    from diffract.core.data.nn.aggregates.view import AggregateView
+    from diffract.core.data.nn.params.interface import IParameterView
+    from diffract.session.field_cache import SessionFieldCache
+
+logger = logging.getLogger(__name__)
+
+
+class SessionExportError(Exception):
+    """Base exception for session export errors."""
+
+
+class ResultExporter(IResultExporter):
+    """Main implementation of result export functionality."""
+
+    def export_results(
+        self,
+        *fields: tuple[str, ...],
+        parameters: IParameterView,
+        aggregates: AggregateView | None = None,
+        formatter: IResultFormatter,
+        parallel: ParallelContext | None = None,
+        field_cache: SessionFieldCache | None = None,
+    ) -> Any:
+        """Export requested fields from parameter and aggregate repositories.
+
+        Args:
+            *fields: Field names to export.
+            parameters: Parameter collection to read scalar values from.
+            aggregates: Aggregate view to read contextual/aggregation values from.
+            formatter: Formatter used to convert results to target format.
+            parallel: Optional per-method parallel context for collection operations.
+            field_cache: Optional session-level cache for field availability.
+                If provided and valid, skips expensive list_fields_by_uid() call.
+
+        Returns:
+            Formatter-specific output object (typically StructuredExportResult).
+
+        Raises:
+            ValueError: If no fields are specified.
+        """
+        if not fields:
+            msg = "At least one field must be specified"
+            raise ValueError(msg)
+
+        param_results = self._collect_parameter_results(
+            *fields, parameters=parameters, parallel=parallel, field_cache=field_cache
+        )
+        aggregate_results = self._collect_aggregate_results(
+            *fields, aggregates=aggregates, parallel=parallel
+        )
+
+        return formatter.format_results(param_results, aggregate_results, fields)
+
+    def _collect_parameter_results(
+        self,
+        *fields: tuple[str, ...],
+        parameters: IParameterView,
+        parallel: ParallelContext | None,
+        field_cache: SessionFieldCache | None = None,
+    ) -> ResultData:
+        """Collect results from parameter repository.
+
+        Args:
+            *fields: Field names to collect (exact match only, no regex patterns).
+            parameters: Parameter collection to iterate over.
+            parallel: Optional per-method parallel context for collection operations.
+            field_cache: Optional session-level cache for field availability.
+                If provided and valid, skips expensive list_fields_by_uid() call.
+
+        Returns:
+            Nested mapping with metadata and field values per parameter.
+        """
+        results: ResultData = {}
+        fields_set = set(fields)
+
+        if field_cache is not None and field_cache.is_valid:
+            cached_fields_by_uid = field_cache.get()
+            required_fields = set(parameters.list_uids())
+            to_receive = tuple(set(required_fields) - set(cached_fields_by_uid.keys()))
+            received = parameters[to_receive].list_fields_by_uid(parallel=parallel)
+            field_cache.update(received)
+            cached_fields_by_uid = field_cache.get()
+            fields_by_uid = {uid: cached_fields_by_uid[uid] for uid in required_fields}
+        else:
+            fields_by_uid = parameters.list_fields_by_uid(parallel=parallel)
+            if field_cache is not None:
+                field_cache.set(fields_by_uid)
+
+        to_receive: dict[str, list[str]] = {}
+        for uid, fields_available in fields_by_uid.items():
+            to_receive[uid] = [f for f in fields_available if f in fields_set]
+
+        try:
+            parameters.prefetch_fields(fields_by_uid=to_receive, parallel=parallel)
+        except Exception as e:  # noqa: BLE001
+            logger.debug(
+                "Prefetch failed during export: %s",
+                format_exception_message(e),
+                exc_info=True,
+            )
+
+        for param in parameters:
+            param_results: dict[str, Any] = {}
+            for field in to_receive[param.meta.uid]:
+                param_results[field] = param.get_field(field, auto_prefetch=False)
+
+            results[param.meta.uid] = {
+                "metadata": {
+                    "name": param.meta.name,
+                    "model_id": param.meta.model_id,
+                    "parameter_type": param.meta.ptype.name,
+                    **param.meta.other_meta,
+                },
+                "fields": param_results,
+            }
+
+        return results
+
+    def _collect_aggregate_results(
+        self,
+        *fields: tuple[str, ...],
+        aggregates: AggregateView | None,
+        parallel: ParallelContext | None,
+    ) -> AggregateData:
+        """Collect results from aggregate repository.
+
+        Args:
+            *fields: Base field names to collect.
+            aggregates: Aggregate view to iterate over.
+            parallel: Optional per-method parallel context (unused currently).
+
+        Returns:
+            List of aggregate records with field name, context, and value.
+        """
+        if aggregates is None or not aggregates:
+            return []
+
+        fields_set = set(fields)
+        results: AggregateData = []
+
+        # Filter aggregates by requested field names
+        filtered_view = aggregates.filter_by_field_name(*fields)
+
+        # Prefetch values
+        try:
+            filtered_view.prefetch_fields(fields=["value"], parallel=parallel)
+        except Exception as e:  # noqa: BLE001
+            logger.debug(
+                "Prefetch failed during aggregate export: %s",
+                format_exception_message(e),
+                exc_info=True,
+            )
+
+        for aggregate in filtered_view:
+            if aggregate.meta.field_name not in fields_set:
+                continue
+            if not aggregate.has_field("value"):
+                continue
+
+            results.append(
+                {
+                    "field": aggregate.meta.field_name,
+                    "context_models": aggregate.meta.context_models,
+                    "context_params": aggregate.meta.context_params,
+                    "value": aggregate.get_field("value"),
+                }
+            )
+
+        return results

diffract/core/export/formatters/__init__.py

@@ -0,0 +1,23 @@
+"""Formatters for exporting results to different output formats.
+
+Includes JSON, dictionary, and optional pandas/polars DataFrame formatters along
+with a small registry/factory for runtime selection by name.
+"""
+
+from .dict_formatter import DictFormatter
+from .json_formatter import JsonFormatter
+from .list_formatter import ListFormatter
+from .pandas_formatter import PandasFormatter  # type: ignore[F401]
+from .polars_formatter import PolarsFormatter  # type: ignore[F401]
+from .registry import FORMATTERS, get_formatter, register_formatter
+
+__all__ = [
+    "FORMATTERS",
+    "DictFormatter",
+    "JsonFormatter",
+    "ListFormatter",
+    "PandasFormatter",
+    "PolarsFormatter",
+    "get_formatter",
+    "register_formatter",
+]