jerry-thomas 0.3.0__py3-none-any.whl → 1.0.1__py3-none-any.whl

datapipeline/transforms/vector/drop/orchestrator.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from collections.abc import Iterator
+from typing import Literal
+
+from datapipeline.domain.sample import Sample
+
+from .horizontal import VectorDropHorizontalTransform
+from .vertical import VectorDropVerticalTransform
+
+Axis = Literal["horizontal", "vertical"]
+
+
+class VectorDropTransform:
+    """Drop vectors or features based on coverage thresholds.
+
+    Thin orchestrator that delegates to horizontal or vertical strategies based
+    on the configured axis.
+    """
+
+    def __init__(
+        self,
+        *,
+        axis: Axis = "horizontal",
+        threshold: float,
+        payload: Literal["features", "targets", "both"] = "features",
+        only: list[str] | None = None,
+        exclude: list[str] | None = None,
+    ) -> None:
+        if axis not in {"horizontal", "vertical"}:
+            raise ValueError("axis must be 'horizontal' or 'vertical'")
+        if axis == "vertical" and payload == "both":
+            raise ValueError("axis='vertical' does not support payload='both'")
+        if axis == "horizontal":
+            self._impl: object = VectorDropHorizontalTransform(
+                threshold=threshold,
+                payload=payload,
+                only=only,
+                exclude=exclude,
+            )
+        else:
+            # Vertical drop is partition/feature-oriented and does not support
+            # payload='both'. Payload is validated above.
+            self._impl = VectorDropVerticalTransform(
+                payload=payload if payload != "both" else "features",
+                threshold=threshold,
+            )
+
+    def bind_context(self, context) -> None:
+        binder = getattr(self._impl, "bind_context", None)
+        if binder is not None:
+            binder(context)
+
+    def __call__(self, stream: Iterator[Sample]) -> Iterator[Sample]:
+        return self.apply(stream)
+
+    def apply(self, stream: Iterator[Sample]) -> Iterator[Sample]:
+        return getattr(self._impl, "apply")(stream)
+

datapipeline/transforms/vector/drop/vertical.py

