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

atdata/dataset.py

@@ -99,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
 
 
@@ -305,7 +307,7 @@ 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)
 
 
@@ -1188,3 +1190,154 @@ def packable(cls: type[_T]) -> type[Packable]:
     ##
 
     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

atdata/index/__init__.py

@@ -0,0 +1,54 @@
+"""Index and entry models for atdata datasets.
+
+Key classes:
+
+- ``Index``: Unified index with pluggable providers (SQLite default),
+  named repositories, and optional atmosphere backend.
+- ``LocalDatasetEntry``: Index entry with ATProto-compatible CIDs.
+"""
+
+from atdata.index._entry import (
+    LocalDatasetEntry,
+    BasicIndexEntry,
+    REDIS_KEY_DATASET_ENTRY,
+    REDIS_KEY_SCHEMA,
+)
+from atdata.index._schema import (
+    SchemaNamespace,
+    SchemaFieldType,
+    SchemaField,
+    LocalSchemaRecord,
+    _ATDATA_URI_PREFIX,
+    _LEGACY_URI_PREFIX,
+    _kind_str_for_sample_type,
+    _schema_ref_from_type,
+    _make_schema_ref,
+    _parse_schema_ref,
+    _increment_patch,
+    _python_type_to_field_type,
+    _build_schema_record,
+)
+from atdata.index._index import Index
+
+__all__ = [
+    # Public API
+    "Index",
+    "LocalDatasetEntry",
+    "BasicIndexEntry",
+    "SchemaNamespace",
+    "SchemaFieldType",
+    "SchemaField",
+    "LocalSchemaRecord",
+    "REDIS_KEY_DATASET_ENTRY",
+    "REDIS_KEY_SCHEMA",
+    # Internal helpers (re-exported for backward compatibility)
+    "_ATDATA_URI_PREFIX",
+    "_LEGACY_URI_PREFIX",
+    "_kind_str_for_sample_type",
+    "_schema_ref_from_type",
+    "_make_schema_ref",
+    "_parse_schema_ref",
+    "_increment_patch",
+    "_python_type_to_field_type",
+    "_build_schema_record",
+]