mostlyright-data 0.9.0__py3-none-any.whl

mostlyright/data_harness/preparation/joins.py

@@ -0,0 +1,574 @@
+"""Stable equality and backward as-of left joins with exact audit evidence."""
+
+from __future__ import annotations
+
+from collections import Counter
+from dataclasses import dataclass, replace
+from datetime import datetime
+from typing import Any
+
+from mostlyright.data_harness.preparation.contracts import (
+    JOIN_AUDIT_VERSION,
+    MAX_ROWS,
+    AsOfJoinSpec,
+    ColumnLineage,
+    EqualityJoinSpec,
+    FieldSpec,
+    parse_utc,
+)
+from mostlyright.data_harness.preparation.errors import PreparationError
+from mostlyright.data_harness.preparation.table import TableState, validate_columns
+
+MAX_ASOF_COMPARISONS = 10_000_000
+
+
+@dataclass(frozen=True)
+class SelectedRevision:
+    left_ordinal: int
+    right_ordinal: int
+    event_time: str
+    available_at: str
+    revision_time: str
+
+    def __post_init__(self) -> None:
+        if type(self.left_ordinal) is not int or self.left_ordinal < 0:
+            raise PreparationError(
+                "JOIN_AUDIT_REVISION",
+                "selected_revision.left_ordinal",
+                "must be a nonnegative integer",
+            )
+        if type(self.right_ordinal) is not int or self.right_ordinal < 0:
+            raise PreparationError(
+                "JOIN_AUDIT_REVISION",
+                "selected_revision.right_ordinal",
+                "must be a nonnegative integer",
+            )
+        parse_utc(self.event_time, "selected_revision.event_time")
+        parse_utc(self.available_at, "selected_revision.available_at")
+        parse_utc(self.revision_time, "selected_revision.revision_time")
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "left_ordinal": self.left_ordinal,
+            "right_ordinal": self.right_ordinal,
+            "event_time": self.event_time,
+            "available_at": self.available_at,
+            "revision_time": self.revision_time,
+        }
+
+
+@dataclass(frozen=True)
+class JoinAudit:
+    schema_version: str
+    join_id: str
+    kind: str
+    cardinality: str
+    input_left_rows: int
+    input_right_rows: int
+    output_rows: int
+    matched_left_rows: int
+    unmatched_left_rows: int
+    matched_pairs: int
+    left_duplicate_keys: int
+    right_duplicate_keys: int
+    row_multiplier_numerator: int
+    row_multiplier_denominator: int
+    future_availability_rejections: int
+    future_revision_rejections: int
+    tolerance_rejections: int
+    lag_seconds_minimum: int | None
+    lag_seconds_maximum: int | None
+    selected_revisions: tuple[SelectedRevision, ...]
+
+    def __post_init__(self) -> None:
+        if self.schema_version != JOIN_AUDIT_VERSION:
+            raise PreparationError("JOIN_AUDIT_VERSION", "join_audit.schema_version", "unsupported")
+        count_fields = (
+            "input_left_rows",
+            "input_right_rows",
+            "output_rows",
+            "matched_left_rows",
+            "unmatched_left_rows",
+            "matched_pairs",
+            "left_duplicate_keys",
+            "right_duplicate_keys",
+            "row_multiplier_numerator",
+            "row_multiplier_denominator",
+            "future_availability_rejections",
+            "future_revision_rejections",
+            "tolerance_rejections",
+        )
+        if any(
+            type(getattr(self, name)) is not int or getattr(self, name) < 0 for name in count_fields
+        ):
+            raise PreparationError(
+                "JOIN_AUDIT_COUNT",
+                "join_audit",
+                "contains a negative or non-integer count",
+            )
+        if self.matched_left_rows + self.unmatched_left_rows != self.input_left_rows:
+            raise PreparationError(
+                "JOIN_AUDIT_COUNT",
+                "join_audit",
+                "matched and unmatched counts do not reconcile",
+            )
+        if self.row_multiplier_numerator != self.output_rows:
+            raise PreparationError(
+                "JOIN_AUDIT_COUNT",
+                "join_audit.row_multiplier_numerator",
+                "must equal output rows",
+            )
+        if self.row_multiplier_denominator != self.input_left_rows:
+            raise PreparationError(
+                "JOIN_AUDIT_COUNT",
+                "join_audit.row_multiplier_denominator",
+                "must equal input left rows",
+            )
+        if self.kind == "asof_left":
+            if self.matched_pairs != self.matched_left_rows:
+                raise PreparationError(
+                    "JOIN_AUDIT_COUNT",
+                    "join_audit.matched_pairs",
+                    "as-of joins select at most one right row per matched left row",
+                )
+            if len(self.selected_revisions) != self.matched_left_rows:
+                raise PreparationError(
+                    "JOIN_AUDIT_COUNT",
+                    "join_audit.selected_revisions",
+                    "does not reconcile matched left rows",
+                )
+        elif self.kind != "equality_left":
+            raise PreparationError("JOIN_AUDIT_KIND", "join_audit.kind", "is unsupported")
+        if (self.lag_seconds_minimum is None) != (self.lag_seconds_maximum is None):
+            raise PreparationError(
+                "JOIN_AUDIT_LAG",
+                "join_audit",
+                "lag bounds must be paired",
+            )
+        if (
+            self.lag_seconds_minimum is not None
+            and self.lag_seconds_maximum is not None
+            and (
+                self.lag_seconds_minimum < 0 or self.lag_seconds_maximum < self.lag_seconds_minimum
+            )
+        ):
+            raise PreparationError("JOIN_AUDIT_LAG", "join_audit", "lag bounds are invalid")
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "schema_version": self.schema_version,
+            "join_id": self.join_id,
+            "kind": self.kind,
+            "cardinality": self.cardinality,
+            "input_left_rows": self.input_left_rows,
+            "input_right_rows": self.input_right_rows,
+            "output_rows": self.output_rows,
+            "matched_left_rows": self.matched_left_rows,
+            "unmatched_left_rows": self.unmatched_left_rows,
+            "matched_pairs": self.matched_pairs,
+            "left_duplicate_keys": self.left_duplicate_keys,
+            "right_duplicate_keys": self.right_duplicate_keys,
+            "row_multiplier_numerator": self.row_multiplier_numerator,
+            "row_multiplier_denominator": self.row_multiplier_denominator,
+            "future_availability_rejections": self.future_availability_rejections,
+            "future_revision_rejections": self.future_revision_rejections,
+            "tolerance_rejections": self.tolerance_rejections,
+            "lag_seconds_minimum": self.lag_seconds_minimum,
+            "lag_seconds_maximum": self.lag_seconds_maximum,
+            "selected_revisions": [item.to_dict() for item in self.selected_revisions],
+        }
+
+
+def equality_left_join(
+    left: TableState,
+    right: TableState,
+    spec: EqualityJoinSpec,
+) -> tuple[TableState, JoinAudit]:
+    """Perform one declared stable equality left join."""
+
+    if right.table_id != spec.right_source_id:
+        raise PreparationError("JOIN_SOURCE", spec.join_id, "right source does not match join")
+    left_indexes = validate_columns(left, spec.left_keys, f"{spec.join_id}.left_keys")
+    right_indexes = validate_columns(right, spec.right_keys, f"{spec.join_id}.right_keys")
+    _validate_key_types(left, right, left_indexes, right_indexes, spec.join_id)
+    left_keys = _keys(left, left_indexes, spec.join_id, "left")
+    right_keys = _keys(right, right_indexes, spec.join_id, "right")
+    left_duplicates = _duplicate_count(left_keys)
+    right_duplicates = _duplicate_count(right_keys)
+    if spec.cardinality in {"one_to_one", "one_to_many"} and left_duplicates:
+        raise PreparationError(
+            "JOIN_CARDINALITY",
+            spec.join_id,
+            "declared cardinality requires unique left keys",
+            severity="critical",
+        )
+    if spec.cardinality in {"one_to_one", "many_to_one"} and right_duplicates:
+        raise PreparationError(
+            "JOIN_CARDINALITY",
+            spec.join_id,
+            "declared cardinality requires unique right keys",
+            severity="critical",
+        )
+    right_lookup: dict[tuple[Any, ...], list[int]] = {}
+    for ordinal, key in enumerate(right_keys):
+        right_lookup.setdefault(key, []).append(ordinal)
+    output_fields, output_lineage, right_output_indexes = _joined_schema(
+        left,
+        right,
+        skipped_right=frozenset(spec.right_keys),
+        suffix=spec.suffix,
+        join_id=spec.join_id,
+    )
+    output_rows: list[tuple[Any, ...]] = []
+    matched_left = 0
+    matched_pairs = 0
+    for left_row, key in zip(left.rows, left_keys, strict=True):
+        matches = right_lookup.get(key, ())
+        if matches:
+            matched_left += 1
+            matched_pairs += len(matches)
+            if len(output_rows) + len(matches) > MAX_ROWS:
+                raise PreparationError(
+                    "JOIN_ROW_LIMIT",
+                    spec.join_id,
+                    "join output exceeds the deterministic row budget",
+                    severity="critical",
+                )
+            for right_ordinal in matches:
+                output_rows.append(
+                    (
+                        *left_row,
+                        *(right.rows[right_ordinal][index] for index in right_output_indexes),
+                    )
+                )
+        else:
+            output_rows.append((*left_row, *(None for _ in right_output_indexes)))
+    output = TableState(
+        table_id=f"{left.table_id}.{spec.join_id}",
+        snapshot_digests=tuple(sorted(set((*left.snapshot_digests, *right.snapshot_digests)))),
+        fields=output_fields,
+        rows=tuple(output_rows),
+        lineage=output_lineage,
+    )
+    audit = JoinAudit(
+        schema_version=JOIN_AUDIT_VERSION,
+        join_id=spec.join_id,
+        kind="equality_left",
+        cardinality=spec.cardinality,
+        input_left_rows=len(left.rows),
+        input_right_rows=len(right.rows),
+        output_rows=len(output.rows),
+        matched_left_rows=matched_left,
+        unmatched_left_rows=len(left.rows) - matched_left,
+        matched_pairs=matched_pairs,
+        left_duplicate_keys=left_duplicates,
+        right_duplicate_keys=right_duplicates,
+        row_multiplier_numerator=len(output.rows),
+        row_multiplier_denominator=len(left.rows),
+        future_availability_rejections=0,
+        future_revision_rejections=0,
+        tolerance_rejections=0,
+        lag_seconds_minimum=None,
+        lag_seconds_maximum=None,
+        selected_revisions=(),
+    )
+    return output, audit
+
+
+def backward_asof_join(
+    left: TableState,
+    right: TableState,
+    spec: AsOfJoinSpec,
+) -> tuple[TableState, JoinAudit]:
+    """Perform a partitioned cutoff-safe backward as-of join."""
+
+    if right.table_id != spec.right_source_id:
+        raise PreparationError("JOIN_SOURCE", spec.join_id, "right source does not match join")
+    left_by = validate_columns(left, spec.left_by, f"{spec.join_id}.left_by")
+    right_by = validate_columns(right, spec.right_by, f"{spec.join_id}.right_by")
+    _validate_key_types(left, right, left_by, right_by, spec.join_id)
+    left_keys = _keys(left, left_by, spec.join_id, "left")
+    right_keys = _keys(right, right_by, spec.join_id, "right")
+    left_duplicates = _duplicate_count(
+        tuple(
+            (*key, left.rows[index][left.index(spec.left_time)])
+            for index, key in enumerate(left_keys)
+        )
+    )
+    if spec.cardinality == "one_to_one" and left_duplicates:
+        raise PreparationError(
+            "JOIN_CARDINALITY",
+            spec.join_id,
+            "one_to_one as-of join requires unique left partition/time keys",
+            severity="critical",
+        )
+    right_time_index = right.index(spec.right_time)
+    available_index = right.index(spec.right_available_at)
+    revision_index = right.index(spec.right_revision_time)
+    left_time_index = left.index(spec.left_time)
+    left_cutoff_index = left.index(spec.left_cutoff)
+    for state, indexes, label in (
+        (left, (left_time_index, left_cutoff_index), "left"),
+        (right, (right_time_index, available_index, revision_index), "right"),
+    ):
+        for index in indexes:
+            if state.fields[index].logical_type != "timestamp_utc":
+                raise PreparationError(
+                    "JOIN_TIME_TYPE",
+                    spec.join_id,
+                    f"{label} time fields must be UTC timestamps",
+                )
+    right_identity = tuple(
+        (*right_keys[index], row[right_time_index], row[revision_index])
+        for index, row in enumerate(right.rows)
+    )
+    right_duplicates = _duplicate_count(right_identity)
+    if right_duplicates:
+        raise PreparationError(
+            "JOIN_UNSTABLE_TIE",
+            spec.join_id,
+            "right partition/event/revision identity is not unique",
+            severity="critical",
+        )
+    output_fields, output_lineage, right_output_indexes = _joined_schema(
+        left,
+        right,
+        skipped_right=frozenset(spec.right_by),
+        suffix=spec.suffix,
+        join_id=spec.join_id,
+    )
+    right_by_partition: dict[tuple[Any, ...], list[int]] = {}
+    for ordinal, key in enumerate(right_keys):
+        right_by_partition.setdefault(key, []).append(ordinal)
+    question_cutoff = parse_utc(spec.question_cutoff, f"{spec.join_id}.question_cutoff")
+    output_rows: list[tuple[Any, ...]] = []
+    selected: list[SelectedRevision] = []
+    lags: list[int] = []
+    matched_left = 0
+    future_availability = 0
+    future_revision = 0
+    tolerance_rejections = 0
+    comparisons = 0
+    for left_ordinal, (left_row, partition) in enumerate(zip(left.rows, left_keys, strict=True)):
+        left_time = _instant(left_row[left_time_index], spec.join_id)
+        row_cutoff = _instant(left_row[left_cutoff_index], spec.join_id)
+        if row_cutoff > question_cutoff:
+            raise PreparationError(
+                "TEMPORAL_CUTOFF",
+                f"{spec.join_id}.left[{left_ordinal}]",
+                "row cutoff exceeds the question cutoff",
+            )
+        eligible: list[tuple[datetime, datetime, int, datetime]] = []
+        for right_ordinal in right_by_partition.get(partition, ()):
+            comparisons += 1
+            if comparisons > MAX_ASOF_COMPARISONS:
+                raise PreparationError(
+                    "JOIN_COMPARISON_LIMIT",
+                    spec.join_id,
+                    "as-of candidate comparisons exceed the deterministic budget",
+                    severity="critical",
+                )
+            right_row = right.rows[right_ordinal]
+            event_time = _instant(right_row[right_time_index], spec.join_id)
+            if event_time > left_time:
+                continue
+            lag = int((left_time - event_time).total_seconds())
+            if lag > spec.tolerance_seconds:
+                tolerance_rejections += 1
+                continue
+            available_at = _instant(right_row[available_index], spec.join_id)
+            if available_at > row_cutoff or available_at > question_cutoff:
+                future_availability += 1
+                continue
+            revision_time = _instant(right_row[revision_index], spec.join_id)
+            if revision_time > row_cutoff or revision_time > question_cutoff:
+                future_revision += 1
+                continue
+            eligible.append((event_time, revision_time, right_ordinal, available_at))
+        chosen: tuple[datetime, datetime, int, datetime] | None = None
+        if eligible:
+            latest_event = max(item[0] for item in eligible)
+            revisions = [item for item in eligible if item[0] == latest_event]
+            if spec.revision_policy == "latest_available":
+                chosen = max(revisions, key=lambda item: (item[1], -item[2]))
+            else:
+                chosen = min(revisions, key=lambda item: (item[1], item[2]))
+        if chosen is None:
+            output_rows.append((*left_row, *(None for _ in right_output_indexes)))
+            continue
+        event_time, revision_time, right_ordinal, available_at = chosen
+        matched_left += 1
+        right_row = right.rows[right_ordinal]
+        output_rows.append((*left_row, *(right_row[index] for index in right_output_indexes)))
+        lag = int((left_time - event_time).total_seconds())
+        lags.append(lag)
+        selected.append(
+            SelectedRevision(
+                left_ordinal=left_ordinal,
+                right_ordinal=right_ordinal,
+                event_time=_utc_text(event_time),
+                available_at=_utc_text(available_at),
+                revision_time=_utc_text(revision_time),
+            )
+        )
+    output = TableState(
+        table_id=f"{left.table_id}.{spec.join_id}",
+        snapshot_digests=tuple(sorted(set((*left.snapshot_digests, *right.snapshot_digests)))),
+        fields=output_fields,
+        rows=tuple(output_rows),
+        lineage=output_lineage,
+    )
+    audit = JoinAudit(
+        schema_version=JOIN_AUDIT_VERSION,
+        join_id=spec.join_id,
+        kind="asof_left",
+        cardinality=spec.cardinality,
+        input_left_rows=len(left.rows),
+        input_right_rows=len(right.rows),
+        output_rows=len(output.rows),
+        matched_left_rows=matched_left,
+        unmatched_left_rows=len(left.rows) - matched_left,
+        matched_pairs=matched_left,
+        left_duplicate_keys=left_duplicates,
+        right_duplicate_keys=right_duplicates,
+        row_multiplier_numerator=len(output.rows),
+        row_multiplier_denominator=len(left.rows),
+        future_availability_rejections=future_availability,
+        future_revision_rejections=future_revision,
+        tolerance_rejections=tolerance_rejections,
+        lag_seconds_minimum=min(lags) if lags else None,
+        lag_seconds_maximum=max(lags) if lags else None,
+        selected_revisions=tuple(selected),
+    )
+    return output, audit
+
+
+def _joined_schema(
+    left: TableState,
+    right: TableState,
+    *,
+    skipped_right: frozenset[str],
+    suffix: str,
+    join_id: str,
+) -> tuple[tuple[FieldSpec, ...], tuple[ColumnLineage, ...], tuple[int, ...]]:
+    output_names = set(left.columns)
+    right_indexes: list[int] = []
+    right_names: dict[str, str] = {}
+    for index, field in enumerate(right.fields):
+        if field.name in skipped_right:
+            continue
+        output_name = field.name if field.name not in output_names else f"{field.name}_{suffix}"
+        if output_name in output_names:
+            raise PreparationError("JOIN_COLLISION", join_id, "suffix does not resolve collision")
+        output_names.add(output_name)
+        right_indexes.append(index)
+        right_names[field.name] = output_name
+    right_fields = tuple(
+        replace(
+            right.fields[index],
+            name=right_names[right.fields[index].name],
+            nullable=True,
+            availability_column=(
+                right_names.get(
+                    right.fields[index].availability_column,
+                    right.fields[index].availability_column,
+                )
+                if right.fields[index].availability_column is not None
+                else None
+            ),
+        )
+        for index in right_indexes
+    )
+    right_lineage_by_name = {item.output_column: item for item in right.lineage}
+    right_lineage = tuple(
+        ColumnLineage(
+            output_column=right_names[right.fields[index].name],
+            inputs=right_lineage_by_name[right.fields[index].name].inputs,
+            operation_ids=(
+                *right_lineage_by_name[right.fields[index].name].operation_ids,
+                join_id,
+            ),
+        )
+        for index in right_indexes
+    )
+    left_lineage = tuple(
+        ColumnLineage(
+            output_column=item.output_column,
+            inputs=item.inputs,
+            operation_ids=(*item.operation_ids, join_id),
+        )
+        for item in left.lineage
+    )
+    return (*left.fields, *right_fields), (*left_lineage, *right_lineage), tuple(right_indexes)
+
+
+def _validate_key_types(
+    left: TableState,
+    right: TableState,
+    left_indexes: tuple[int, ...],
+    right_indexes: tuple[int, ...],
+    join_id: str,
+) -> None:
+    for left_index, right_index in zip(left_indexes, right_indexes, strict=True):
+        left_field = left.fields[left_index]
+        right_field = right.fields[right_index]
+        left_type = (
+            left_field.logical_type,
+            left_field.physical_type,
+            left_field.semantic_type,
+            left_field.unit,
+            left_field.entity,
+            left_field.temporal_role,
+        )
+        right_type = (
+            right_field.logical_type,
+            right_field.physical_type,
+            right_field.semantic_type,
+            right_field.unit,
+            right_field.entity,
+            right_field.temporal_role,
+        )
+        if left_type != right_type:
+            raise PreparationError("JOIN_KEY_TYPE", join_id, "join key types are not exact")
+
+
+def _keys(
+    state: TableState,
+    indexes: tuple[int, ...],
+    join_id: str,
+    side: str,
+) -> tuple[tuple[Any, ...], ...]:
+    result: list[tuple[Any, ...]] = []
+    for row_index, row in enumerate(state.rows):
+        key = tuple(row[index] for index in indexes)
+        if any(value is None for value in key):
+            raise PreparationError(
+                "JOIN_NULL_KEY",
+                f"{join_id}.{side}[{row_index}]",
+                "join keys cannot be null",
+                severity="critical",
+            )
+        result.append(key)
+    return tuple(result)
+
+
+def _duplicate_count(keys: tuple[tuple[Any, ...], ...]) -> int:
+    return sum(count - 1 for count in Counter(keys).values())
+
+
+def _instant(value: Any, join_id: str) -> datetime:
+    if not isinstance(value, str):
+        raise PreparationError("JOIN_TIME_TYPE", join_id, "join time value is not text")
+    return parse_utc(value, join_id)
+
+
+def _utc_text(value: datetime) -> str:
+    return value.isoformat(timespec="microseconds").replace("+00:00", "Z").replace(".000000Z", "Z")
+
+
+__all__ = [
+    "JoinAudit",
+    "SelectedRevision",
+    "backward_asof_join",
+    "equality_left_join",
+]