@@ -0,0 +1,182 @@
+from __future__ import annotations
+
+from collections.abc import Iterator
+from typing import Literal
+
+from datapipeline.config.metadata import (
+    FEATURE_VECTORS_COUNT_KEY,
+    TARGET_VECTORS_COUNT_KEY,
+    VectorMetadata,
+)
+from datapipeline.domain.sample import Sample
+from datapipeline.domain.vector import Vector
+from datapipeline.services.artifacts import (
+    ArtifactNotRegisteredError,
+    VECTOR_METADATA_SPEC,
+)
+
+from ..common import (
+    VectorContextMixin,
+    replace_vector,
+    select_vector,
+    try_get_current_context,
+)
+
+
+class VectorDropVerticalTransform(VectorContextMixin):
+    required_artifacts = {VECTOR_METADATA_SPEC.key}
+    """Drop partitions/features when metadata coverage falls below configured thresholds.
+
+    Requires the optional `metadata.json` artifact generated by the
+    `metadata` build task. The transform evaluates coverage using the recorded
+    `present_count`/`null_count` metrics and prunes the schema cache once so
+    downstream coverage checks stop expecting bad partitions.
+    """
+
+    def __init__(
+        self,
+        *,
+        payload: Literal["features", "targets"] = "features",
+        threshold: float,
+    ) -> None:
+        super().__init__(payload=payload)
+        if not 0.0 <= threshold <= 1.0:
+            raise ValueError("threshold must be between 0 and 1.")
+        self._threshold = threshold
+        self._drop_ids: set[str] | None = None
+        self._schema_pruned = False
+
+    def __call__(self, stream: Iterator[Sample]) -> Iterator[Sample]:
+        return self.apply(stream)
+
+    def apply(self, stream: Iterator[Sample]) -> Iterator[Sample]:
+        drop_ids = self._resolve_drop_ids()
+        if not drop_ids:
+            yield from stream
+            return
+        self._maybe_prune_schema(drop_ids)
+        for sample in stream:
+            if not self._schema_pruned:
+                self._maybe_prune_schema(drop_ids)
+            vector = select_vector(sample, self._payload)
+            if vector is None or not vector.values:
+                yield sample
+                continue
+            retained: dict[str, object] = {}
+            changed = False
+            for fid, value in vector.values.items():
+                if fid in drop_ids:
+                    changed = True
+                    continue
+                retained[fid] = value
+            if not changed:
+                yield sample
+            else:
+                yield replace_vector(sample, self._payload, Vector(values=retained))
+
+    def _resolve_drop_ids(self) -> set[str]:
+        if self._drop_ids is not None:
+            return self._drop_ids
+        context = self._context or try_get_current_context()
+        if not context:
+            raise RuntimeError("VectorDropVerticalTransform requires an active pipeline context.")
+        try:
+            raw = context.require_artifact(VECTOR_METADATA_SPEC)
+        except ArtifactNotRegisteredError as exc:
+            raise RuntimeError(
+                "Vector metadata artifact missing. Enable the `metadata` build task "
+                "and rerun `jerry build --project <project.yaml>`."
+            ) from exc
+        meta = VectorMetadata.model_validate(raw)
+        section_key = "targets" if self._payload == "targets" else "features"
+        counts_key = (
+            TARGET_VECTORS_COUNT_KEY
+            if self._payload == "targets"
+            else FEATURE_VECTORS_COUNT_KEY
+        )
+
+        entries = getattr(meta, section_key) or []
+        window_size = self._window_size(getattr(meta, "window", None))
+        total = window_size if window_size is not None else meta.counts.get(counts_key)
+        if not isinstance(total, (int, float)) or total <= 0:
+            if self._payload == "targets":
+                raise RuntimeError(
+                    "Vector metadata artifact missing counts for targets; "
+                    "ensure your dataset defines target streams and rebuild."
+                )
+            raise RuntimeError(
+                "Vector metadata artifact missing counts for features; "
+                "rerun `jerry build --project <project.yaml>` to refresh metadata."
+            )
+        expected_buckets = float(total)
+        drop_ids: set[str] = set()
+        for entry in entries:
+            if not isinstance(entry, dict):
+                continue
+            fid = entry.get("id")
+            if not isinstance(fid, str):
+                continue
+            coverage = self._coverage_for_entry(entry, expected_buckets)
+            if coverage < self._threshold:
+                drop_ids.add(fid)
+        self._drop_ids = drop_ids
+        return drop_ids
+
+    @staticmethod
+    def _window_size(window) -> float | None:
+        if window is None:
+            return None
+        if isinstance(window, dict):
+            return window.get("size")
+        return getattr(window, "size", None)
+
+    @staticmethod
+    def _coverage_for_entry(entry: dict, expected_buckets: float) -> float:
+        if expected_buckets <= 0:
+            return 0.0
+        present = float(entry.get("present_count") or 0.0)
+        nulls = float(entry.get("null_count") or 0.0)
+        cadence_doc = entry.get("cadence")
+        cadence = cadence_doc.get("target") if isinstance(cadence_doc, dict) else None
+        observed_elements = entry.get("observed_elements")
+
+        if isinstance(observed_elements, (int, float)) and cadence:
+            # Base expected elements on buckets where this feature actually appeared
+            # to avoid over-crediting sparse sequences.
+            expected_elements = float(max(present, 0.0)) * float(cadence)
+            if expected_elements > 0:
+                return max(
+                    0.0,
+                    min(1.0, float(observed_elements) / expected_elements),
+                )
+
+        coverage = (present - nulls) / expected_buckets
+        return max(0.0, min(1.0, coverage))
+
+    def _maybe_prune_schema(self, drop_ids: set[str]) -> None:
+        if self._schema_pruned or not drop_ids:
+            return
+        context = self._context or try_get_current_context()
+        if not context:
+            self._schema_pruned = True
+            return
+        cache = getattr(context, "_cache", None)
+        if cache is None:
+            self._schema_pruned = True
+            return
+        schema_key = f"schema:{self._payload}"
+        if schema_key not in cache:
+            return
+        entries = cache.get(schema_key)
+        if not entries:
+            self._schema_pruned = True
+            return
+        kept = [entry for entry in entries if entry.get("id") not in drop_ids]
+        cache[schema_key] = kept
+        ids_key = f"expected_ids:{self._payload}"
+        cache[ids_key] = [
+            entry.get("id")
+            for entry in kept
+            if isinstance(entry.get("id"), str)
+        ]
+        self._schema_pruned = True

