atdata 0.2.3b1__py3-none-any.whl → 0.3.1b1__py3-none-any.whl

atdata/manifest/__init__.py

@@ -0,0 +1,28 @@
+"""Per-shard manifest and query system.
+
+Provides manifest generation during shard writes and efficient
+query execution over large datasets without full scans.
+
+Components:
+
+- ``ManifestField``: Annotation marker for manifest-included fields
+- ``ManifestBuilder``: Accumulates sample metadata during writes
+- ``ShardManifest``: Loaded manifest representation
+- ``ManifestWriter``: Serializes manifests to JSON + parquet
+- ``QueryExecutor``: Two-phase query over manifest metadata
+- ``SampleLocation``: Address of a sample within a shard
+"""
+
+from ._fields import ManifestField as ManifestField
+from ._fields import AggregateKind as AggregateKind
+from ._fields import resolve_manifest_fields as resolve_manifest_fields
+from ._aggregates import CategoricalAggregate as CategoricalAggregate
+from ._aggregates import NumericAggregate as NumericAggregate
+from ._aggregates import SetAggregate as SetAggregate
+from ._aggregates import create_aggregate as create_aggregate
+from ._builder import ManifestBuilder as ManifestBuilder
+from ._manifest import ShardManifest as ShardManifest
+from ._manifest import MANIFEST_FORMAT_VERSION as MANIFEST_FORMAT_VERSION
+from ._writer import ManifestWriter as ManifestWriter
+from ._query import QueryExecutor as QueryExecutor
+from ._query import SampleLocation as SampleLocation

atdata/manifest/_aggregates.py

@@ -0,0 +1,156 @@
+"""Statistical aggregate collectors for manifest fields.
+
+Each aggregate type tracks running statistics during shard writing and
+produces a summary dict for inclusion in the manifest JSON header.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class CategoricalAggregate:
+    """Aggregate for categorical (string/enum) fields.
+
+    Tracks value counts and cardinality across all samples in a shard.
+
+    Examples:
+        >>> agg = CategoricalAggregate()
+        >>> agg.add("dog")
+        >>> agg.add("cat")
+        >>> agg.add("dog")
+        >>> agg.to_dict()
+        {'type': 'categorical', 'cardinality': 2, 'value_counts': {'dog': 2, 'cat': 1}}
+    """
+
+    value_counts: dict[str, int] = field(default_factory=dict)
+
+    @property
+    def cardinality(self) -> int:
+        """Number of distinct values observed."""
+        return len(self.value_counts)
+
+    def add(self, value: Any) -> None:
+        """Record a value observation."""
+        key = str(value)
+        self.value_counts[key] = self.value_counts.get(key, 0) + 1
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize to a JSON-compatible dict."""
+        return {
+            "type": "categorical",
+            "cardinality": self.cardinality,
+            "value_counts": dict(self.value_counts),
+        }
+
+
+@dataclass
+class NumericAggregate:
+    """Aggregate for numeric (int/float) fields.
+
+    Tracks min, max, sum, and count for computing summary statistics.
+
+    Examples:
+        >>> agg = NumericAggregate()
+        >>> agg.add(1.0)
+        >>> agg.add(3.0)
+        >>> agg.add(2.0)
+        >>> agg.to_dict()
+        {'type': 'numeric', 'min': 1.0, 'max': 3.0, 'mean': 2.0, 'count': 3}
+    """
+
+    _min: float = field(default=float("inf"))
+    _max: float = field(default=float("-inf"))
+    _sum: float = 0.0
+    count: int = 0
+
+    @property
+    def min(self) -> float:
+        """Minimum observed value."""
+        return self._min
+
+    @property
+    def max(self) -> float:
+        """Maximum observed value."""
+        return self._max
+
+    @property
+    def mean(self) -> float:
+        """Running mean of observed values."""
+        if self.count == 0:
+            return 0.0
+        return self._sum / self.count
+
+    def add(self, value: float | int) -> None:
+        """Record a numeric observation."""
+        v = float(value)
+        if v < self._min:
+            self._min = v
+        if v > self._max:
+            self._max = v
+        self._sum += v
+        self.count += 1
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize to a JSON-compatible dict."""
+        return {
+            "type": "numeric",
+            "min": self._min,
+            "max": self._max,
+            "mean": self.mean,
+            "count": self.count,
+        }
+
+
+@dataclass
+class SetAggregate:
+    """Aggregate for set/tag (list) fields.
+
+    Tracks the union of all observed values across samples.
+
+    Examples:
+        >>> agg = SetAggregate()
+        >>> agg.add(["outdoor", "day"])
+        >>> agg.add(["indoor"])
+        >>> agg.to_dict()
+        {'type': 'set', 'all_values': ['day', 'indoor', 'outdoor']}
+    """
+
+    all_values: set[str] = field(default_factory=set)
+
+    def add(self, values: list | set | tuple) -> None:
+        """Record a collection of values."""
+        for v in values:
+            self.all_values.add(str(v))
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize to a JSON-compatible dict."""
+        return {
+            "type": "set",
+            "all_values": sorted(self.all_values),
+        }
+
+
+def create_aggregate(
+    kind: str,
+) -> CategoricalAggregate | NumericAggregate | SetAggregate:
+    """Create an aggregate collector for the given kind.
+
+    Args:
+        kind: One of ``"categorical"``, ``"numeric"``, ``"set"``.
+
+    Returns:
+        A new aggregate collector instance.
+
+    Raises:
+        ValueError: If kind is not recognized.
+    """
+    if kind == "categorical":
+        return CategoricalAggregate()
+    if kind == "numeric":
+        return NumericAggregate()
+    if kind == "set":
+        return SetAggregate()
+    raise ValueError(f"Unknown aggregate kind: {kind!r}")

