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

atdata/local/_schema.py

@@ -0,0 +1,380 @@
+"""Schema models and helper functions for local storage."""
+
+from atdata._type_utils import (
+    PRIMITIVE_TYPE_MAP,
+    unwrap_optional,
+    is_ndarray_type,
+    extract_ndarray_dtype,
+    parse_semver,
+)
+from atdata._protocols import Packable
+
+from dataclasses import dataclass, fields, is_dataclass
+from datetime import datetime, timezone
+from typing import (
+    Any,
+    Type,
+    TypeVar,
+    Iterator,
+    Optional,
+    Literal,
+    get_type_hints,
+    get_origin,
+    get_args,
+)
+
+T = TypeVar("T", bound=Packable)
+
+# URI scheme prefixes
+_ATDATA_URI_PREFIX = "atdata://local/sampleSchema/"
+_LEGACY_URI_PREFIX = "local://schemas/"
+
+
+class SchemaNamespace:
+    """Namespace for accessing loaded schema types as attributes.
+
+    After ``index.load_schema(uri)``, the type is available as an attribute.
+    Supports attribute access, iteration, ``len()``, and ``in`` checks.
+
+    Examples:
+        >>> index.load_schema("atdata://local/sampleSchema/MySample@1.0.0")
+        >>> MyType = index.types.MySample
+        >>> sample = MyType(field1="hello", field2=42)
+
+    Note:
+        For full IDE autocomplete, enable ``auto_stubs=True`` and add
+        ``index.stub_dir`` to your IDE's extraPaths.
+    """
+
+    def __init__(self) -> None:
+        self._types: dict[str, Type[Packable]] = {}
+
+    def _register(self, name: str, cls: Type[Packable]) -> None:
+        """Register a schema type in the namespace."""
+        self._types[name] = cls
+
+    def __getattr__(self, name: str) -> Any:
+        # Returns Any to avoid IDE complaints about unknown attributes.
+        # For full IDE support, import from the generated module instead.
+        if name.startswith("_"):
+            raise AttributeError(f"'{type(self).__name__}' has no attribute '{name}'")
+        if name not in self._types:
+            raise AttributeError(
+                f"Schema '{name}' not loaded. "
+                f"Call index.load_schema() first to load the schema."
+            )
+        return self._types[name]
+
+    def __dir__(self) -> list[str]:
+        return list(self._types.keys()) + ["_types", "_register", "get"]
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._types)
+
+    def __len__(self) -> int:
+        return len(self._types)
+
+    def __contains__(self, name: str) -> bool:
+        return name in self._types
+
+    def __repr__(self) -> str:
+        if not self._types:
+            return "SchemaNamespace(empty)"
+        names = ", ".join(sorted(self._types.keys()))
+        return f"SchemaNamespace({names})"
+
+    def get(self, name: str, default: T | None = None) -> Type[Packable] | T | None:
+        """Get a type by name, returning default if not found.
+
+        Args:
+            name: The schema class name to look up.
+            default: Value to return if not found (default: None).
+
+        Returns:
+            The schema class, or default if not loaded.
+        """
+        return self._types.get(name, default)
+
+
+##
+# Schema types
+
+
+@dataclass
+class SchemaFieldType:
+    """Schema field type definition for local storage.
+
+    Represents a type in the schema type system, supporting primitives,
+    ndarrays, arrays, and references to other schemas.
+    """
+
+    kind: Literal["primitive", "ndarray", "ref", "array"]
+    """The category of type."""
+
+    primitive: Optional[str] = None
+    """For kind='primitive': one of 'str', 'int', 'float', 'bool', 'bytes'."""
+
+    dtype: Optional[str] = None
+    """For kind='ndarray': numpy dtype string (e.g., 'float32')."""
+
+    ref: Optional[str] = None
+    """For kind='ref': URI of referenced schema."""
+
+    items: Optional["SchemaFieldType"] = None
+    """For kind='array': type of array elements."""
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "SchemaFieldType":
+        """Create from a dictionary (e.g., from Redis storage)."""
+        type_str = data.get("$type", "")
+        if "#" in type_str:
+            kind = type_str.split("#")[-1]
+        else:
+            kind = data.get("kind", "primitive")
+
+        items = None
+        if "items" in data and data["items"]:
+            items = cls.from_dict(data["items"])
+
+        return cls(
+            kind=kind,  # type: ignore[arg-type]
+            primitive=data.get("primitive"),
+            dtype=data.get("dtype"),
+            ref=data.get("ref"),
+            items=items,
+        )
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for storage."""
+        result: dict[str, Any] = {"$type": f"local#{self.kind}"}
+        if self.kind == "primitive":
+            result["primitive"] = self.primitive
+        elif self.kind == "ndarray":
+            result["dtype"] = self.dtype
+        elif self.kind == "ref":
+            result["ref"] = self.ref
+        elif self.kind == "array" and self.items:
+            result["items"] = self.items.to_dict()
+        return result
+
+
+@dataclass
+class SchemaField:
+    """Schema field definition for local storage."""
+
+    name: str
+    """Field name."""
+
+    field_type: SchemaFieldType
+    """Type of this field."""
+
+    optional: bool = False
+    """Whether this field can be None."""
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "SchemaField":
+        """Create from a dictionary."""
+        return cls(
+            name=data["name"],
+            field_type=SchemaFieldType.from_dict(data["fieldType"]),
+            optional=data.get("optional", False),
+        )
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for storage."""
+        return {
+            "name": self.name,
+            "fieldType": self.field_type.to_dict(),
+            "optional": self.optional,
+        }
+
+
+@dataclass
+class LocalSchemaRecord:
+    """Schema record for local storage.
+
+    Represents a PackableSample schema stored in the local index.
+    Aligns with the atmosphere SchemaRecord structure for seamless promotion.
+    """
+
+    name: str
+    """Schema name (typically the class name)."""
+
+    version: str
+    """Semantic version string (e.g., '1.0.0')."""
+
+    fields: list[SchemaField]
+    """List of field definitions."""
+
+    ref: str
+    """Schema reference URI (atdata://local/sampleSchema/{name}@{version})."""
+
+    description: Optional[str] = None
+    """Human-readable description."""
+
+    created_at: Optional[datetime] = None
+    """When this schema was published."""
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "LocalSchemaRecord":
+        """Create from a dictionary (e.g., from Redis storage)."""
+        created_at = None
+        if "createdAt" in data:
+            try:
+                created_at = datetime.fromisoformat(data["createdAt"])
+            except (ValueError, TypeError):
+                created_at = None  # Invalid datetime format, leave as None
+
+        return cls(
+            name=data["name"],
+            version=data["version"],
+            fields=[SchemaField.from_dict(f) for f in data.get("fields", [])],
+            ref=data.get("$ref", ""),
+            description=data.get("description"),
+            created_at=created_at,
+        )
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for storage."""
+        result: dict[str, Any] = {
+            "name": self.name,
+            "version": self.version,
+            "fields": [f.to_dict() for f in self.fields],
+            "$ref": self.ref,
+        }
+        if self.description:
+            result["description"] = self.description
+        if self.created_at:
+            result["createdAt"] = self.created_at.isoformat()
+        return result
+
+
+##
+# Schema helpers
+
+
+def _kind_str_for_sample_type(st: Type[Packable]) -> str:
+    """Return fully-qualified 'module.name' string for a sample type."""
+    return f"{st.__module__}.{st.__name__}"
+
+
+def _schema_ref_from_type(sample_type: Type[Packable], version: str) -> str:
+    """Generate 'atdata://local/sampleSchema/{name}@{version}' reference."""
+    return _make_schema_ref(sample_type.__name__, version)
+
+
+def _make_schema_ref(name: str, version: str) -> str:
+    """Generate schema reference URI from name and version."""
+    return f"{_ATDATA_URI_PREFIX}{name}@{version}"
+
+
+def _parse_schema_ref(ref: str) -> tuple[str, str]:
+    """Parse schema reference into (name, version).
+
+    Supports both new format: 'atdata://local/sampleSchema/{name}@{version}'
+    and legacy format: 'local://schemas/{module.Class}@{version}'
+    """
+    if ref.startswith(_ATDATA_URI_PREFIX):
+        path = ref[len(_ATDATA_URI_PREFIX) :]
+    elif ref.startswith(_LEGACY_URI_PREFIX):
+        path = ref[len(_LEGACY_URI_PREFIX) :]
+    else:
+        raise ValueError(f"Invalid schema reference: {ref}")
+
+    if "@" not in path:
+        raise ValueError(f"Schema reference must include version (@version): {ref}")
+
+    name, version = path.rsplit("@", 1)
+    # For legacy format, extract just the class name from module.Class
+    if "." in name:
+        name = name.rsplit(".", 1)[1]
+    return name, version
+
+
+def _increment_patch(version: str) -> str:
+    """Increment patch version: 1.0.0 -> 1.0.1"""
+    major, minor, patch = parse_semver(version)
+    return f"{major}.{minor}.{patch + 1}"
+
+
+def _python_type_to_field_type(python_type: Any) -> dict:
+    """Convert Python type annotation to schema field type dict."""
+    if python_type in PRIMITIVE_TYPE_MAP:
+        return {
+            "$type": "local#primitive",
+            "primitive": PRIMITIVE_TYPE_MAP[python_type],
+        }
+
+    if is_ndarray_type(python_type):
+        return {"$type": "local#ndarray", "dtype": extract_ndarray_dtype(python_type)}
+
+    origin = get_origin(python_type)
+    if origin is list:
+        args = get_args(python_type)
+        items = (
+            _python_type_to_field_type(args[0])
+            if args
+            else {"$type": "local#primitive", "primitive": "str"}
+        )
+        return {"$type": "local#array", "items": items}
+
+    if is_dataclass(python_type):
+        raise TypeError(
+            f"Nested dataclass types not yet supported: {python_type.__name__}. "
+            "Publish nested types separately and use references."
+        )
+
+    raise TypeError(f"Unsupported type for schema field: {python_type}")
+
+
+def _build_schema_record(
+    sample_type: Type[Packable],
+    *,
+    version: str,
+    description: str | None = None,
+) -> dict:
+    """Build a schema record dict from a PackableSample type.
+
+    Args:
+        sample_type: The PackableSample subclass to introspect.
+        version: Semantic version string.
+        description: Optional human-readable description. If None, uses the
+            class docstring.
+
+    Returns:
+        Schema record dict suitable for Redis storage.
+
+    Raises:
+        ValueError: If sample_type is not a dataclass.
+        TypeError: If a field type is not supported.
+    """
+    if not is_dataclass(sample_type):
+        raise ValueError(f"{sample_type.__name__} must be a dataclass (use @packable)")
+
+    # Use docstring as fallback for description
+    if description is None:
+        description = sample_type.__doc__
+
+    field_defs = []
+    type_hints = get_type_hints(sample_type)
+
+    for f in fields(sample_type):
+        field_type = type_hints.get(f.name, f.type)
+        field_type, is_optional = unwrap_optional(field_type)
+        field_type_dict = _python_type_to_field_type(field_type)
+
+        field_defs.append(
+            {
+                "name": f.name,
+                "fieldType": field_type_dict,
+                "optional": is_optional,
+            }
+        )
+
+    return {
+        "name": sample_type.__name__,
+        "version": version,
+        "fields": field_defs,
+        "description": description,
+        "createdAt": datetime.now(timezone.utc).isoformat(),
+    }

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}")