datapipeline/transforms/vector/ensure_schema.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+from collections import OrderedDict
+from collections.abc import Iterator
+from typing import Any, Literal
+
+from datapipeline.domain.sample import Sample
+from datapipeline.domain.vector import Vector
+from datapipeline.transforms.vector_utils import clone
+
+from .common import VectorContextMixin, replace_vector, select_vector
+
+MissingPolicy = Literal["error", "drop", "fill"]
+ExtraPolicy = Literal["error", "drop", "keep"]
+
+
+class VectorEnsureSchemaTransform(VectorContextMixin):
+    """Ensure vectors conform to the vector schema (`schema.json`) artifact.
+
+    Options allow filling or dropping rows with missing identifiers and
+    pruning/raising on unexpected identifiers.
+    """
+
+    def __init__(
+        self,
+        *,
+        payload: Literal["features", "targets"] = "features",
+        on_missing: MissingPolicy = "error",
+        fill_value: Any = None,
+        on_extra: ExtraPolicy = "error",
+    ) -> None:
+        super().__init__(payload=payload)
+        if on_missing not in {"error", "drop", "fill"}:
+            raise ValueError("on_missing must be one of: 'error', 'drop', 'fill'")
+        if on_extra not in {"error", "drop", "keep"}:
+            raise ValueError("on_extra must be one of: 'error', 'drop', 'keep'")
+        self._on_missing = on_missing
+        self._fill_value = fill_value
+        self._on_extra = on_extra
+        self._baseline: list[str] | None = None
+        self._schema_entries: list[dict[str, Any]] | None = None
+        self._schema_meta: dict[str, dict[str, Any]] = {}
+
+    def __call__(self, stream: Iterator[Sample]) -> Iterator[Sample]:
+        return self.apply(stream)
+
+    def apply(self, stream: Iterator[Sample]) -> Iterator[Sample]:
+        baseline = self._schema_ids()
+        baseline_set = set(baseline)
+
+        for sample in stream:
+            vector = select_vector(sample, self._payload)
+            if vector is None:
+                yield sample
+                continue
+
+            values = vector.values
+            working = None
+
+            missing = [fid for fid in baseline if fid not in values]
+            if missing:
+                decision = self._on_missing
+                if decision == "error":
+                    raise ValueError(
+                        f"Vector missing required identifiers {missing} "
+                        f"for payload '{self._payload}'."
+                    )
+                if decision == "drop":
+                    continue
+                working = clone(values)
+                for fid in missing:
+                    working[fid] = self._fill_value
+
+            extras = [fid for fid in values if fid not in baseline_set]
+            if extras:
+                decision = self._on_extra
+                if decision == "error":
+                    raise ValueError(
+                        f"Vector contains unexpected identifiers {extras} "
+                        f"for payload '{self._payload}'."
+                    )
+                if decision == "drop":
+                    working = working or clone(values)
+                    for fid in extras:
+                        working.pop(fid, None)
+
+            current_values = working or values
+
+            # Optionally enforce per-id cadence from schema metadata
+            current_values = self._enforce_cadence(current_values)
+
+            ordered = OrderedDict()
+            for fid in baseline:
+                ordered[fid] = current_values.get(fid)
+            if self._on_extra == "keep":
+                for fid, value in current_values.items():
+                    if fid not in baseline_set:
+                        ordered[fid] = value
+            current_values = ordered
+
+            if current_values is not values:
+                updated_vector = Vector(values=dict(current_values))
+                sample = replace_vector(sample, self._payload, updated_vector)
+
+            yield sample
+
+    def _schema_ids(self) -> list[str]:
+        if self._baseline is None:
+            entries = self._load_schema_entries()
+            ordered = [entry["id"] for entry in entries if isinstance(entry.get("id"), str)]
+            if not ordered:
+                raise RuntimeError(
+                    "Vector schema artifact is empty or unavailable; run `jerry build` "
+                    "to materialize `schema.json` via the `vector_schema` task."
+                )
+            self._baseline = ordered
+            self._schema_meta = {
+                entry["id"]: entry for entry in entries if isinstance(entry.get("id"), str)
+            }
+        return list(self._baseline)
+
+    def _load_schema_entries(self) -> list[dict[str, Any]]:
+        if self._schema_entries is None:
+            context = getattr(self, "_context", None)
+            if not context:
+                entries = []
+            else:
+                entries = context.load_schema(payload=self._payload)
+            self._schema_entries = entries or []
+        return self._schema_entries
+
+    def _enforce_cadence(self, values: dict[str, Any]) -> dict[str, Any]:
+        if not values or not self._schema_meta:
+            return values
+        adjusted = None
+        for fid, value in values.items():
+            meta = self._schema_meta.get(fid)
+            if not meta or meta.get("kind") != "list":
+                continue
+            expected = self._expected_lengths(meta)
+            if not expected:
+                continue
+            current_len = len(value) if isinstance(value, list) else (0 if value is None else 1)
+            if current_len in expected:
+                continue
+            decision = self._on_missing
+            if decision == "error":
+                raise ValueError(
+                    f"List feature '{fid}' length {current_len} violates schema cadence {sorted(expected)}"
+                )
+            if decision == "drop":
+                return {}
+            # fill: pad or truncate to the closest expected length
+            target_len = expected[0]
+            adjusted = adjusted or clone(values)
+            if isinstance(value, list):
+                seq = value[:target_len]
+            elif value is None:
+                seq = []
+            else:
+                seq = [value]
+            if len(seq) < target_len:
+                seq = seq + [self._fill_value] * (target_len - len(seq))
+            adjusted[fid] = seq
+        return adjusted or values
+
+    def _expected_lengths(self, meta: dict[str, Any]) -> list[int]:
+        cadence = meta.get("cadence")
+        if isinstance(cadence, dict):
+            target = cadence.get("target")
+            if isinstance(target, (int, float)) and target > 0:
+                return [int(target)]
+        modes = meta.get("list_length", {}).get("modes")
+        if isinstance(modes, (list, tuple)) and modes:
+            ints = [int(m) for m in modes if isinstance(m, (int, float))]
+            if ints:
+                return sorted(ints)
+        expected = meta.get("expected_length")
+        if isinstance(expected, (int, float)):
+            return [int(expected)]
+        max_len = meta.get("list_length", {}).get("max")
+        if isinstance(max_len, (int, float)) and max_len > 0:
+            return [int(max_len)]
+        return []

