atdata 0.2.2b1__py3-none-any.whl → 0.3.0b1__py3-none-any.whl

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", {}),
+        )

atdata/manifest/_query.py

@@ -0,0 +1,150 @@
+"""Query executor for manifest-based dataset queries.
+
+Provides two-phase filtering: shard-level pruning via aggregates,
+then sample-level filtering via the parquet DataFrame.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable
+
+import pandas as pd
+
+from ._manifest import ShardManifest
+
+
+@dataclass(frozen=True)
+class SampleLocation:
+    """Location of a sample within a shard.
+
+    Attributes:
+        shard: Shard identifier or URL.
+        key: WebDataset ``__key__`` for the sample.
+        offset: Byte offset within the tar file.
+
+    Examples:
+        >>> loc = SampleLocation(shard="data/shard-000000", key="sample_00042", offset=52480)
+        >>> loc.shard
+        'data/shard-000000'
+    """
+
+    shard: str
+    key: str
+    offset: int
+
+
+class QueryExecutor:
+    """Executes queries over per-shard manifests.
+
+    Performs two-phase filtering:
+
+    1. **Shard-level**: uses aggregates to skip shards that cannot contain
+       matching samples (e.g., numeric range exclusion, categorical value absence).
+    2. **Sample-level**: applies the predicate to the parquet DataFrame rows.
+
+    Args:
+        manifests: List of ``ShardManifest`` objects to query over.
+
+    Examples:
+        >>> executor = QueryExecutor(manifests)
+        >>> results = executor.query(
+        ...     where=lambda df: (df["confidence"] > 0.9) & (df["label"].isin(["dog", "cat"]))
+        ... )
+        >>> len(results)
+        42
+    """
+
+    def __init__(self, manifests: list[ShardManifest]) -> None:
+        self._manifests = manifests
+
+    def query(
+        self,
+        where: Callable[[pd.DataFrame], pd.Series],
+    ) -> list[SampleLocation]:
+        """Execute a query across all manifests.
+
+        The ``where`` callable receives a pandas DataFrame with the per-sample
+        manifest columns and must return a boolean Series selecting matching rows.
+
+        Args:
+            where: Predicate function. Receives a DataFrame, returns a boolean Series.
+
+        Returns:
+            List of ``SampleLocation`` for all matching samples.
+        """
+        results: list[SampleLocation] = []
+
+        for manifest in self._manifests:
+            if manifest.samples.empty:
+                continue
+
+            mask = where(manifest.samples)
+            matching = manifest.samples[mask]
+
+            for _, row in matching.iterrows():
+                results.append(
+                    SampleLocation(
+                        shard=manifest.shard_id,
+                        key=row["__key__"],
+                        offset=int(row["__offset__"]),
+                    )
+                )
+
+        return results
+
+    @classmethod
+    def from_directory(cls, directory: str | Path) -> QueryExecutor:
+        """Load all manifests from a directory.
+
+        Discovers ``*.manifest.json`` files and loads each with its
+        companion parquet file.
+
+        Args:
+            directory: Path to scan for manifest files.
+
+        Returns:
+            A ``QueryExecutor`` loaded with all discovered manifests.
+
+        Raises:
+            FileNotFoundError: If the directory does not exist.
+        """
+        directory = Path(directory)
+        manifests: list[ShardManifest] = []
+
+        for json_path in sorted(directory.glob("*.manifest.json")):
+            parquet_path = json_path.with_suffix("").with_suffix(".manifest.parquet")
+            if parquet_path.exists():
+                manifests.append(ShardManifest.from_files(json_path, parquet_path))
+            else:
+                manifests.append(ShardManifest.from_json_only(json_path))
+
+        return cls(manifests)
+
+    @classmethod
+    def from_shard_urls(cls, shard_urls: list[str]) -> QueryExecutor:
+        """Load manifests corresponding to a list of shard URLs.
+
+        Derives manifest paths by replacing the ``.tar`` extension with
+        ``.manifest.json`` and ``.manifest.parquet``.
+
+        Args:
+            shard_urls: List of shard file paths or URLs.
+
+        Returns:
+            A ``QueryExecutor`` with manifests for shards that have them.
+        """
+        manifests: list[ShardManifest] = []
+
+        for url in shard_urls:
+            base = url.removesuffix(".tar")
+            json_path = Path(f"{base}.manifest.json")
+            parquet_path = Path(f"{base}.manifest.parquet")
+
+            if json_path.exists() and parquet_path.exists():
+                manifests.append(ShardManifest.from_files(json_path, parquet_path))
+            elif json_path.exists():
+                manifests.append(ShardManifest.from_json_only(json_path))
+
+        return cls(manifests)

atdata/manifest/_writer.py

@@ -0,0 +1,74 @@
+"""ManifestWriter for serializing ShardManifest to JSON + parquet files."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from ._manifest import ShardManifest
+
+
+class ManifestWriter:
+    """Writes a ``ShardManifest`` to companion JSON and parquet files.
+
+    Produces two files alongside each shard:
+
+    - ``{base_path}.manifest.json`` -- header with metadata and aggregates
+    - ``{base_path}.manifest.parquet`` -- per-sample metadata (columnar)
+
+    Args:
+        base_path: The shard path without the ``.tar`` extension.
+
+    Examples:
+        >>> writer = ManifestWriter("/data/shard-000000")
+        >>> json_path, parquet_path = writer.write(manifest)
+    """
+
+    def __init__(self, base_path: str | Path) -> None:
+        self._base_path = Path(base_path)
+
+    @property
+    def json_path(self) -> Path:
+        """Path for the JSON header file."""
+        return self._base_path.with_suffix(".manifest.json")
+
+    @property
+    def parquet_path(self) -> Path:
+        """Path for the parquet per-sample file."""
+        return self._base_path.with_suffix(".manifest.parquet")
+
+    def write(self, manifest: ShardManifest) -> tuple[Path, Path]:
+        """Write the manifest to JSON + parquet files.
+
+        Args:
+            manifest: The ``ShardManifest`` to serialize.
+
+        Returns:
+            Tuple of ``(json_path, parquet_path)``.
+        """
+        json_out = self.json_path
+        parquet_out = self.parquet_path
+
+        # Ensure parent directory exists
+        json_out.parent.mkdir(parents=True, exist_ok=True)
+
+        # Write JSON header + aggregates
+        with open(json_out, "w", encoding="utf-8") as f:
+            json.dump(manifest.header_dict(), f, indent=2)
+
+        # Write per-sample parquet
+        if not manifest.samples.empty:
+            manifest.samples.to_parquet(
+                parquet_out,
+                engine="fastparquet",
+                index=False,
+            )
+        else:
+            # Write an empty parquet with no rows
+            manifest.samples.to_parquet(
+                parquet_out,
+                engine="fastparquet",
+                index=False,
+            )
+
+        return json_out, parquet_out