atdata/manifest/_builder.py

@@ -0,0 +1,163 @@
+"""ManifestBuilder for accumulating sample metadata during shard writes.
+
+Creates one ``ManifestBuilder`` per shard. Call ``add_sample()`` for each
+sample written, then ``build()`` to produce a finalized ``ShardManifest``.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Any
+
+import pandas as pd
+
+from ._aggregates import (
+    create_aggregate,
+    CategoricalAggregate,
+    NumericAggregate,
+    SetAggregate,
+)
+from ._fields import resolve_manifest_fields
+from ._manifest import ShardManifest
+
+
+@dataclass
+class _SampleRow:
+    """Internal per-sample metadata row."""
+
+    key: str
+    offset: int
+    size: int
+    fields: dict[str, Any]
+
+
+class ManifestBuilder:
+    """Accumulates sample metadata during shard writing.
+
+    Extracts manifest-annotated fields from each sample, feeds running
+    aggregate collectors, and accumulates per-sample metadata rows.
+    Call ``build()`` after all samples are written to produce a
+    ``ShardManifest``.
+
+    Args:
+        sample_type: The Packable type being written.
+        shard_id: Identifier for this shard (e.g., path without extension).
+        schema_version: Version string for the schema.
+        source_job_id: Optional provenance job identifier.
+        parent_shards: Optional list of input shard identifiers.
+        pipeline_version: Optional pipeline version string.
+
+    Examples:
+        >>> builder = ManifestBuilder(
+        ...     sample_type=ImageSample,
+        ...     shard_id="train/shard-00042",
+        ... )
+        >>> builder.add_sample(key="abc", offset=0, size=1024, sample=my_sample)
+        >>> manifest = builder.build()
+        >>> manifest.num_samples
+        1
+    """
+
+    def __init__(
+        self,
+        sample_type: type,
+        shard_id: str,
+        schema_version: str = "1.0.0",
+        source_job_id: str | None = None,
+        parent_shards: list[str] | None = None,
+        pipeline_version: str | None = None,
+    ) -> None:
+        self._sample_type = sample_type
+        self._shard_id = shard_id
+        self._schema_version = schema_version
+        self._source_job_id = source_job_id
+        self._parent_shards = parent_shards or []
+        self._pipeline_version = pipeline_version
+
+        self._manifest_fields = resolve_manifest_fields(sample_type)
+        self._aggregates: dict[
+            str, CategoricalAggregate | NumericAggregate | SetAggregate
+        ] = {
+            name: create_aggregate(mf.aggregate)
+            for name, mf in self._manifest_fields.items()
+        }
+        self._rows: list[_SampleRow] = []
+        self._total_size: int = 0
+
+    def add_sample(
+        self,
+        *,
+        key: str,
+        offset: int,
+        size: int,
+        sample: Any,
+    ) -> None:
+        """Record a sample's metadata.
+
+        Extracts manifest-annotated fields from the sample, updates
+        running aggregates, and appends a row to the internal list.
+
+        Args:
+            key: The WebDataset ``__key__`` for this sample.
+            offset: Byte offset within the tar file.
+            size: Size in bytes of this sample's tar entry.
+            sample: The sample instance (dataclass with manifest fields).
+        """
+        field_values: dict[str, Any] = {}
+        for name in self._manifest_fields:
+            value = getattr(sample, name, None)
+            if value is not None:
+                self._aggregates[name].add(value)
+                field_values[name] = value
+
+        self._rows.append(
+            _SampleRow(key=key, offset=offset, size=size, fields=field_values)
+        )
+        self._total_size += size
+
+    def build(self) -> ShardManifest:
+        """Finalize aggregates and produce a ``ShardManifest``.
+
+        Returns:
+            A complete ``ShardManifest`` with header, aggregates, and
+            per-sample DataFrame.
+        """
+        # Build aggregates dict
+        aggregates = {name: agg.to_dict() for name, agg in self._aggregates.items()}
+
+        # Build per-sample DataFrame
+        records: list[dict[str, Any]] = []
+        for row in self._rows:
+            record: dict[str, Any] = {
+                "__key__": row.key,
+                "__offset__": row.offset,
+                "__size__": row.size,
+            }
+            record.update(row.fields)
+            records.append(record)
+
+        samples_df = pd.DataFrame(records) if records else pd.DataFrame()
+
+        # Build provenance
+        provenance: dict[str, Any] = {}
+        if self._source_job_id:
+            provenance["source_job_id"] = self._source_job_id
+        if self._parent_shards:
+            provenance["parent_shards"] = self._parent_shards
+        if self._pipeline_version:
+            provenance["pipeline_version"] = self._pipeline_version
+
+        schema_name = self._sample_type.__name__
+
+        return ShardManifest(
+            shard_id=self._shard_id,
+            schema_type=schema_name,
+            schema_version=self._schema_version,
+            num_samples=len(self._rows),
+            size_bytes=self._total_size,
+            created_at=datetime.now(timezone.utc),
+            aggregates=aggregates,
+            samples=samples_df,
+            provenance=provenance,
+        )

atdata/manifest/_fields.py

@@ -0,0 +1,154 @@
+"""Manifest field annotation and introspection.
+
+Provides the ``ManifestField`` marker for annotating which sample fields
+should appear in per-shard manifests, and ``resolve_manifest_fields()``
+for introspecting a sample type to discover its manifest-included fields.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+from dataclasses import dataclass
+from typing import Any, Literal, get_args, get_origin, get_type_hints
+
+from atdata._type_utils import PRIMITIVE_TYPE_MAP, is_ndarray_type, unwrap_optional
+
+AggregateKind = Literal["categorical", "numeric", "set"]
+
+
+@dataclass(frozen=True)
+class ManifestField:
+    """Marker for manifest-included fields.
+
+    Use with ``Annotated`` to control which fields appear in per-shard
+    manifests and what aggregate statistics to compute.
+
+    Args:
+        aggregate: The type of statistical aggregate to compute.
+        exclude: If True, explicitly exclude this field from the manifest
+            even if it would be auto-inferred.
+
+    Examples:
+        >>> from typing import Annotated
+        >>> from numpy.typing import NDArray
+        >>> @atdata.packable
+        ... class ImageSample:
+        ...     image: NDArray
+        ...     label: Annotated[str, ManifestField("categorical")]
+        ...     confidence: Annotated[float, ManifestField("numeric")]
+        ...     tags: Annotated[list[str], ManifestField("set")]
+    """
+
+    aggregate: AggregateKind
+    exclude: bool = False
+
+
+def _extract_manifest_field(annotation: Any) -> ManifestField | None:
+    """Extract a ManifestField from an Annotated type, if present."""
+    if get_origin(annotation) is not None:
+        # Check for Annotated[T, ManifestField(...)]
+        # In Python 3.12+, typing.Annotated has __metadata__
+        metadata = getattr(annotation, "__metadata__", None)
+        if metadata is not None:
+            for item in metadata:
+                if isinstance(item, ManifestField):
+                    return item
+    return None
+
+
+def _infer_aggregate_kind(python_type: Any) -> AggregateKind | None:
+    """Infer the aggregate kind from a Python type annotation.
+
+    Returns:
+        The inferred aggregate kind, or None if the type should be excluded.
+    """
+    # Unwrap Optional
+    inner_type, _ = unwrap_optional(python_type)
+
+    # Exclude NDArray and bytes
+    if is_ndarray_type(inner_type):
+        return None
+    if inner_type is bytes:
+        return None
+
+    # Check primitives
+    if inner_type in PRIMITIVE_TYPE_MAP:
+        type_name = PRIMITIVE_TYPE_MAP[inner_type]
+        if type_name in ("str", "bool"):
+            return "categorical"
+        if type_name in ("int", "float"):
+            return "numeric"
+        return None
+
+    # Check list types -> set aggregate
+    origin = get_origin(inner_type)
+    if origin is list:
+        return "set"
+
+    return None
+
+
+def _get_base_type(annotation: Any) -> Any:
+    """Get the base type from an Annotated type or return as-is."""
+    args = get_args(annotation)
+    metadata = getattr(annotation, "__metadata__", None)
+    if metadata is not None and args:
+        return args[0]
+    return annotation
+
+
+def resolve_manifest_fields(sample_type: type) -> dict[str, ManifestField]:
+    """Extract manifest field descriptors from a Packable type.
+
+    Inspects type hints for ``Annotated[..., ManifestField(...)]`` markers.
+    For fields without explicit markers, applies auto-inference rules:
+
+    - ``str``, ``bool`` -> categorical
+    - ``int``, ``float`` -> numeric
+    - ``list[T]`` -> set
+    - ``NDArray``, ``bytes`` -> excluded
+
+    Args:
+        sample_type: A ``@packable`` or ``PackableSample`` subclass.
+
+    Returns:
+        Dict mapping field name to ``ManifestField`` descriptor for all
+        manifest-included fields.
+
+    Examples:
+        >>> from typing import Annotated
+        >>> @atdata.packable
+        ... class MySample:
+        ...     label: Annotated[str, ManifestField("categorical")]
+        ...     score: float
+        >>> fields = resolve_manifest_fields(MySample)
+        >>> fields["label"].aggregate
+        'categorical'
+        >>> fields["score"].aggregate
+        'numeric'
+    """
+    if not dataclasses.is_dataclass(sample_type):
+        raise TypeError(f"{sample_type} is not a dataclass")
+
+    hints = get_type_hints(sample_type, include_extras=True)
+    dc_fields = {f.name for f in dataclasses.fields(sample_type)}
+    result: dict[str, ManifestField] = {}
+
+    for name, annotation in hints.items():
+        if name not in dc_fields:
+            continue
+
+        # Check for explicit ManifestField annotation
+        explicit = _extract_manifest_field(annotation)
+        if explicit is not None:
+            if not explicit.exclude:
+                result[name] = explicit
+            continue
+
+        # Auto-infer from base type
+        base_type = _get_base_type(annotation)
+        kind = _infer_aggregate_kind(base_type)
+        if kind is not None:
+            result[name] = ManifestField(aggregate=kind)
+
+    return result