datapipeline/transforms/vector/fill.py

@@ -0,0 +1,87 @@
+from collections import deque
+from collections.abc import Iterator
+from statistics import mean, median
+from typing import Any, Literal
+
+from datapipeline.domain.sample import Sample
+from datapipeline.domain.vector import Vector
+from datapipeline.transforms.vector_utils import clone, is_missing
+
+from .common import VectorPostprocessBase, replace_vector, select_vector
+
+
+class VectorFillTransform(VectorPostprocessBase):
+    """Fill missing entries using running statistics from prior buckets."""
+
+    def __init__(
+        self,
+        *,
+        statistic: Literal["mean", "median"] = "median",
+        window: int | None = None,
+        min_samples: int = 1,
+        payload: Literal["features", "targets", "both"] = "features",
+        only: list[str] | None = None,
+        exclude: list[str] | None = None,
+    ) -> None:
+        super().__init__(payload=payload, only=only, exclude=exclude)
+        if window is not None and window <= 0:
+            raise ValueError("window must be positive when provided")
+        if min_samples <= 0:
+            raise ValueError("min_samples must be positive")
+        self.statistic = statistic
+        self.window = window
+        self.min_samples = min_samples
+        self.history: dict[str, deque[float]] = {}
+
+    def _compute(self, feature_id: str) -> float | None:
+        values = self.history.get(feature_id)
+        if not values or len(values) < self.min_samples:
+            return None
+        if self.statistic == "mean":
+            return float(mean(values))
+        return float(median(values))
+
+    def _push(self, feature_id: str, value: Any) -> None:
+        if is_missing(value):
+            return
+        try:
+            num = float(value)
+        except (TypeError, ValueError):
+            return
+        bucket = self.history.setdefault(str(feature_id), deque(maxlen=self.window))
+        bucket.append(num)
+
+    def __call__(self, stream: Iterator[Sample]) -> Iterator[Sample]:
+        return self.apply(stream)
+
+    def apply(self, stream: Iterator[Sample]) -> Iterator[Sample]:
+        for sample in stream:
+            for kind in self._payload_kinds():
+                ids = self._ids_for(kind)
+                if ids:
+                    sample = self._apply_to_payload(sample, kind, ids)
+            yield sample
+
+    def _apply_to_payload(
+        self,
+        sample: Sample,
+        payload: Literal["features", "targets"],
+        ids: list[str],
+    ) -> Sample:
+        vector = select_vector(sample, payload)
+        if vector is None:
+            return sample
+        data = clone(vector.values)
+        updated = False
+        for feature in ids:
+            if feature in data and not is_missing(data[feature]):
+                continue
+            fill = self._compute(feature)
+            if fill is not None:
+                data[feature] = fill
+                updated = True
+        for fid, value in data.items():
+            self._push(fid, value)
+        if not updated:
+            return sample
+        return replace_vector(sample, payload, Vector(values=data))

datapipeline/transforms/vector/replace.py

