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

datapipeline/build/tasks/metadata.py

@@ -0,0 +1,170 @@
+from __future__ import annotations
+
+import json
+from collections import defaultdict
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Dict, Tuple
+
+from datapipeline.config.dataset.loader import load_dataset
+from datapipeline.config.metadata import (
+    VectorMetadata,
+    Window,
+    FEATURE_VECTORS_COUNT_KEY,
+    TARGET_VECTORS_COUNT_KEY,
+)
+from datapipeline.config.tasks import MetadataTask
+from datapipeline.runtime import Runtime
+from datapipeline.utils.paths import ensure_parent
+from datapipeline.config.dataset.normalize import floor_time_to_bucket
+from datapipeline.utils.time import parse_timecode
+
+from .utils import collect_schema_entries, metadata_entries_from_stats
+
+
+def _entry_window(entry: dict) -> tuple[datetime | None, datetime | None]:
+    start = entry.get("first_ts")
+    end = entry.get("last_ts")
+    return (start if isinstance(start, datetime) else None, end if isinstance(end, datetime) else None)
+
+
+def _group_ranges(entries: list[dict], key_name: str) -> list[tuple[datetime, datetime]]:
+    grouped: dict[str, list[tuple[datetime, datetime]]] = defaultdict(list)
+    for entry in entries:
+        start, end = _entry_window(entry)
+        if start is None or end is None:
+            continue
+        group_key = entry.get(key_name) or entry.get("id")
+        if not isinstance(group_key, str):
+            continue
+        grouped[group_key].append((start, end))
+    ranges: list[tuple[datetime, datetime]] = []
+    for values in grouped.values():
+        group_start = min(start for start, _ in values)
+        group_end = max(end for _, end in values)
+        ranges.append((group_start, group_end))
+    return ranges
+
+
+def _range_union(ranges):
+    if not ranges:
+        return None, None
+    start = min(r[0] for r in ranges)
+    end = max(r[1] for r in ranges)
+    if start >= end:
+        return None, None
+    return start, end
+
+
+def _range_intersection(ranges):
+    if not ranges:
+        return None, None
+    start = max(r[0] for r in ranges)
+    end = min(r[1] for r in ranges)
+    if start >= end:
+        return None, None
+    return start, end
+
+
+def _window_bounds_from_stats(
+    feature_stats: list[dict],
+    target_stats: list[dict],
+    *,
+    mode: str,
+) -> tuple[datetime | None, datetime | None]:
+    base_ranges = _group_ranges(
+        feature_stats, "base_id") + _group_ranges(target_stats, "base_id")
+    partition_ranges = _group_ranges(
+        feature_stats, "id") + _group_ranges(target_stats, "id")
+
+    if mode == "intersection":
+        return _range_intersection(base_ranges)
+    if mode == "strict":
+        return _range_intersection(partition_ranges)
+    if mode == "relaxed":
+        return _range_union(partition_ranges)
+    # default to union
+    return _range_union(base_ranges if base_ranges else partition_ranges)
+
+
+def _window_size(start: datetime | None, end: datetime | None, cadence: str | None) -> int | None:
+    if start is None or end is None or cadence is None:
+        return None
+    try:
+        anchored_start = floor_time_to_bucket(start, cadence)
+        anchored_end = floor_time_to_bucket(end, cadence)
+        step = parse_timecode(cadence)
+        if anchored_end < anchored_start:
+            return None
+        return int(((anchored_end - anchored_start) / step)) + 1
+    except Exception:
+        return None
+
+
+def materialize_metadata(runtime: Runtime, task_cfg: MetadataTask) -> Tuple[str, Dict[str, object]] | None:
+    if not task_cfg.enabled:
+        return None
+    dataset = load_dataset(runtime.project_yaml, "vectors")
+    features_cfgs = list(dataset.features or [])
+    feature_stats, feature_vectors, feature_min, feature_max = collect_schema_entries(
+        runtime,
+        features_cfgs,
+        dataset.group_by,
+        cadence_strategy=task_cfg.cadence_strategy,
+        collect_metadata=True,
+    )
+    target_meta: list[dict] = []
+    target_vectors = 0
+    target_cfgs = list(dataset.targets or [])
+    target_stats: list[dict] = []
+    target_min = target_max = None
+    if target_cfgs:
+        target_stats, target_vectors, target_min, target_max = collect_schema_entries(
+            runtime,
+            target_cfgs,
+            dataset.group_by,
+            cadence_strategy=task_cfg.cadence_strategy,
+            collect_metadata=True,
+        )
+        target_meta = metadata_entries_from_stats(
+            target_stats, task_cfg.cadence_strategy)
+    feature_meta = metadata_entries_from_stats(
+        feature_stats, task_cfg.cadence_strategy)
+
+    generated_at = datetime.now(timezone.utc)
+    window_obj: Window | None = None
+    computed_start, computed_end = _window_bounds_from_stats(
+        feature_stats,
+        target_stats if target_cfgs else [],
+        mode=task_cfg.window_mode,
+    )
+    start = computed_start
+    end = computed_end
+    if start is not None and end is not None and start < end:
+        size = _window_size(start, end, dataset.group_by)
+        window_obj = Window(start=start, end=end,
+                            mode=task_cfg.window_mode, size=size)
+
+    doc = VectorMetadata(
+        schema_version=1,
+        generated_at=generated_at,
+        features=feature_meta,
+        targets=target_meta,
+        counts={
+            FEATURE_VECTORS_COUNT_KEY: feature_vectors,
+            TARGET_VECTORS_COUNT_KEY: target_vectors,
+        },
+        window=window_obj,
+    )
+
+    relative_path = Path(task_cfg.output)
+    destination = (runtime.artifacts_root / relative_path).resolve()
+    ensure_parent(destination)
+    with destination.open("w", encoding="utf-8") as fh:
+        json.dump(doc.model_dump(mode="json"), fh, indent=2)
+
+    meta: Dict[str, object] = {
+        "features": len(feature_meta),
+        "targets": len(target_meta),
+    }
+    return str(relative_path), meta