atdata/manifest/_manifest.py

@@ -0,0 +1,146 @@
+"""ShardManifest data model.
+
+Represents a loaded manifest with JSON header (metadata + aggregates)
+and per-sample metadata (as a pandas DataFrame).
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+
+MANIFEST_FORMAT_VERSION = "1.0.0"
+
+
+@dataclass
+class ShardManifest:
+    """In-memory representation of a shard's manifest.
+
+    Contains the JSON header (metadata + aggregates) and per-sample
+    metadata stored as a pandas DataFrame for efficient columnar filtering.
+
+    Attributes:
+        shard_id: Shard identifier (path without extension).
+        schema_type: Schema class name.
+        schema_version: Schema version string.
+        num_samples: Number of samples in the shard.
+        size_bytes: Total shard size in bytes.
+        created_at: When the manifest was created.
+        aggregates: Dict of field name to aggregate summary dict.
+        samples: DataFrame with ``__key__``, ``__offset__``, ``__size__``,
+            and manifest field columns.
+        provenance: Optional provenance metadata (job ID, parent shards, etc.).
+
+    Examples:
+        >>> manifest = ShardManifest.from_files(
+        ...     "data/shard-000000.manifest.json",
+        ...     "data/shard-000000.manifest.parquet",
+        ... )
+        >>> manifest.num_samples
+        1000
+        >>> manifest.aggregates["label"]["cardinality"]
+        3
+    """
+
+    shard_id: str
+    schema_type: str
+    schema_version: str
+    num_samples: int
+    size_bytes: int
+    created_at: datetime
+    aggregates: dict[str, dict[str, Any]]
+    samples: pd.DataFrame
+    provenance: dict[str, Any] = field(default_factory=dict)
+
+    def header_dict(self) -> dict[str, Any]:
+        """Return the JSON-serializable header including aggregates.
+
+        Returns:
+            Dict suitable for writing as the ``.manifest.json`` file.
+        """
+        header: dict[str, Any] = {
+            "manifest_version": MANIFEST_FORMAT_VERSION,
+            "shard_id": self.shard_id,
+            "schema_type": self.schema_type,
+            "schema_version": self.schema_version,
+            "num_samples": self.num_samples,
+            "size_bytes": self.size_bytes,
+            "created_at": self.created_at.isoformat(),
+            "aggregates": self.aggregates,
+        }
+        if self.provenance:
+            header["provenance"] = self.provenance
+        return header
+
+    @classmethod
+    def from_files(
+        cls, json_path: str | Path, parquet_path: str | Path
+    ) -> ShardManifest:
+        """Load a manifest from its JSON + parquet companion files.
+
+        Args:
+            json_path: Path to the ``.manifest.json`` file.
+            parquet_path: Path to the ``.manifest.parquet`` file.
+
+        Returns:
+            A fully loaded ``ShardManifest``.
+
+        Raises:
+            FileNotFoundError: If either file does not exist.
+            json.JSONDecodeError: If the JSON file is malformed.
+        """
+        json_path = Path(json_path)
+        parquet_path = Path(parquet_path)
+
+        with open(json_path, "r", encoding="utf-8") as f:
+            header = json.load(f)
+
+        samples = pd.read_parquet(parquet_path, engine="fastparquet")
+
+        return cls(
+            shard_id=header["shard_id"],
+            schema_type=header["schema_type"],
+            schema_version=header["schema_version"],
+            num_samples=header["num_samples"],
+            size_bytes=header["size_bytes"],
+            created_at=datetime.fromisoformat(header["created_at"]),
+            aggregates=header.get("aggregates", {}),
+            samples=samples,
+            provenance=header.get("provenance", {}),
+        )
+
+    @classmethod
+    def from_json_only(cls, json_path: str | Path) -> ShardManifest:
+        """Load header-only manifest for shard-level filtering.
+
+        Loads just the JSON header without the parquet per-sample data.
+        Useful for fast shard pruning via aggregates before loading
+        the full parquet file.
+
+        Args:
+            json_path: Path to the ``.manifest.json`` file.
+
+        Returns:
+            A ``ShardManifest`` with an empty ``samples`` DataFrame.
+        """
+        json_path = Path(json_path)
+
+        with open(json_path, "r", encoding="utf-8") as f:
+            header = json.load(f)
+
+        return cls(
+            shard_id=header["shard_id"],
+            schema_type=header["schema_type"],
+            schema_version=header["schema_version"],
+            num_samples=header["num_samples"],
+            size_bytes=header["size_bytes"],
+            created_at=datetime.fromisoformat(header["created_at"]),
+            aggregates=header.get("aggregates", {}),
+            samples=pd.DataFrame(),
+            provenance=header.get("provenance", {}),
+        )