@@ -0,0 +1,62 @@
+from collections.abc import Iterator
+from typing import Any, Literal
+
+from datapipeline.domain.sample import Sample
+from datapipeline.domain.vector import Vector
+from datapipeline.transforms.vector_utils import clone, is_missing
+
+from .common import VectorPostprocessBase, replace_vector, select_vector
+
+
+class VectorReplaceTransform(VectorPostprocessBase):
+    """Fill missing entries with a constant value."""
+
+    def __init__(
+        self,
+        *,
+        value: Any,
+        payload: Literal["features", "targets", "both"] = "features",
+        only: list[str] | None = None,
+        exclude: list[str] | None = None,
+        target: Any | None = None,
+    ) -> None:
+        super().__init__(payload=payload, only=only, exclude=exclude)
+        self.value = value
+        self._target = target
+
+    def __call__(self, stream: Iterator[Sample]) -> Iterator[Sample]:
+        return self.apply(stream)
+
+    def apply(self, stream: Iterator[Sample]) -> Iterator[Sample]:
+        for sample in stream:
+            for kind in self._payload_kinds():
+                ids = self._ids_for(kind)
+                if ids:
+                    sample = self._apply_to_payload(sample, kind, ids)
+            yield sample
+
+    def _should_replace(self, value: Any) -> bool:
+        if self._target is None:
+            return is_missing(value)
+        return value == self._target
+
+    def _apply_to_payload(
+        self,
+        sample: Sample,
+        payload: Literal["features", "targets"],
+        ids: list[str],
+    ) -> Sample:
+        vector = select_vector(sample, payload)
+        if vector is None:
+            return sample
+        data = clone(vector.values)
+        updated = False
+        for feature in ids:
+            current = data.get(feature)
+            if not self._should_replace(current):
+                continue
+            data[feature] = self.value
+            updated = True
+        if not updated:
+            return sample
+        return replace_vector(sample, payload, Vector(values=data))

datapipeline/utils/load.py

@@ -1,19 +1,40 @@
+import importlib
 import importlib.metadata as md
 from functools import lru_cache
-import yaml
 from pathlib import Path
 
+import yaml
+
+# Local fallback map so newly added entrypoints remain usable in editable installs
+# before package metadata is refreshed.
+_EP_OVERRIDES = {}
+
 
 @lru_cache
 def load_ep(group: str, name: str):
+    target = _EP_OVERRIDES.get((group, name))
+    if target:
+        module, attr = target.split(":")
+        return getattr(importlib.import_module(module), attr)
+
     eps = md.entry_points().select(group=group, name=name)
     if not eps:
         available = ", ".join(
-            sorted(ep.name for ep in md.entry_points().select(group=group)))
+            sorted(ep.name for ep in md.entry_points().select(group=group))
+        )
         raise ValueError(
             f"No entry point '{name}' in '{group}'. Available: {available or '(none)'}")
     if len(eps) > 1:
-        mods = ", ".join(f"{ep.module}:{ep.attr}" for ep in eps)
+        def describe(ep):
+            value = getattr(ep, "value", None)
+            if value:
+                return value
+            module = getattr(ep, "module", None)
+            attr = getattr(ep, "attr", None)
+            if module and attr:
+                return f"{module}:{attr}"
+            return repr(ep)
+        mods = ", ".join(describe(ep) for ep in eps)
         raise ValueError(
             f"Ambiguous entry point '{name}' in '{group}': {mods}")
     # EntryPoints in newer Python versions are mapping-like; avoid integer indexing

datapipeline/utils/rich_compat.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+
+def suppress_file_proxy_shutdown_errors() -> None:
+    """Patch rich.file_proxy.FileProxy.flush to ignore shutdown ImportErrors.
+
+    Rich leaves behind FileProxy instances that may flush while the interpreter
+    is tearing down, which triggers `ImportError: sys.meta_path is None`.
+    Swallow those benign errors so CLI commands exit cleanly.
+    """
+    try:
+        from rich.file_proxy import FileProxy
+    except Exception:
+        return
+
+    if getattr(FileProxy, "_datapipeline_safe_flush", False):
+        return
+
+    original_flush = FileProxy.flush
+
+    def _safe_flush(self) -> None:  # type: ignore[override]
+        try:
+            original_flush(self)
+        except ImportError as exc:
+            if "sys.meta_path is None" in str(exc):
+                return
+            raise
+        except RuntimeError as exc:
+            message = str(exc)
+            if "shutting down" in message.lower():
+                return
+            raise
+
+    FileProxy.flush = _safe_flush  # type: ignore[assignment]
+    setattr(FileProxy, "_datapipeline_safe_flush", True)
+
+
+__all__ = ["suppress_file_proxy_shutdown_errors"]