datapipeline/build/tasks/scaler.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Dict, Iterator, Tuple
+
+from datapipeline.config.tasks import ScalerTask
+from datapipeline.config.dataset.loader import load_dataset
+from datapipeline.domain.sample import Sample
+from datapipeline.pipeline.context import PipelineContext
+from datapipeline.pipeline.pipelines import build_vector_pipeline
+from datapipeline.pipeline.split import build_labeler
+from datapipeline.runtime import Runtime
+from datapipeline.transforms.feature.scaler import StandardScaler
+from datapipeline.utils.paths import ensure_parent
+
+
+def materialize_scaler_statistics(runtime: Runtime, task_cfg: ScalerTask) -> Tuple[str, Dict[str, object]] | None:
+    if not task_cfg.enabled:
+        return None
+
+    dataset = load_dataset(runtime.project_yaml, "vectors")
+    feature_cfgs = list(dataset.features or [])
+    target_cfgs = list(dataset.targets or [])
+    if not feature_cfgs and not target_cfgs:
+        return None
+
+    sanitized_features = [cfg.model_copy(update={"scale": False}) for cfg in feature_cfgs]
+    sanitized_targets = [cfg.model_copy(update={"scale": False}) for cfg in target_cfgs]
+
+    context = PipelineContext(runtime)
+    vectors = build_vector_pipeline(
+        context,
+        sanitized_features,
+        dataset.group_by,
+        target_configs=sanitized_targets,
+        rectangular=False,
+    )
+
+    cfg = getattr(runtime, "split", None)
+    labeler = build_labeler(cfg) if cfg else None
+    if not labeler and task_cfg.split_label != "all":
+        raise RuntimeError(
+            f"Cannot compute scaler statistics for split '{task_cfg.split_label}' "
+            "when no split configuration is defined in the project."
+        )
+
+    def _train_stream() -> Iterator[Sample]:
+        for sample in vectors:
+            if labeler and labeler.label(sample.key, sample.features) != task_cfg.split_label:
+                continue
+            yield sample
+
+    scaler = StandardScaler()
+    total_observations = scaler.fit(_train_stream())
+
+    if not scaler.statistics:
+        raise RuntimeError(
+            f"No scaler statistics computed for split '{task_cfg.split_label}'."
+        )
+
+    relative_path = Path(task_cfg.output)
+    destination = (runtime.artifacts_root / relative_path).resolve()
+    ensure_parent(destination)
+
+    scaler.save(destination)
+
+    meta: Dict[str, object] = {
+        "features": len(scaler.statistics),
+        "split": task_cfg.split_label,
+        "observations": total_observations,
+    }
+
+    return str(relative_path), meta

