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

atdata/dataset.py

@@ -31,6 +31,7 @@ Examples:
 import webdataset as wds
 
 from pathlib import Path
+import itertools
 import uuid
 
 import dataclasses
@@ -42,15 +43,16 @@ from dataclasses import (
 from abc import ABC
 
 from ._sources import URLSource
-from ._protocols import DataSource
+from ._protocols import DataSource, Packable
+from ._exceptions import SampleKeyError, PartialFailureError
 
-from tqdm import tqdm
 import numpy as np
 import pandas as pd
 import requests
 
 import typing
 from typing import (
+    TYPE_CHECKING,
     Any,
     Optional,
     Dict,
@@ -66,6 +68,9 @@ from typing import (
     dataclass_transform,
     overload,
 )
+
+if TYPE_CHECKING:
+    from .manifest._query import SampleLocation
 from numpy.typing import NDArray
 
 import msgpack
@@ -94,9 +99,11 @@ DT = TypeVar("DT")
 
 
 def _make_packable(x):
-    """Convert numpy arrays to bytes; pass through other values unchanged."""
+    """Convert numpy arrays to bytes; coerce numpy scalars to Python natives."""
     if isinstance(x, np.ndarray):
         return eh.array_to_bytes(x)
+    if isinstance(x, np.generic):
+        return x.item()
     return x
 
 
@@ -157,37 +164,17 @@ class DictSample:
 
     @classmethod
     def from_data(cls, data: dict[str, Any]) -> "DictSample":
-        """Create a DictSample from unpacked msgpack data.
-
-        Args:
-            data: Dictionary with field names as keys.
-
-        Returns:
-            New DictSample instance wrapping the data.
-        """
+        """Create a DictSample from unpacked msgpack data."""
         return cls(_data=data)
 
     @classmethod
     def from_bytes(cls, bs: bytes) -> "DictSample":
-        """Create a DictSample from raw msgpack bytes.
-
-        Args:
-            bs: Raw bytes from a msgpack-serialized sample.
-
-        Returns:
-            New DictSample instance with the unpacked data.
-        """
+        """Create a DictSample from raw msgpack bytes."""
         return cls.from_data(ormsgpack.unpackb(bs))
 
     def __getattr__(self, name: str) -> Any:
         """Access a field by attribute name.
 
-        Args:
-            name: Field name to access.
-
-        Returns:
-            The field value.
-
         Raises:
             AttributeError: If the field doesn't exist.
         """
@@ -203,21 +190,9 @@ class DictSample:
             ) from None
 
     def __getitem__(self, key: str) -> Any:
-        """Access a field by dict key.
-
-        Args:
-            key: Field name to access.
-
-        Returns:
-            The field value.
-
-        Raises:
-            KeyError: If the field doesn't exist.
-        """
         return self._data[key]
 
     def __contains__(self, key: str) -> bool:
-        """Check if a field exists."""
         return key in self._data
 
     def keys(self) -> list[str]:
@@ -225,23 +200,13 @@ class DictSample:
         return list(self._data.keys())
 
     def values(self) -> list[Any]:
-        """Return list of field values."""
         return list(self._data.values())
 
     def items(self) -> list[tuple[str, Any]]:
-        """Return list of (field_name, value) tuples."""
         return list(self._data.items())
 
     def get(self, key: str, default: Any = None) -> Any:
-        """Get a field value with optional default.
-
-        Args:
-            key: Field name to access.
-            default: Value to return if field doesn't exist.
-
-        Returns:
-            The field value or default.
-        """
+        """Get a field value, returning *default* if missing."""
         return self._data.get(key, default)
 
     def to_dict(self) -> dict[str, Any]:
@@ -250,20 +215,12 @@ class DictSample:
 
     @property
     def packed(self) -> bytes:
-        """Pack this sample's data into msgpack bytes.
-
-        Returns:
-            Raw msgpack bytes representing this sample's data.
-        """
+        """Serialize to msgpack bytes."""
         return msgpack.packb(self._data)
 
     @property
     def as_wds(self) -> "WDSRawSample":
-        """Pack this sample's data for writing to WebDataset.
-
-        Returns:
-            A dictionary with ``__key__`` and ``msgpack`` fields.
-        """
+        """Serialize for writing to WebDataset (``__key__`` + ``msgpack``)."""
         return {
             "__key__": str(uuid.uuid1(0, 0)),
             "msgpack": self.packed,
@@ -300,31 +257,13 @@ class PackableSample(ABC):
 
     def _ensure_good(self):
         """Convert bytes to NDArray for fields annotated as NDArray or NDArray | None."""
-
-        # Auto-convert known types when annotated
-        # for var_name, var_type in vars( self.__class__ )['__annotations__'].items():
         for field in dataclasses.fields(self):
-            var_name = field.name
-            var_type = field.type
-
-            # Annotation for this variable is to be an NDArray
-            if _is_possibly_ndarray_type(var_type):
-                # ... so, we'll always auto-convert to numpy
-
-                var_cur_value = getattr(self, var_name)
-
-                # Execute the appropriate conversion for intermediate data
-                # based on what is provided
-
-                if isinstance(var_cur_value, np.ndarray):
-                    # Already the correct type, no conversion needed
+            if _is_possibly_ndarray_type(field.type):
+                value = getattr(self, field.name)
+                if isinstance(value, np.ndarray):
                     continue
-
-                elif isinstance(var_cur_value, bytes):
-                    # Design note: bytes in NDArray-typed fields are always interpreted
-                    # as serialized arrays. This means raw bytes fields must not be
-                    # annotated as NDArray.
-                    setattr(self, var_name, eh.bytes_to_array(var_cur_value))
+                elif isinstance(value, bytes):
+                    setattr(self, field.name, eh.bytes_to_array(value))
 
     def __post_init__(self):
         self._ensure_good()
@@ -333,67 +272,31 @@ class PackableSample(ABC):
 
     @classmethod
     def from_data(cls, data: WDSRawSample) -> Self:
-        """Create a sample instance from unpacked msgpack data.
-
-        Args:
-            data: Dictionary with keys matching the sample's field names.
-
-        Returns:
-            New instance with NDArray fields auto-converted from bytes.
-        """
+        """Create an instance from unpacked msgpack data."""
         return cls(**data)
 
     @classmethod
     def from_bytes(cls, bs: bytes) -> Self:
-        """Create a sample instance from raw msgpack bytes.
-
-        Args:
-            bs: Raw bytes from a msgpack-serialized sample.
-
-        Returns:
-            A new instance of this sample class deserialized from the bytes.
-        """
+        """Create an instance from raw msgpack bytes."""
         return cls.from_data(ormsgpack.unpackb(bs))
 
     @property
     def packed(self) -> bytes:
-        """Pack this sample's data into msgpack bytes.
-
-        NDArray fields are automatically converted to bytes before packing.
-        All other fields are packed as-is if they're msgpack-compatible.
-
-        Returns:
-            Raw msgpack bytes representing this sample's data.
+        """Serialize to msgpack bytes. NDArray fields are auto-converted.
 
         Raises:
             RuntimeError: If msgpack serialization fails.
         """
-
-        # Make sure that all of our (possibly unpackable) data is in a packable
-        # format
         o = {k: _make_packable(v) for k, v in vars(self).items()}
-
         ret = msgpack.packb(o)
-
         if ret is None:
             raise RuntimeError(f"Failed to pack sample to bytes: {o}")
-
         return ret
 
     @property
     def as_wds(self) -> WDSRawSample:
-        """Pack this sample's data for writing to WebDataset.
-
-        Returns:
-            A dictionary with ``__key__`` (UUID v1 for sortable keys) and
-            ``msgpack`` (packed sample data) fields suitable for WebDataset.
-
-        Note:
-            Keys are auto-generated as UUID v1 for time-sortable ordering.
-            Custom key specification is not currently supported.
-        """
+        """Serialize for writing to WebDataset (``__key__`` + ``msgpack``)."""
         return {
-            # Generates a UUID that is timelike-sortable
             "__key__": str(uuid.uuid1(0, 0)),
             "msgpack": self.packed,
         }
@@ -404,82 +307,45 @@ def _batch_aggregate(xs: Sequence):
     if not xs:
         return []
     if isinstance(xs[0], np.ndarray):
-        return np.array(list(xs))
+        return np.stack(xs)
     return list(xs)
 
 
 class SampleBatch(Generic[DT]):
     """A batch of samples with automatic attribute aggregation.
 
-    This class wraps a sequence of samples and provides magic ``__getattr__``
-    access to aggregate sample attributes. When you access an attribute that
-    exists on the sample type, it automatically aggregates values across all
-    samples in the batch.
-
-    NDArray fields are stacked into a numpy array with a batch dimension.
-    Other fields are aggregated into a list.
+    Accessing an attribute aggregates that field across all samples:
+    NDArray fields are stacked into a numpy array with a batch dimension;
+    other fields are collected into a list. Results are cached.
 
     Parameters:
         DT: The sample type, must derive from ``PackableSample``.
 
-    Attributes:
-        samples: The list of sample instances in this batch.
-
     Examples:
         >>> batch = SampleBatch[MyData]([sample1, sample2, sample3])
-        >>> batch.embeddings  # Returns stacked numpy array of shape (3, ...)
-        >>> batch.names  # Returns list of names
-
-    Note:
-        This class uses Python's ``__orig_class__`` mechanism to extract the
-        type parameter at runtime. Instances must be created using the
-        subscripted syntax ``SampleBatch[MyType](samples)`` rather than
-        calling the constructor directly with an unsubscripted class.
+        >>> batch.embeddings  # Stacked numpy array of shape (3, ...)
+        >>> batch.names  # List of names
     """
 
-    # Design note: The docstring uses "Parameters:" for type parameters because
-    # quartodoc doesn't yet support "Type Parameters:" sections in generated docs.
-
     def __init__(self, samples: Sequence[DT]):
-        """Create a batch from a sequence of samples.
-
-        Args:
-            samples: A sequence of sample instances to aggregate into a batch.
-                Each sample must be an instance of a type derived from
-                ``PackableSample``.
-        """
+        """Create a batch from a sequence of samples."""
         self.samples = list(samples)
         self._aggregate_cache = dict()
         self._sample_type_cache: Type | None = None
 
     @property
     def sample_type(self) -> Type:
-        """The type of each sample in this batch.
-
-        Returns:
-            The type parameter ``DT`` used when creating this ``SampleBatch[DT]``.
-        """
+        """The type parameter ``DT`` used when creating this batch."""
         if self._sample_type_cache is None:
             self._sample_type_cache = typing.get_args(self.__orig_class__)[0]
-            assert self._sample_type_cache is not None
+            if self._sample_type_cache is None:
+                raise TypeError(
+                    "SampleBatch requires a type parameter, e.g. SampleBatch[MySample]"
+                )
         return self._sample_type_cache
 
     def __getattr__(self, name):
-        """Aggregate an attribute across all samples in the batch.
-
-        This magic method enables attribute-style access to aggregated sample
-        fields. Results are cached for efficiency.
-
-        Args:
-            name: The attribute name to aggregate across samples.
-
-        Returns:
-            For NDArray fields: a stacked numpy array with batch dimension.
-            For other fields: a list of values from each sample.
-
-        Raises:
-            AttributeError: If the attribute doesn't exist on the sample type.
-        """
+        """Aggregate a field across all samples (cached)."""
         # Aggregate named params of sample type
         if name in vars(self.sample_type)["__annotations__"]:
             if name not in self._aggregate_cache:
@@ -492,8 +358,8 @@ class SampleBatch(Generic[DT]):
         raise AttributeError(f"No sample attribute named {name}")
 
 
-ST = TypeVar("ST", bound=PackableSample)
-RT = TypeVar("RT", bound=PackableSample)
+ST = TypeVar("ST", bound=Packable)
+RT = TypeVar("RT", bound=Packable)
 
 
 class _ShardListStage(wds.utils.PipelineStage):
@@ -571,23 +437,18 @@ class Dataset(Generic[ST]):
 
     @property
     def sample_type(self) -> Type:
-        """The type of each returned sample from this dataset's iterator.
-
-        Returns:
-            The type parameter ``ST`` used when creating this ``Dataset[ST]``.
-        """
+        """The type parameter ``ST`` used when creating this dataset."""
         if self._sample_type_cache is None:
             self._sample_type_cache = typing.get_args(self.__orig_class__)[0]
-            assert self._sample_type_cache is not None
+            if self._sample_type_cache is None:
+                raise TypeError(
+                    "Dataset requires a type parameter, e.g. Dataset[MySample]"
+                )
         return self._sample_type_cache
 
     @property
     def batch_type(self) -> Type:
-        """The type of batches produced by this dataset.
-
-        Returns:
-            ``SampleBatch[ST]`` where ``ST`` is this dataset's sample type.
-        """
+        """``SampleBatch[ST]`` where ``ST`` is this dataset's sample type."""
         return SampleBatch[self.sample_type]
 
     def __init__(
@@ -614,28 +475,21 @@ class Dataset(Generic[ST]):
         """
         super().__init__()
 
-        # Handle backward compatibility: url= keyword argument
         if source is None and url is not None:
             source = url
         elif source is None:
             raise TypeError("Dataset() missing required argument: 'source' or 'url'")
 
-        # Normalize source: strings become URLSource for backward compatibility
         if isinstance(source, str):
             self._source: DataSource = URLSource(source)
             self.url = source
         else:
             self._source = source
-            # For compatibility, expose URL if source has list_shards
             shards = source.list_shards()
-            # Design note: Using first shard as url for legacy compatibility.
-            # Full shard list is available via list_shards() method.
             self.url = shards[0] if shards else ""
 
         self._metadata: dict[str, Any] | None = None
         self.metadata_url: str | None = metadata_url
-        """Optional URL to msgpack-encoded metadata for this dataset."""
-
         self._output_lens: Lens | None = None
         self._sample_type_cache: Type | None = None
 
@@ -645,47 +499,23 @@ class Dataset(Generic[ST]):
         return self._source
 
     def as_type(self, other: Type[RT]) -> "Dataset[RT]":
-        """View this dataset through a different sample type using a registered lens.
-
-        Args:
-            other: The target sample type to transform into. Must be a type
-                derived from ``PackableSample``.
-
-        Returns:
-            A new ``Dataset`` instance that yields samples of type ``other``
-            by applying the appropriate lens transformation from the global
-            ``LensNetwork`` registry.
+        """View this dataset through a different sample type via a registered lens.
 
         Raises:
-            ValueError: If no registered lens exists between the current
-                sample type and the target type.
+            ValueError: If no lens exists between the current and target types.
         """
         ret = Dataset[other](self._source)
-        # Get the singleton lens registry
         lenses = LensNetwork()
         ret._output_lens = lenses.transform(self.sample_type, ret.sample_type)
         return ret
 
     @property
     def shards(self) -> Iterator[str]:
-        """Lazily iterate over shard identifiers.
-
-        Yields:
-            Shard identifiers (e.g., 'train-000000.tar', 'train-000001.tar').
-
-        Examples:
-            >>> for shard in ds.shards:
-            ...     print(f"Processing {shard}")
-        """
+        """Lazily iterate over shard identifiers."""
         return iter(self._source.list_shards())
 
     def list_shards(self) -> list[str]:
-        """Get list of individual dataset shards.
-
-        Returns:
-            A full (non-lazy) list of the individual ``tar`` files within the
-            source WebDataset.
-        """
+        """Return all shard paths/URLs as a list."""
         return self._source.list_shards()
 
     # Legacy alias for backwards compatibility
@@ -707,14 +537,7 @@ class Dataset(Generic[ST]):
 
     @property
     def metadata(self) -> dict[str, Any] | None:
-        """Fetch and cache metadata from metadata_url.
-
-        Returns:
-            Deserialized metadata dictionary, or None if no metadata_url is set.
-
-        Raises:
-            requests.HTTPError: If metadata fetch fails.
-        """
+        """Fetch and cache metadata from metadata_url, or ``None`` if unset."""
         if self.metadata_url is None:
             return None
 
@@ -726,6 +549,367 @@ class Dataset(Generic[ST]):
         # Use our cached values
         return self._metadata
 
+    ##
+    # Convenience methods (GH#38 developer experience)
+
+    @property
+    def schema(self) -> dict[str, type]:
+        """Field names and types for this dataset's sample type.
+
+        Examples:
+            >>> ds = Dataset[MyData]("data.tar")
+            >>> ds.schema
+            {'name': <class 'str'>, 'embedding': numpy.ndarray}
+        """
+        st = self.sample_type
+        if st is DictSample:
+            return {"_data": dict}
+        if dataclasses.is_dataclass(st):
+            return {f.name: f.type for f in dataclasses.fields(st)}
+        return {}
+
+    @property
+    def column_names(self) -> list[str]:
+        """List of field names for this dataset's sample type."""
+        st = self.sample_type
+        if dataclasses.is_dataclass(st):
+            return [f.name for f in dataclasses.fields(st)]
+        return []
+
+    def __iter__(self) -> Iterator[ST]:
+        """Shorthand for ``ds.ordered()``."""
+        return iter(self.ordered())
+
+    def __len__(self) -> int:
+        """Total sample count (iterates all shards on first call, then cached)."""
+        if not hasattr(self, "_len_cache"):
+            self._len_cache: int = sum(1 for _ in self.ordered())
+        return self._len_cache
+
+    def head(self, n: int = 5) -> list[ST]:
+        """Return the first *n* samples from the dataset.
+
+        Args:
+            n: Number of samples to return. Default: 5.
+
+        Returns:
+            List of up to *n* samples in shard order.
+
+        Examples:
+            >>> samples = ds.head(3)
+            >>> len(samples)
+            3
+        """
+        return list(itertools.islice(self.ordered(), n))
+
+    def get(self, key: str) -> ST:
+        """Retrieve a single sample by its ``__key__``.
+
+        Scans shards sequentially until a sample with a matching key is found.
+        This is O(n) for streaming datasets.
+
+        Args:
+            key: The WebDataset ``__key__`` string to search for.
+
+        Returns:
+            The matching sample.
+
+        Raises:
+            SampleKeyError: If no sample with the given key exists.
+
+        Examples:
+            >>> sample = ds.get("00000001-0001-1000-8000-010000000000")
+        """
+        pipeline = wds.pipeline.DataPipeline(
+            _ShardListStage(self._source),
+            wds.shardlists.split_by_worker,
+            _StreamOpenerStage(self._source),
+            wds.tariterators.tar_file_expander,
+            wds.tariterators.group_by_keys,
+        )
+        for raw_sample in pipeline:
+            if raw_sample.get("__key__") == key:
+                return self.wrap(raw_sample)
+        raise SampleKeyError(key)
+
+    def describe(self) -> dict[str, Any]:
+        """Summary statistics: sample_type, fields, num_shards, shards, url, metadata."""
+        shards = self.list_shards()
+        return {
+            "sample_type": self.sample_type.__name__,
+            "fields": self.schema,
+            "num_shards": len(shards),
+            "shards": shards,
+            "url": self.url,
+            "metadata": self.metadata,
+        }
+
+    def filter(self, predicate: Callable[[ST], bool]) -> "Dataset[ST]":
+        """Return a new dataset that yields only samples matching *predicate*.
+
+        The filter is applied lazily during iteration — no data is copied.
+
+        Args:
+            predicate: A function that takes a sample and returns ``True``
+                to keep it or ``False`` to discard it.
+
+        Returns:
+            A new ``Dataset`` whose iterators apply the filter.
+
+        Examples:
+            >>> long_names = ds.filter(lambda s: len(s.name) > 10)
+            >>> for sample in long_names:
+            ...     assert len(sample.name) > 10
+        """
+        filtered = Dataset[self.sample_type](self._source, self.metadata_url)
+        filtered._sample_type_cache = self._sample_type_cache
+        filtered._output_lens = self._output_lens
+        filtered._filter_fn = predicate
+        # Preserve any existing filters
+        parent_filters = getattr(self, "_filter_fn", None)
+        if parent_filters is not None:
+            outer = parent_filters
+            filtered._filter_fn = lambda s: outer(s) and predicate(s)
+        # Preserve any existing map
+        if hasattr(self, "_map_fn"):
+            filtered._map_fn = self._map_fn
+        return filtered
+
+    def map(self, fn: Callable[[ST], Any]) -> "Dataset":
+        """Return a new dataset that applies *fn* to each sample during iteration.
+
+        The mapping is applied lazily during iteration — no data is copied.
+
+        Args:
+            fn: A function that takes a sample of type ``ST`` and returns
+                a transformed value.
+
+        Returns:
+            A new ``Dataset`` whose iterators apply the mapping.
+
+        Examples:
+            >>> names = ds.map(lambda s: s.name)
+            >>> for name in names:
+            ...     print(name)
+        """
+        mapped = Dataset[self.sample_type](self._source, self.metadata_url)
+        mapped._sample_type_cache = self._sample_type_cache
+        mapped._output_lens = self._output_lens
+        mapped._map_fn = fn
+        # Preserve any existing map
+        if hasattr(self, "_map_fn"):
+            outer = self._map_fn
+            mapped._map_fn = lambda s: fn(outer(s))
+        # Preserve any existing filter
+        if hasattr(self, "_filter_fn"):
+            mapped._filter_fn = self._filter_fn
+        return mapped
+
+    def process_shards(
+        self,
+        fn: Callable[[list[ST]], Any],
+        *,
+        shards: list[str] | None = None,
+    ) -> dict[str, Any]:
+        """Process each shard independently, collecting per-shard results.
+
+        Unlike :meth:`map` (which is lazy and per-sample), this method eagerly
+        processes each shard in turn, calling *fn* with the full list of samples
+        from that shard. If some shards fail, raises
+        :class:`~atdata._exceptions.PartialFailureError` containing both the
+        successful results and the per-shard errors.
+
+        Args:
+            fn: Function receiving a list of samples from one shard and
+                returning an arbitrary result.
+            shards: Optional list of shard identifiers to process. If ``None``,
+                processes all shards in the dataset. Useful for retrying only
+                the failed shards from a previous ``PartialFailureError``.
+
+        Returns:
+            Dict mapping shard identifier to *fn*'s return value for each shard.
+
+        Raises:
+            PartialFailureError: If at least one shard fails. The exception
+                carries ``.succeeded_shards``, ``.failed_shards``, ``.errors``,
+                and ``.results`` for inspection and retry.
+
+        Examples:
+            >>> results = ds.process_shards(lambda samples: len(samples))
+            >>> # On partial failure, retry just the failed shards:
+            >>> try:
+            ...     results = ds.process_shards(expensive_fn)
+            ... except PartialFailureError as e:
+            ...     retry = ds.process_shards(expensive_fn, shards=e.failed_shards)
+        """
+        from ._logging import get_logger
+
+        log = get_logger()
+        shard_ids = shards or self.list_shards()
+        log.info("process_shards: starting %d shards", len(shard_ids))
+
+        succeeded: list[str] = []
+        failed: list[str] = []
+        errors: dict[str, Exception] = {}
+        results: dict[str, Any] = {}
+
+        for shard_id in shard_ids:
+            try:
+                shard_ds = Dataset[self.sample_type](shard_id)
+                shard_ds._sample_type_cache = self._sample_type_cache
+                samples = list(shard_ds.ordered())
+                results[shard_id] = fn(samples)
+                succeeded.append(shard_id)
+                log.debug("process_shards: shard ok %s", shard_id)
+            except Exception as exc:
+                failed.append(shard_id)
+                errors[shard_id] = exc
+                log.warning("process_shards: shard failed %s: %s", shard_id, exc)
+
+        if failed:
+            log.error(
+                "process_shards: %d/%d shards failed",
+                len(failed),
+                len(shard_ids),
+            )
+            raise PartialFailureError(
+                succeeded_shards=succeeded,
+                failed_shards=failed,
+                errors=errors,
+                results=results,
+            )
+
+        log.info("process_shards: all %d shards succeeded", len(shard_ids))
+        return results
+
+    def select(self, indices: Sequence[int]) -> list[ST]:
+        """Return samples at the given integer indices.
+
+        Iterates through the dataset in order and collects samples whose
+        positional index matches. This is O(n) for streaming datasets.
+
+        Args:
+            indices: Sequence of zero-based indices to select.
+
+        Returns:
+            List of samples at the requested positions, in index order.
+
+        Examples:
+            >>> samples = ds.select([0, 5, 10])
+            >>> len(samples)
+            3
+        """
+        if not indices:
+            return []
+        target = set(indices)
+        max_idx = max(indices)
+        result: dict[int, ST] = {}
+        for i, sample in enumerate(self.ordered()):
+            if i in target:
+                result[i] = sample
+            if i >= max_idx:
+                break
+        return [result[i] for i in indices if i in result]
+
+    def query(
+        self,
+        where: "Callable[[pd.DataFrame], pd.Series]",
+    ) -> "list[SampleLocation]":
+        """Query this dataset using per-shard manifest metadata.
+
+        Requires manifests to have been generated during shard writing.
+        Discovers manifest files alongside the tar shards, loads them,
+        and executes a two-phase query (shard-level aggregate pruning,
+        then sample-level parquet filtering).
+
+        Args:
+            where: Predicate function that receives a pandas DataFrame
+                of manifest fields and returns a boolean Series selecting
+                matching rows.
+
+        Returns:
+            List of ``SampleLocation`` for matching samples.
+
+        Raises:
+            FileNotFoundError: If no manifest files are found alongside shards.
+
+        Examples:
+            >>> locs = ds.query(where=lambda df: df["confidence"] > 0.9)
+            >>> len(locs)
+            42
+        """
+        from .manifest import QueryExecutor
+
+        shard_urls = self.list_shards()
+        executor = QueryExecutor.from_shard_urls(shard_urls)
+        return executor.query(where=where)
+
+    def to_pandas(self, limit: int | None = None) -> "pd.DataFrame":
+        """Materialize the dataset (or first *limit* samples) as a DataFrame.
+
+        Args:
+            limit: Maximum number of samples to include. ``None`` means all
+                samples (may use significant memory for large datasets).
+
+        Returns:
+            A pandas DataFrame with one row per sample and columns matching
+            the sample fields.
+
+        Warning:
+            With ``limit=None`` this loads the entire dataset into memory.
+
+        Examples:
+            >>> df = ds.to_pandas(limit=100)
+            >>> df.columns.tolist()
+            ['name', 'embedding']
+        """
+        samples = self.head(limit) if limit is not None else list(self.ordered())
+        rows = [
+            asdict(s) if dataclasses.is_dataclass(s) else s.to_dict() for s in samples
+        ]
+        return pd.DataFrame(rows)
+
+    def to_dict(self, limit: int | None = None) -> dict[str, list[Any]]:
+        """Materialize the dataset as a column-oriented dictionary.
+
+        Args:
+            limit: Maximum number of samples to include. ``None`` means all.
+
+        Returns:
+            Dictionary mapping field names to lists of values (one entry
+            per sample).
+
+        Warning:
+            With ``limit=None`` this loads the entire dataset into memory.
+
+        Examples:
+            >>> d = ds.to_dict(limit=10)
+            >>> d.keys()
+            dict_keys(['name', 'embedding'])
+            >>> len(d['name'])
+            10
+        """
+        samples = self.head(limit) if limit is not None else list(self.ordered())
+        if not samples:
+            return {}
+        if dataclasses.is_dataclass(samples[0]):
+            fields = [f.name for f in dataclasses.fields(samples[0])]
+            return {f: [getattr(s, f) for s in samples] for f in fields}
+        # DictSample path
+        keys = samples[0].keys()
+        return {k: [s[k] for s in samples] for k in keys}
+
+    def _post_wrap_stages(self) -> list:
+        """Build extra pipeline stages for filter/map set via .filter()/.map()."""
+        stages: list = []
+        filter_fn = getattr(self, "_filter_fn", None)
+        if filter_fn is not None:
+            stages.append(wds.filters.select(filter_fn))
+        map_fn = getattr(self, "_map_fn", None)
+        if map_fn is not None:
+            stages.append(wds.filters.map(map_fn))
+        return stages
+
     @overload
     def ordered(
         self,
@@ -769,6 +953,7 @@ class Dataset(Generic[ST]):
                 wds.tariterators.tar_file_expander,
                 wds.tariterators.group_by_keys,
                 wds.filters.map(self.wrap),
+                *self._post_wrap_stages(),
             )
 
         return wds.pipeline.DataPipeline(
@@ -839,6 +1024,7 @@ class Dataset(Generic[ST]):
                 wds.tariterators.group_by_keys,
                 wds.filters.shuffle(buffer_samples),
                 wds.filters.map(self.wrap),
+                *self._post_wrap_stages(),
             )
 
         return wds.pipeline.DataPipeline(
@@ -862,100 +1048,47 @@ class Dataset(Generic[ST]):
         maxcount: Optional[int] = None,
         **kwargs,
     ):
-        """Export dataset contents to parquet format.
-
-        Converts all samples to a pandas DataFrame and saves to parquet file(s).
-        Useful for interoperability with data analysis tools.
+        """Export dataset to parquet file(s).
 
         Args:
-            path: Output path for the parquet file. If ``maxcount`` is specified,
-                files are named ``{stem}-{segment:06d}.parquet``.
-            sample_map: Optional function to convert samples to dictionaries.
-                Defaults to ``dataclasses.asdict``.
-            maxcount: If specified, split output into multiple files with at most
-                this many samples each. Recommended for large datasets.
-            **kwargs: Additional arguments passed to ``pandas.DataFrame.to_parquet()``.
-                Common options include ``compression``, ``index``, ``engine``.
-
-        Warning:
-            **Memory Usage**: When ``maxcount=None`` (default), this method loads
-            the **entire dataset into memory** as a pandas DataFrame before writing.
-            For large datasets, this can cause memory exhaustion.
-
-            For datasets larger than available RAM, always specify ``maxcount``::
-
-                # Safe for large datasets - processes in chunks
-                ds.to_parquet("output.parquet", maxcount=10000)
-
-            This creates multiple parquet files: ``output-000000.parquet``,
-            ``output-000001.parquet``, etc.
+            path: Output path. With *maxcount*, files are named
+                ``{stem}-{segment:06d}.parquet``.
+            sample_map: Convert sample to dict. Defaults to ``dataclasses.asdict``.
+            maxcount: Split into files of at most this many samples.
+                Without it, the entire dataset is loaded into memory.
+            **kwargs: Passed to ``pandas.DataFrame.to_parquet()``.
 
         Examples:
-            >>> ds = Dataset[MySample]("data.tar")
-            >>> # Small dataset - load all at once
-            >>> ds.to_parquet("output.parquet")
-            >>>
-            >>> # Large dataset - process in chunks
             >>> ds.to_parquet("output.parquet", maxcount=50000)
         """
-        ##
-
-        # Normalize args
         path = Path(path)
         if sample_map is None:
             sample_map = asdict
 
-        verbose = kwargs.get("verbose", False)
-
-        it = self.ordered(batch_size=None)
-        if verbose:
-            it = tqdm(it)
-
-        #
-
         if maxcount is None:
-            # Load and save full dataset
             df = pd.DataFrame([sample_map(x) for x in self.ordered(batch_size=None)])
             df.to_parquet(path, **kwargs)
-
         else:
-            # Load and save dataset in segments of size `maxcount`
-
             cur_segment = 0
-            cur_buffer = []
+            cur_buffer: list = []
             path_template = (
                 path.parent / f"{path.stem}-{{:06d}}{path.suffix}"
             ).as_posix()
 
             for x in self.ordered(batch_size=None):
                 cur_buffer.append(sample_map(x))
-
                 if len(cur_buffer) >= maxcount:
-                    # Write current segment
                     cur_path = path_template.format(cur_segment)
-                    df = pd.DataFrame(cur_buffer)
-                    df.to_parquet(cur_path, **kwargs)
-
+                    pd.DataFrame(cur_buffer).to_parquet(cur_path, **kwargs)
                     cur_segment += 1
                     cur_buffer = []
 
-            if len(cur_buffer) > 0:
-                # Write one last segment with remainder
+            if cur_buffer:
                 cur_path = path_template.format(cur_segment)
-                df = pd.DataFrame(cur_buffer)
-                df.to_parquet(cur_path, **kwargs)
+                pd.DataFrame(cur_buffer).to_parquet(cur_path, **kwargs)
 
     def wrap(self, sample: WDSRawSample) -> ST:
-        """Wrap a raw msgpack sample into the appropriate dataset-specific type.
-
-        Args:
-            sample: A dictionary containing at minimum a ``'msgpack'`` key with
-                serialized sample bytes.
-
-        Returns:
-            A deserialized sample of type ``ST``, optionally transformed through
-            a lens if ``as_type()`` was called.
-        """
+        """Deserialize a raw WDS sample dict into type ``ST``."""
         if "msgpack" not in sample:
             raise ValueError(
                 f"Sample missing 'msgpack' key, got keys: {list(sample.keys())}"
@@ -972,20 +1105,7 @@ class Dataset(Generic[ST]):
         return self._output_lens(source_sample)
 
     def wrap_batch(self, batch: WDSRawBatch) -> SampleBatch[ST]:
-        """Wrap a batch of raw msgpack samples into a typed SampleBatch.
-
-        Args:
-            batch: A dictionary containing a ``'msgpack'`` key with a list of
-                serialized sample bytes.
-
-        Returns:
-            A ``SampleBatch[ST]`` containing deserialized samples, optionally
-            transformed through a lens if ``as_type()`` was called.
-
-        Note:
-            This implementation deserializes samples one at a time, then
-            aggregates them into a batch.
-        """
+        """Deserialize a raw WDS batch dict into ``SampleBatch[ST]``."""
 
         if "msgpack" not in batch:
             raise ValueError(
@@ -1009,24 +1129,12 @@ _T = TypeVar("_T")
 
 
 @dataclass_transform()
-def packable(cls: type[_T]) -> type[_T]:
-    """Decorator to convert a regular class into a ``PackableSample``.
-
-    This decorator transforms a class into a dataclass that inherits from
-    ``PackableSample``, enabling automatic msgpack serialization/deserialization
-    with special handling for NDArray fields.
+def packable(cls: type[_T]) -> type[Packable]:
+    """Convert a class into a ``PackableSample`` dataclass with msgpack serialization.
 
-    The resulting class satisfies the ``Packable`` protocol, making it compatible
-    with all atdata APIs that accept packable types (e.g., ``publish_schema``,
-    lens transformations, etc.).
-
-    Args:
-        cls: The class to convert. Should have type annotations for its fields.
-
-    Returns:
-        A new dataclass that inherits from ``PackableSample`` with the same
-        name and annotations as the original class. The class satisfies the
-        ``Packable`` protocol and can be used with ``Type[Packable]`` signatures.
+    The resulting class gains ``packed``, ``as_wds``, ``from_bytes``, and
+    ``from_data`` methods, and satisfies the ``Packable`` protocol.
+    NDArray fields are automatically handled during serialization.
 
     Examples:
         >>> @packable
@@ -1035,11 +1143,7 @@ def packable(cls: type[_T]) -> type[_T]:
         ...     values: NDArray
         ...
         >>> sample = MyData(name="test", values=np.array([1, 2, 3]))
-        >>> bytes_data = sample.packed
-        >>> restored = MyData.from_bytes(bytes_data)
-        >>>
-        >>> # Works with Packable-typed APIs
-        >>> index.publish_schema(MyData, version="1.0.0")  # Type-safe
+        >>> restored = MyData.from_bytes(sample.packed)
     """
 
     ##
@@ -1086,3 +1190,154 @@ def packable(cls: type[_T]) -> type[_T]:
     ##
 
     return as_packable
+
+
+# ---------------------------------------------------------------------------
+# write_samples — convenience function for writing samples to tar files
+# ---------------------------------------------------------------------------
+
+
+def write_samples(
+    samples: Iterable[ST],
+    path: str | Path,
+    *,
+    maxcount: int | None = None,
+    maxsize: int | None = None,
+    manifest: bool = False,
+) -> "Dataset[ST]":
+    """Write an iterable of samples to WebDataset tar file(s).
+
+    Args:
+        samples: Iterable of ``PackableSample`` instances. Must be non-empty.
+        path: Output path for the tar file. For sharded output (when
+            *maxcount* or *maxsize* is set), a ``%06d`` pattern is
+            auto-appended if the path does not already contain ``%``.
+        maxcount: Maximum samples per shard. Triggers multi-shard output.
+        maxsize: Maximum bytes per shard. Triggers multi-shard output.
+        manifest: If True, write per-shard manifest sidecar files
+            (``.manifest.json`` + ``.manifest.parquet``) alongside each
+            tar file. Manifests enable metadata queries via
+            ``QueryExecutor`` without opening the tars.
+
+    Returns:
+        A ``Dataset`` wrapping the written file(s), typed to the sample
+        type of the input samples.
+
+    Raises:
+        ValueError: If *samples* is empty.
+
+    Examples:
+        >>> samples = [MySample(key="0", text="hello")]
+        >>> ds = write_samples(samples, "out.tar")
+        >>> list(ds.ordered())
+        [MySample(key='0', text='hello')]
+    """
+    from ._hf_api import _shards_to_wds_url
+
+    if manifest:
+        from .manifest._builder import ManifestBuilder
+        from .manifest._writer import ManifestWriter
+
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    use_shard_writer = maxcount is not None or maxsize is not None
+    sample_type: type | None = None
+    written_paths: list[str] = []
+
+    # Manifest tracking state
+    _current_builder: list = []  # single-element list for nonlocal mutation
+    _builders: list[tuple[str, "ManifestBuilder"]] = []
+    _running_offset: list[int] = [0]
+
+    def _finalize_builder() -> None:
+        """Finalize the current manifest builder and stash it."""
+        if _current_builder:
+            shard_path = written_paths[-1] if written_paths else ""
+            _builders.append((shard_path, _current_builder[0]))
+            _current_builder.clear()
+
+    def _start_builder(shard_path: str) -> None:
+        """Start a new manifest builder for a shard."""
+        _finalize_builder()
+        shard_id = Path(shard_path).stem
+        _current_builder.append(
+            ManifestBuilder(sample_type=sample_type, shard_id=shard_id)
+        )
+        _running_offset[0] = 0
+
+    def _record_sample(sample: "PackableSample", wds_dict: dict) -> None:
+        """Record a sample in the active manifest builder."""
+        if not _current_builder:
+            return
+        packed_bytes = wds_dict["msgpack"]
+        size = len(packed_bytes)
+        _current_builder[0].add_sample(
+            key=wds_dict["__key__"],
+            offset=_running_offset[0],
+            size=size,
+            sample=sample,
+        )
+        _running_offset[0] += size
+
+    if use_shard_writer:
+        # Build shard pattern from path
+        if "%" not in str(path):
+            pattern = str(path.parent / f"{path.stem}-%06d{path.suffix}")
+        else:
+            pattern = str(path)
+
+        writer_kwargs: dict[str, Any] = {}
+        if maxcount is not None:
+            writer_kwargs["maxcount"] = maxcount
+        if maxsize is not None:
+            writer_kwargs["maxsize"] = maxsize
+
+        def _track(p: str) -> None:
+            written_paths.append(str(Path(p).resolve()))
+            if manifest and sample_type is not None:
+                _start_builder(p)
+
+        with wds.writer.ShardWriter(pattern, post=_track, **writer_kwargs) as sink:
+            for sample in samples:
+                if sample_type is None:
+                    sample_type = type(sample)
+                wds_dict = sample.as_wds
+                sink.write(wds_dict)
+                if manifest:
+                    # The first sample triggers _track before we get here when
+                    # ShardWriter opens the first shard, but just in case:
+                    if not _current_builder and sample_type is not None:
+                        _start_builder(str(path))
+                    _record_sample(sample, wds_dict)
+    else:
+        with wds.writer.TarWriter(str(path)) as sink:
+            for sample in samples:
+                if sample_type is None:
+                    sample_type = type(sample)
+                wds_dict = sample.as_wds
+                sink.write(wds_dict)
+                if manifest:
+                    if not _current_builder and sample_type is not None:
+                        _current_builder.append(
+                            ManifestBuilder(sample_type=sample_type, shard_id=path.stem)
+                        )
+                    _record_sample(sample, wds_dict)
+        written_paths.append(str(path.resolve()))
+
+    if sample_type is None:
+        raise ValueError("samples must be non-empty")
+
+    # Finalize and write manifests
+    if manifest:
+        _finalize_builder()
+        for shard_path, builder in _builders:
+            m = builder.build()
+            base = str(Path(shard_path).with_suffix(""))
+            writer = ManifestWriter(base)
+            writer.write(m)
+
+    url = _shards_to_wds_url(written_paths)
+    ds: Dataset = Dataset(url)
+    ds._sample_type_cache = sample_type
+    return ds