mostlyright-data 0.9.0__py3-none-any.whl

mostlyright/data_harness/preparation/transforms.py

@@ -0,0 +1,568 @@
+"""Closed deterministic cleaning, normalization, resolution, and projection operations."""
+
+from __future__ import annotations
+
+import math
+import re
+from dataclasses import dataclass, replace
+from datetime import date
+from decimal import Decimal, InvalidOperation
+from fractions import Fraction
+from functools import cmp_to_key
+from typing import Any
+
+from mostlyright.data_harness.preparation.contracts import (
+    OPERATION_AUDIT_VERSION,
+    PHYSICAL_TYPES,
+    ColumnLineage,
+    OperationSpec,
+    parse_utc,
+)
+from mostlyright.data_harness.preparation.errors import PreparationError
+from mostlyright.data_harness.preparation.profile import TableProfile, profile_table
+from mostlyright.data_harness.preparation.table import TableState, validate_columns
+
+_ASCII_TRIM = "\t\n\v\f\r "
+_DIGEST = re.compile(r"\A[0-9a-f]{64}\Z")
+_INT64 = re.compile(r"\A-?[0-9]+\Z")
+_FLOAT64 = re.compile(r"\A-?(?:0|[1-9][0-9]*)(?:\.[0-9]+)?(?:[eE]-?[0-9]+)?\Z")
+_UNIT_CONVERSIONS: dict[tuple[str, str], tuple[int, int, int, int]] = {
+    ("celsius", "fahrenheit"): (9, 5, 32, 1),
+    ("fahrenheit", "celsius"): (5, 9, -160, 9),
+    ("celsius", "kelvin"): (1, 1, 27315, 100),
+    ("kelvin", "celsius"): (1, 1, -27315, 100),
+    ("meter", "kilometer"): (1, 1000, 0, 1),
+    ("kilometer", "meter"): (1000, 1, 0, 1),
+    ("mile", "kilometer"): (25146, 15625, 0, 1),
+    ("kilometer", "mile"): (15625, 25146, 0, 1),
+    ("second", "minute"): (1, 60, 0, 1),
+    ("minute", "second"): (60, 1, 0, 1),
+    ("minute", "hour"): (1, 60, 0, 1),
+    ("hour", "minute"): (60, 1, 0, 1),
+    ("ratio", "percent"): (100, 1, 0, 1),
+    ("percent", "ratio"): (1, 100, 0, 1),
+}
+
+
+@dataclass(frozen=True)
+class OperationAudit:
+    schema_version: str
+    operation_id: str
+    kind: str
+    parameters_digest: str
+    input_profile_digest: str
+    output_profile_digest: str
+    input_rows: int
+    output_rows: int
+    changed_cells: int
+    removed_rows: int
+
+    def __post_init__(self) -> None:
+        if self.schema_version != OPERATION_AUDIT_VERSION:
+            raise PreparationError("OPERATION_AUDIT_VERSION", "audit.schema_version", "unsupported")
+        for name in (
+            "parameters_digest",
+            "input_profile_digest",
+            "output_profile_digest",
+        ):
+            value = getattr(self, name)
+            if not isinstance(value, str) or _DIGEST.fullmatch(value) is None:
+                raise PreparationError(
+                    "OPERATION_AUDIT_DIGEST",
+                    f"audit.{name}",
+                    "must be a lowercase SHA-256 digest",
+                )
+        for name in (
+            "input_rows",
+            "output_rows",
+            "changed_cells",
+            "removed_rows",
+        ):
+            value = getattr(self, name)
+            if type(value) is not int or value < 0:
+                raise PreparationError(
+                    "OPERATION_AUDIT_COUNT",
+                    f"audit.{name}",
+                    "must be nonnegative",
+                )
+        if self.removed_rows != self.input_rows - self.output_rows:
+            raise PreparationError(
+                "OPERATION_AUDIT_COUNT",
+                "audit.removed_rows",
+                "does not reconcile input and output rows",
+            )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "schema_version": self.schema_version,
+            "operation_id": self.operation_id,
+            "kind": self.kind,
+            "parameters_digest": self.parameters_digest,
+            "input_profile_digest": self.input_profile_digest,
+            "output_profile_digest": self.output_profile_digest,
+            "input_rows": self.input_rows,
+            "output_rows": self.output_rows,
+            "changed_cells": self.changed_cells,
+            "removed_rows": self.removed_rows,
+        }
+
+
+@dataclass(frozen=True)
+class PreparedSource:
+    state: TableState
+    before_profile: TableProfile
+    after_profile: TableProfile
+    operation_audits: tuple[OperationAudit, ...]
+
+
+def apply_operations(
+    state: TableState,
+    operations: tuple[OperationSpec, ...],
+) -> PreparedSource:
+    """Apply a fixed ordered operation list and bind every before/after profile."""
+
+    if not isinstance(operations, tuple):
+        raise PreparationError("OPERATION_COLLECTION", "operations", "must be an immutable tuple")
+    before = profile_table(state)
+    audits: list[OperationAudit] = []
+    current = state
+    for operation in operations:
+        if operation.source_id != state.table_id:
+            raise PreparationError(
+                "OPERATION_SOURCE",
+                operation.operation_id,
+                "operation source does not match the prepared source",
+            )
+        input_profile = profile_table(current)
+        output, changed = _apply(current, operation)
+        output_profile = profile_table(output)
+        audits.append(
+            OperationAudit(
+                schema_version=OPERATION_AUDIT_VERSION,
+                operation_id=operation.operation_id,
+                kind=operation.kind,
+                parameters_digest=operation.parameters_digest,
+                input_profile_digest=input_profile.digest,
+                output_profile_digest=output_profile.digest,
+                input_rows=len(current.rows),
+                output_rows=len(output.rows),
+                changed_cells=changed,
+                removed_rows=len(current.rows) - len(output.rows),
+            )
+        )
+        current = output
+    return PreparedSource(
+        state=current,
+        before_profile=before,
+        after_profile=profile_table(current),
+        operation_audits=tuple(audits),
+    )
+
+
+def _apply(state: TableState, operation: OperationSpec) -> tuple[TableState, int]:
+    dispatch = {
+        "trim": _trim,
+        "case_normalize": _case_normalize,
+        "null_tokens": _null_tokens,
+        "strict_cast": _strict_cast,
+        "unit_convert": _unit_convert,
+        "entity_resolve": _entity_resolve,
+        "rename": _rename,
+        "project": _project,
+        "deduplicate": _deduplicate,
+    }
+    return dispatch[operation.kind](state, operation)
+
+
+def _map_column(
+    state: TableState,
+    operation: OperationSpec,
+    transform: Any,
+    *,
+    field_transform: Any | None = None,
+) -> tuple[TableState, int]:
+    column = operation.parameter("column")
+    index = state.index(column)
+    changed = 0
+    rows: list[tuple[Any, ...]] = []
+    for row_index, row in enumerate(state.rows):
+        value = transform(row[index], row_index)
+        if _typed_key(value) != _typed_key(row[index]):
+            changed += 1
+        mutable = list(row)
+        mutable[index] = value
+        rows.append(tuple(mutable))
+    fields = state.fields
+    if field_transform is not None:
+        mutable_fields = list(fields)
+        mutable_fields[index] = field_transform(fields[index])
+        fields = tuple(mutable_fields)
+    lineage = _append_operation(state.lineage, (column,), operation.operation_id)
+    return state.replace(fields=fields, rows=tuple(rows), lineage=lineage), changed
+
+
+def _trim(state: TableState, operation: OperationSpec) -> tuple[TableState, int]:
+    field = state.field(operation.parameter("column"))
+    if field.logical_type != "string":
+        raise PreparationError("OPERATION_TYPE", operation.operation_id, "trim requires string")
+    return _map_column(
+        state,
+        operation,
+        lambda value, _row: value.strip(_ASCII_TRIM) if value is not None else None,
+    )
+
+
+def _case_normalize(state: TableState, operation: OperationSpec) -> tuple[TableState, int]:
+    field = state.field(operation.parameter("column"))
+    if field.logical_type != "string":
+        raise PreparationError(
+            "OPERATION_TYPE",
+            operation.operation_id,
+            "case normalization requires string",
+        )
+    mode = operation.parameter("mode")
+
+    def normalize(value: Any, _row: int) -> Any:
+        if value is None:
+            return None
+        if mode == "lower":
+            return value.lower()
+        if mode == "upper":
+            return value.upper()
+        return value.casefold()
+
+    return _map_column(state, operation, normalize)
+
+
+def _null_tokens(state: TableState, operation: OperationSpec) -> tuple[TableState, int]:
+    field = state.field(operation.parameter("column"))
+    if field.logical_type != "string":
+        raise PreparationError(
+            "OPERATION_TYPE",
+            operation.operation_id,
+            "null tokens require string",
+        )
+    tokens = frozenset(operation.parameter("tokens"))
+    return _map_column(
+        state,
+        operation,
+        lambda value, _row: None if value in tokens else value,
+    )
+
+
+def _strict_cast(state: TableState, operation: OperationSpec) -> tuple[TableState, int]:
+    column = operation.parameter("column")
+    target = operation.parameter("to_type")
+    field = state.field(column)
+    if target == field.logical_type:
+        raise PreparationError("OPERATION_CAST", operation.operation_id, "cast must change type")
+
+    def cast(value: Any, row_index: int) -> Any:
+        if value is None:
+            return None
+        path = f"{operation.operation_id}.row[{row_index}]"
+        if target == "string":
+            if type(value) is bool:
+                return "true" if value else "false"
+            if type(value) is int:
+                return str(value)
+            if type(value) is float and math.isfinite(value):
+                return format(value, ".17g")
+            if isinstance(value, str):
+                return value
+        if not isinstance(value, str):
+            raise PreparationError("CAST_VALUE", path, "cast input must be a canonical string")
+        if target == "int64":
+            if _INT64.fullmatch(value) is None:
+                raise PreparationError("CAST_INT64", path, "value is not strict int64 syntax")
+            parsed = int(value)
+            if not -(1 << 63) <= parsed <= (1 << 63) - 1:
+                raise PreparationError("CAST_INT64", path, "value is outside int64 range")
+            return parsed
+        if target == "float64":
+            if _FLOAT64.fullmatch(value) is None:
+                raise PreparationError(
+                    "CAST_FLOAT64",
+                    path,
+                    "value is not strict finite float syntax",
+                )
+            parsed = float(value)
+            if not math.isfinite(parsed):
+                raise PreparationError("CAST_FLOAT64", path, "value is not finite")
+            return parsed
+        if target == "boolean":
+            if value not in {"false", "true"}:
+                raise PreparationError("CAST_BOOLEAN", path, "value must be true or false")
+            return value == "true"
+        if target == "date":
+            try:
+                parsed_date = date.fromisoformat(value)
+            except ValueError:
+                raise PreparationError("CAST_DATE", path, "value is not a real ISO date") from None
+            if parsed_date.isoformat() != value or len(value) != 10:
+                raise PreparationError("CAST_DATE", path, "value is not canonical ISO date")
+            return value
+        if target == "timestamp_utc":
+            parse_utc(value, path)
+            return value
+        raise PreparationError("CAST_VALUE", path, "cast conversion is unsupported")
+
+    return _map_column(
+        state,
+        operation,
+        cast,
+        field_transform=lambda current: replace(
+            current,
+            logical_type=target,
+            physical_type=PHYSICAL_TYPES[target],
+        ),
+    )
+
+
+def _unit_convert(state: TableState, operation: OperationSpec) -> tuple[TableState, int]:
+    column = operation.parameter("column")
+    source = operation.parameter("from_unit")
+    target = operation.parameter("to_unit")
+    field = state.field(column)
+    if field.logical_type not in {"int64", "float64"} or field.semantic_type != "measure":
+        raise PreparationError(
+            "OPERATION_TYPE",
+            operation.operation_id,
+            "unit conversion requires a numeric measure",
+        )
+    if field.unit != source:
+        raise PreparationError(
+            "OPERATION_UNIT",
+            operation.operation_id,
+            "declared source unit differs",
+        )
+    try:
+        scale_numerator, scale_denominator, offset_numerator, offset_denominator = (
+            _UNIT_CONVERSIONS[(source, target)]
+        )
+    except KeyError:
+        raise PreparationError(
+            "OPERATION_UNIT",
+            operation.operation_id,
+            "conversion is not in the closed rational registry",
+        ) from None
+
+    def convert(value: Any, row_index: int) -> Any:
+        if value is None:
+            return None
+        try:
+            converted = Fraction(Decimal(str(value))) * Fraction(
+                scale_numerator, scale_denominator
+            ) + Fraction(offset_numerator, offset_denominator)
+        except (InvalidOperation, ValueError):
+            raise PreparationError(
+                "OPERATION_UNIT",
+                f"{operation.operation_id}.row[{row_index}]",
+                "numeric value cannot be converted",
+            ) from None
+        result = float(converted)
+        if not math.isfinite(result):
+            raise PreparationError(
+                "OPERATION_UNIT",
+                f"{operation.operation_id}.row[{row_index}]",
+                "conversion produced a non-finite value",
+            )
+        return result
+
+    return _map_column(
+        state,
+        operation,
+        convert,
+        field_transform=lambda current: replace(
+            current,
+            logical_type="float64",
+            physical_type=PHYSICAL_TYPES["float64"],
+            unit=target,
+        ),
+    )
+
+
+def _entity_resolve(state: TableState, operation: OperationSpec) -> tuple[TableState, int]:
+    column = operation.parameter("column")
+    field = state.field(column)
+    if field.logical_type != "string" or field.semantic_type != "entity":
+        raise PreparationError(
+            "OPERATION_TYPE",
+            operation.operation_id,
+            "entity resolution requires a string entity field",
+        )
+    if field.entity != operation.parameter("entity_type"):
+        raise PreparationError(
+            "ENTITY_TYPE",
+            operation.operation_id,
+            "entity type differs from the field declaration",
+        )
+    aliases = dict(operation.parameter("aliases"))
+    canonical = frozenset(aliases.values())
+
+    def resolve(value: Any, row_index: int) -> Any:
+        if value is None or value in canonical:
+            return value
+        try:
+            return aliases[value]
+        except KeyError:
+            raise PreparationError(
+                "ENTITY_UNRESOLVED",
+                f"{operation.operation_id}.row[{row_index}]",
+                "entity value is neither canonical nor an allowlisted alias",
+            ) from None
+
+    return _map_column(state, operation, resolve)
+
+
+def _rename(state: TableState, operation: OperationSpec) -> tuple[TableState, int]:
+    mappings = dict(operation.parameter("mappings"))
+    unknown = set(mappings) - set(state.columns)
+    if unknown:
+        raise PreparationError(
+            "OPERATION_RENAME",
+            operation.operation_id,
+            "references unknown columns",
+        )
+    output_names = tuple(mappings.get(name, name) for name in state.columns)
+    if len(set(output_names)) != len(output_names):
+        raise PreparationError(
+            "OPERATION_RENAME",
+            operation.operation_id,
+            "creates a name collision",
+        )
+    fields = tuple(
+        replace(
+            field,
+            name=mappings.get(field.name, field.name),
+            availability_column=(
+                mappings.get(field.availability_column, field.availability_column)
+                if field.availability_column is not None
+                else None
+            ),
+        )
+        for field in state.fields
+    )
+    lineage = tuple(
+        ColumnLineage(
+            output_column=mappings.get(item.output_column, item.output_column),
+            inputs=item.inputs,
+            operation_ids=(*item.operation_ids, operation.operation_id),
+        )
+        for item in state.lineage
+    )
+    return state.replace(fields=fields, lineage=lineage), len(state.rows) * len(mappings)
+
+
+def _project(state: TableState, operation: OperationSpec) -> tuple[TableState, int]:
+    columns = operation.parameter("columns")
+    indexes = validate_columns(state, columns, operation.operation_id)
+    fields = tuple(state.fields[index] for index in indexes)
+    output_names = frozenset(columns)
+    for field in fields:
+        if field.availability_column is not None and field.availability_column not in output_names:
+            raise PreparationError(
+                "TEMPORAL_SEMANTICS",
+                operation.operation_id,
+                f"projection drops availability column for {field.name}",
+            )
+    rows = tuple(tuple(row[index] for index in indexes) for row in state.rows)
+    lineage_by_name = {item.output_column: item for item in state.lineage}
+    lineage = tuple(
+        ColumnLineage(
+            output_column=column,
+            inputs=lineage_by_name[column].inputs,
+            operation_ids=(*lineage_by_name[column].operation_ids, operation.operation_id),
+        )
+        for column in columns
+    )
+    removed_columns = len(state.columns) - len(columns)
+    return state.replace(fields=fields, rows=rows, lineage=lineage), len(rows) * removed_columns
+
+
+def _deduplicate(state: TableState, operation: OperationSpec) -> tuple[TableState, int]:
+    key_indexes = validate_columns(state, operation.parameter("keys"), operation.operation_id)
+    authority = operation.parameter("authority")
+    authority_indexes = tuple((state.index(column), direction) for column, direction in authority)
+    groups: dict[tuple[Any, ...], list[tuple[int, tuple[Any, ...]]]] = {}
+    for ordinal, row in enumerate(state.rows):
+        key = tuple(row[index] for index in key_indexes)
+        if any(value is None for value in key):
+            raise PreparationError(
+                "DEDUPLICATE_NULL_KEY",
+                f"{operation.operation_id}.row[{ordinal}]",
+                "deduplication keys cannot be null",
+            )
+        if any(row[index] is None for index, _ in authority_indexes):
+            raise PreparationError(
+                "DEDUPLICATE_NULL_AUTHORITY",
+                f"{operation.operation_id}.row[{ordinal}]",
+                "authority values cannot be null",
+            )
+        groups.setdefault(key, []).append((ordinal, row))
+    selected: list[tuple[int, tuple[Any, ...]]] = []
+    for candidates in groups.values():
+        candidates.sort(
+            key=cmp_to_key(
+                lambda left, right: _compare_authority(
+                    left,
+                    right,
+                    authority_indexes,
+                )
+            )
+        )
+        selected.append(candidates[0])
+    selected.sort(key=lambda item: item[0])
+    lineage = _append_operation(state.lineage, state.columns, operation.operation_id)
+    rows = tuple(row for _, row in selected)
+    return state.replace(rows=rows, lineage=lineage), len(state.rows) - len(rows)
+
+
+def _compare_authority(
+    left: tuple[int, tuple[Any, ...]],
+    right: tuple[int, tuple[Any, ...]],
+    authority: tuple[tuple[int, str], ...],
+) -> int:
+    for index, direction in authority:
+        left_key = _sortable(left[1][index])
+        right_key = _sortable(right[1][index])
+        if left_key < right_key:
+            return -1 if direction == "ascending" else 1
+        if left_key > right_key:
+            return 1 if direction == "ascending" else -1
+    return (left[0] > right[0]) - (left[0] < right[0])
+
+
+def _sortable(value: Any) -> tuple[int, Any]:
+    if type(value) is bool:
+        return (0, value)
+    if type(value) in {int, float}:
+        return (1, Decimal(str(value)))
+    return (2, value)
+
+
+def _append_operation(
+    lineage: tuple[ColumnLineage, ...],
+    columns: tuple[str, ...],
+    operation_id: str,
+) -> tuple[ColumnLineage, ...]:
+    selected = frozenset(columns)
+    return tuple(
+        ColumnLineage(
+            output_column=item.output_column,
+            inputs=item.inputs,
+            operation_ids=(
+                (*item.operation_ids, operation_id)
+                if item.output_column in selected
+                else item.operation_ids
+            ),
+        )
+        for item in lineage
+    )
+
+
+def _typed_key(value: Any) -> tuple[type[Any], Any]:
+    return (type(value), value)
+
+
+__all__ = ["OperationAudit", "PreparedSource", "apply_operations"]