datapipeline/build/tasks/schema.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+import json
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Dict, Tuple
+
+from datapipeline.config.tasks import SchemaTask
+from datapipeline.config.dataset.loader import load_dataset
+from datapipeline.runtime import Runtime
+from datapipeline.utils.paths import ensure_parent
+from datapipeline.utils.window import resolve_window_bounds
+
+from .utils import collect_schema_entries, schema_entries_from_stats
+
+
+def materialize_vector_schema(runtime: Runtime, task_cfg: SchemaTask) -> Tuple[str, Dict[str, object]] | None:
+    if not task_cfg.enabled:
+        return None
+    dataset = load_dataset(runtime.project_yaml, "vectors")
+    features_cfgs = list(dataset.features or [])
+    feature_stats, feature_vectors, feature_min, feature_max = collect_schema_entries(
+        runtime,
+        features_cfgs,
+        dataset.group_by,
+        cadence_strategy=task_cfg.cadence_strategy,
+        collect_metadata=False,
+    )
+    target_entries: list[dict] = []
+    target_cfgs = list(dataset.targets or [])
+    target_min = target_max = None
+    if target_cfgs:
+        target_stats, _, target_min, target_max = collect_schema_entries(
+            runtime,
+            target_cfgs,
+            dataset.group_by,
+            cadence_strategy=task_cfg.cadence_strategy,
+            collect_metadata=False,
+        )
+        target_entries = schema_entries_from_stats(target_stats, task_cfg.cadence_strategy)
+    feature_entries = schema_entries_from_stats(feature_stats, task_cfg.cadence_strategy)
+
+    doc = {
+        "schema_version": 1,
+        "generated_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
+    }
+    doc["features"] = feature_entries
+    doc["targets"] = target_entries
+
+    relative_path = Path(task_cfg.output)
+    destination = (runtime.artifacts_root / relative_path).resolve()
+    ensure_parent(destination)
+    with destination.open("w", encoding="utf-8") as fh:
+        json.dump(doc, fh, indent=2)
+
+    meta: Dict[str, object] = {
+        "features": len(feature_entries),
+        "targets": len(target_entries),
+    }
+    return str(relative_path), meta

datapipeline/build/tasks/utils.py

@@ -0,0 +1,169 @@
+from __future__ import annotations
+
+from collections import Counter, OrderedDict
+from datetime import datetime
+from typing import Any
+
+from datapipeline.pipeline.context import PipelineContext
+from datapipeline.pipeline.pipelines import build_vector_pipeline
+from datapipeline.runtime import Runtime
+from datapipeline.transforms.vector_utils import base_id as _base_feature_id
+from datapipeline.transforms.utils import is_missing
+
+
+def _type_name(value: object) -> str:
+    if value is None:
+        return "null"
+    return type(value).__name__
+
+
+def collect_schema_entries(
+    runtime: Runtime,
+    configs,
+    group_by: str,
+    *,
+    cadence_strategy: str,
+    collect_metadata: bool,
+) -> tuple[list[dict], int, datetime | None, datetime | None]:
+    configs = list(configs or [])
+    if not configs:
+        return [], 0, None, None
+    sanitized = [cfg.model_copy(update={"scale": False}) for cfg in configs]
+    context = PipelineContext(runtime)
+    vectors = build_vector_pipeline(
+        context,
+        sanitized,
+        group_by,
+        rectangular=False,
+    )
+
+    stats: OrderedDict[str, dict] = OrderedDict()
+    vector_count = 0
+    min_time: datetime | None = None
+    max_time: datetime | None = None
+    for sample in vectors:
+        vector_count += 1
+        ts = sample.key[0] if isinstance(sample.key, tuple) and sample.key else None
+        if isinstance(ts, datetime):
+            min_time = ts if min_time is None else min(min_time, ts)
+            max_time = ts if max_time is None else max(max_time, ts)
+        payload = sample.features
+        for fid, value in payload.values.items():
+            entry = stats.get(fid)
+            if not entry:
+                entry = stats[fid] = {
+                    "id": fid,
+                    "base_id": _base_feature_id(fid),
+                    "kind": None,
+                    "max_length": None,
+                    "present_count": 0,
+                    "null_count": 0,
+                    "scalar_types": set(),
+                    "element_types": set(),
+                    "min_length": None,
+                    "lengths": Counter(),
+                    "first_ts": None,
+                    "last_ts": None,
+                }
+            if isinstance(ts, datetime):
+                prev_start = entry.get("first_ts")
+                entry["first_ts"] = ts if prev_start is None else min(prev_start, ts)
+                prev_end = entry.get("last_ts")
+                entry["last_ts"] = ts if prev_end is None else max(prev_end, ts)
+            if collect_metadata:
+                entry["present_count"] += 1
+            if is_missing(value):
+                if collect_metadata:
+                    entry["null_count"] += 1
+                continue
+            if isinstance(value, list):
+                entry["kind"] = "list"
+                length = len(value)
+                entry["min_length"] = length if entry["min_length"] is None else min(
+                    entry["min_length"], length
+                )
+                entry["max_length"] = length if entry["max_length"] is None else max(
+                    entry["max_length"], length
+                )
+                if collect_metadata:
+                    entry["lengths"][length] += 1
+                    entry["observed_elements"] = entry.get("observed_elements", 0) + sum(
+                        1 for v in value if not is_missing(v)
+                    )
+                    if not value:
+                        entry["element_types"].add("empty")
+                    else:
+                        entry["element_types"].update(_type_name(v) for v in value)
+            else:
+                if entry["kind"] != "list":
+                    entry["kind"] = "scalar"
+                if collect_metadata:
+                    entry["scalar_types"].add(_type_name(value))
+
+    return list(stats.values()), vector_count, min_time, max_time
+
+
+def _resolve_cadence_target(stats: dict, strategy: str) -> int | None:
+    if strategy == "max":
+        max_len = stats.get("max_length")
+        if isinstance(max_len, (int, float)) and max_len > 0:
+            return int(max_len)
+    return None
+
+
+def schema_entries_from_stats(entries: list[dict], cadence_strategy: str) -> list[dict]:
+    doc: list[dict] = []
+    for entry in entries:
+        kind = entry.get("kind") or "scalar"
+        item = {
+            "id": entry["id"],
+            "base_id": entry["base_id"],
+            "kind": kind,
+        }
+        if kind == "list":
+            target = _resolve_cadence_target(entry, cadence_strategy)
+            if target is not None:
+                item["cadence"] = {"strategy": cadence_strategy, "target": target}
+        doc.append(item)
+    return doc
+
+
+def _to_iso(ts: datetime | None) -> str | None:
+    if isinstance(ts, datetime):
+        text = ts.isoformat()
+        if text.endswith("+00:00"):
+            return text[:-6] + "Z"
+        return text
+    return None
+
+
+def metadata_entries_from_stats(entries: list[dict], cadence_strategy: str) -> list[dict]:
+    meta_entries: list[dict] = []
+    for entry in entries:
+        kind = entry.get("kind") or "scalar"
+        item: dict[str, Any] = {
+            "id": entry["id"],
+            "base_id": entry["base_id"],
+            "kind": kind,
+            "present_count": entry.get("present_count", 0),
+            "null_count": entry.get("null_count", 0),
+        }
+        first_ts = _to_iso(entry.get("first_ts"))
+        last_ts = _to_iso(entry.get("last_ts"))
+        if first_ts:
+            item["first_observed"] = first_ts
+        if last_ts:
+            item["last_observed"] = last_ts
+        if kind == "list":
+            item["element_types"] = sorted(entry.get("element_types", []))
+            lengths = entry.get("lengths") or {}
+            item["lengths"] = {str(length): count for length, count in sorted(lengths.items())}
+            target = _resolve_cadence_target(entry, cadence_strategy)
+            if target is not None:
+                item["cadence"] = {"strategy": cadence_strategy, "target": target}
+            if "observed_elements" in entry:
+                item["observed_elements"] = int(entry.get("observed_elements", 0))
+        else:
+            item["value_types"] = sorted(entry.get("scalar_types", []))
+        meta_entries.append(item)
+    return meta_entries