DataOps-platform/app/runner/rule_polars.py

"""Runner adapter that reconstructs allowlisted Polars LazyFrame operations."""

from __future__ import annotations

from contextlib import suppress
from decimal import Decimal
from typing import Any

import polars as pl

from app.core.common.identifiers import ensure_governance_uid
from app.core.data_rules.compilers.polars import (
    bound_polars_plan_hash,
    validate_bound_polars_plan,
)
from app.runner.nodes import NodeExecutionError

_TYPE_MAP = {
    "boolean": pl.Boolean,
    "date": pl.Date,
    "double": pl.Float64,
    "float": pl.Float32,
    "integer": pl.Int64,
    "string": pl.String,
    "timestamp": pl.Datetime,
}
_IDEMPOTENCY = {
    "deduplication_key",
    "partition_replace",
    "upsert",
}


def _uid(value: Any, label: str) -> str:
    try:
        return ensure_governance_uid({"uid": str(value)})
    except ValueError as exc:
        raise NodeExecutionError(f"{label} is invalid") from exc


class _ExpressionCompiler:
    def compile(self, ast: dict[str, Any]) -> pl.Expr:
        kind = ast["kind"]
        if kind == "identifier":
            return pl.col(ast["name"])
        if kind == "literal":
            value = ast["value"]
            if ast["type"] == "decimal":
                value = Decimal(value)
            return pl.lit(value)
        if kind == "unary":
            operand = self.compile(ast["operand"])
            return ~operand if ast["operator"] == "!" else -operand
        if kind == "binary":
            left = self.compile(ast["left"])
            right = self.compile(ast["right"])
            return {
                "||": lambda: left | right,
                "&&": lambda: left & right,
                "==": lambda: left == right,
                "!=": lambda: left != right,
                "<": lambda: left < right,
                "<=": lambda: left <= right,
                ">": lambda: left > right,
                ">=": lambda: left >= right,
                "+": lambda: left + right,
                "-": lambda: left - right,
                "*": lambda: left * right,
                "/": lambda: left / right,
                "%": lambda: left % right,
            }[ast["operator"]]()
        function = ast["function"]
        arguments = [self.compile(item) for item in ast["arguments"]]
        if function == "matches":
            pattern = ast["arguments"][1]["value"]
            return arguments[0].str.contains(pattern, strict=True)
        if function == "lower":
            return arguments[0].str.to_lowercase()
        if function == "upper":
            return arguments[0].str.to_uppercase()
        if function == "trim":
            return arguments[0].str.strip_chars()
        if function == "length":
            return arguments[0].str.len_chars()
        if function == "coalesce":
            return pl.coalesce(arguments)
        if function == "date":
            return arguments[0].cast(pl.Date, strict=True)
        if function == "timestamp":
            return arguments[0].cast(pl.Datetime, strict=True)
        if function == "abs":
            return arguments[0].abs()
        raise NodeExecutionError("published Polars expression is unsupported")


def _cast_type(field: dict[str, Any]) -> pl.DataType:
    name = field["type"]
    if name == "decimal":
        precision = field.get("precision")
        scale = field.get("scale")
        if precision is None or scale is None:
            raise NodeExecutionError(
                "published decimal cast contract is incomplete"
            )
        return pl.Decimal(precision=precision, scale=scale)
    if name == "timestamptz":
        timezone = field.get("timezone")
        if not timezone:
            raise NodeExecutionError(
                "published timestamptz cast contract is incomplete"
            )
        return pl.Datetime(time_zone=timezone)
    dtype = _TYPE_MAP.get(name)
    if dtype is None:
        raise NodeExecutionError("published Polars cast type is unsupported")
    return dtype


def _bounded_materialize(
    frame: pl.LazyFrame | pl.DataFrame,
    limits: dict[str, int],
    label: str,
) -> pl.DataFrame:
    try:
        collected = (
            frame.lazy() if isinstance(frame, pl.DataFrame) else frame
        ).head(limits["max_rows"] + 1).collect(engine="streaming")
    except Exception as exc:
        raise NodeExecutionError(
            f"published Polars {label} materialization failed"
        ) from exc
    if collected.height > limits["max_rows"]:
        raise NodeExecutionError(
            f"published Polars {label} exceeds its row limit"
        )
    if collected.estimated_size() > limits["memory_limit_bytes"]:
        raise NodeExecutionError(
            f"published Polars {label} exceeds its memory limit"
        )
    return collected


def _resolve_artifact(
    resolver,
    *,
    binding_id: str,
    binding_hash: str,
    correlation_id: str,
) -> dict[str, Any]:
    try:
        artifact = resolver.resolve(
            binding_id=binding_id,
            correlation_id=correlation_id,
        )
    except Exception as exc:
        raise NodeExecutionError(
            "published Polars artifact binding was not resolved"
        ) from exc
    if (
        not isinstance(artifact, dict)
        or artifact.get("binding_hash") != binding_hash
        or not isinstance(artifact.get("artifact_ref"), str)
        or not isinstance(artifact.get("digest"), str)
    ):
        raise NodeExecutionError(
            "published Polars artifact binding does not match"
        )
    return artifact


class PolarsRulePlanAdapter:
    """Execute one digest-bound Polars plan and write one bounded artifact."""

    def __init__(
        self,
        *,
        artifact_store,
        artifact_resolver,
        masking_policies=None,
        artifact_ttl_seconds=3600,
    ):
        self.artifact_store = artifact_store
        self.artifact_resolver = artifact_resolver
        self.masking_policies = dict(masking_policies or {})
        self.artifact_ttl_seconds = int(artifact_ttl_seconds)

    def execute(
        self,
        *,
        plan,
        node,
        parameters,
        write_authorized,
        correlation_id=None,
    ):
        try:
            normalized = validate_bound_polars_plan(plan)
        except ValueError as exc:
            raise NodeExecutionError(
                "published Polars rule plan is invalid"
            ) from exc
        config = node.get("config") or {}
        if config.get("execution_plan_hash") != bound_polars_plan_hash(
            normalized
        ):
            raise NodeExecutionError("published Polars rule plan hash does not match")
        if config.get("rule_version_id") != normalized["rule_version_id"]:
            raise NodeExecutionError("published Polars rule id does not match")
        idempotency = node.get("idempotency")
        if (
            node.get("type") != "rule.apply"
            or node.get("purpose") != "write"
            or not write_authorized
            or not isinstance(idempotency, dict)
            or idempotency.get("strategy") not in _IDEMPOTENCY
            or not str(idempotency.get("key") or "").strip()
        ):
            raise NodeExecutionError(
                "governed write authorization and idempotency are required"
            )
        if parameters not in ({}, None):
            raise NodeExecutionError(
                "bound Polars rule plans do not accept runtime parameters"
            )
        correlation = _uid(correlation_id, "correlation_id")
        try:
            self.artifact_resolver.attest_binding(
                binding_id=normalized["output_binding_id"],
                binding_hash=normalized["output_binding_hash"],
                access_mode="write",
            )
        except Exception as exc:
            raise NodeExecutionError(
                "published Polars output binding no longer matches"
            ) from exc
        source = _resolve_artifact(
            self.artifact_resolver,
            binding_id=normalized["input_binding_id"],
            binding_hash=normalized["input_binding_hash"],
            correlation_id=correlation,
        )
        try:
            frame = self.artifact_store.read(
                source["artifact_ref"],
                source["digest"],
                expected_schema_fields=normalized["input_fields"],
                limits=normalized["resource_limits"],
            )
        except ValueError as exc:
            raise NodeExecutionError(
                "published Polars input artifact is invalid"
            ) from exc
        limits = normalized["resource_limits"]
        initial = _bounded_materialize(frame, limits, "input")
        rows_in = initial.height
        frame = initial.lazy()
        expressions = _ExpressionCompiler()
        violations = []
        rows_rejected = 0
        rows_filtered = 0
        rows_deduplicated = 0
        rows_join_dropped = 0
        rows_aggregated = 0
        output_fields = {
            field["name"]: field for field in normalized["output_fields"]
        }

        for index, operation in enumerate(normalized["operations"]):
            op = operation["op"]
            before = _bounded_materialize(
                frame, limits, f"{op} input {index}"
            ).height
            if op == "normalize_text":
                expression = pl.col(operation["column"])
                if operation["trim"]:
                    expression = expression.str.strip_chars()
                if operation["lowercase"]:
                    expression = expression.str.to_lowercase()
                if operation["uppercase"]:
                    expression = expression.str.to_uppercase()
                frame = frame.with_columns(
                    expression.alias(operation["column"])
                )
            elif op == "regex_replace":
                frame = frame.with_columns(
                    pl.col(operation["column"])
                    .str.replace_all(
                        operation["pattern"],
                        operation["replacement"],
                    )
                    .alias(operation["column"])
                )
            elif op == "fill_null":
                frame = frame.with_columns(
                    pl.col(operation["column"])
                    .fill_null(operation["value"])
                    .alias(operation["column"])
                )
            elif op == "filter":
                frame = frame.filter(
                    expressions.compile(operation["expression_ast"]).fill_null(
                        False
                    )
                )
                after = _bounded_materialize(
                    frame, limits, f"{op} output {index}"
                )
                rows_filtered += before - after.height
                frame = after.lazy()
                continue
            elif op == "assert":
                predicate = expressions.compile(
                    operation["expression_ast"]
                ).fill_null(False)
                invalid = int(
                    frame.select((~predicate).sum().alias("count"))
                    .collect(engine="streaming")
                    .item()
                    or 0
                )
                violations.append(
                    {"step_id": operation["step_id"], "count": invalid}
                )
                rows_rejected += invalid
                frame = frame.filter(predicate)
            elif op == "derive":
                frame = frame.with_columns(
                    expressions.compile(operation["expression_ast"]).alias(
                        operation["target"]
                    )
                )
            elif op == "map_values":
                column = pl.col(operation["column"])
                frame = frame.with_columns(
                    column.replace_strict(
                        operation["mapping"],
                        default=column,
                    ).alias(operation["column"])
                )
            elif op == "cast":
                target_field = output_fields.get(operation["column"])
                if target_field is None:
                    target_field = {
                        "type": operation["to"],
                    }
                column_expression = pl.col(operation["column"])
                source_dtype = frame.collect_schema()[operation["column"]]
                if (
                    source_dtype == pl.String
                    and target_field["type"] == "timestamptz"
                ):
                    column_expression = column_expression.str.to_datetime(
                        time_zone=target_field["timezone"],
                        strict=True,
                    )
                elif (
                    source_dtype == pl.String
                    and target_field["type"] == "timestamp"
                ):
                    column_expression = column_expression.str.to_datetime(
                        strict=True,
                    )
                else:
                    column_expression = column_expression.cast(
                        _cast_type(target_field), strict=True
                    )
                frame = frame.with_columns(
                    column_expression.alias(operation["column"])
                )
            elif op == "deduplicate":
                order = [
                    *operation["order_by"],
                    *sorted(
                        name
                        for name in frame.collect_schema().names()
                        if name not in operation["order_by"]
                    ),
                ]
                descending = operation["keep"] == "last"
                frame = (
                    frame.sort(
                        order,
                        descending=[descending] * len(order),
                        nulls_last=True,
                    )
                    .unique(
                        subset=operation["keys"],
                        keep="first",
                        maintain_order=True,
                    )
                )
                after = _bounded_materialize(
                    frame, limits, f"{op} output {index}"
                )
                rows_deduplicated += before - after.height
                frame = after.lazy()
                continue
            elif op == "mask":
                if self.masking_policies.get(
                    operation["policy_id"]
                ) != operation["policy_kind"]:
                    raise NodeExecutionError(
                        "published masking policy is not registered"
                    )
                column = pl.col(operation["column"])
                if operation["policy_kind"] == "redact":
                    masked = pl.when(column.is_null()).then(None).otherwise(
                        pl.lit("***")
                    )
                elif operation["policy_kind"] == "preserve_last_4":
                    masked = pl.when(column.is_null()).then(None).otherwise(
                        pl.lit("***") + column.str.slice(-4)
                    )
                else:
                    raise NodeExecutionError(
                        "published masking policy is unsupported"
                    )
                frame = frame.with_columns(masked.alias(operation["column"]))
            elif op == "aggregate":
                aggregations = []
                for aggregate in operation["aggregations"]:
                    expression = pl.col(aggregate["column"])
                    function = aggregate["function"]
                    if function == "count":
                        expression = expression.count()
                    elif function == "sum":
                        expression = expression.sum()
                    elif function == "min":
                        expression = expression.min()
                    elif function == "max":
                        expression = expression.max()
                    elif function == "mean":
                        expression = expression.mean()
                    aggregations.append(
                        expression.alias(aggregate["target"])
                    )
                frame = frame.group_by(
                    operation["group_by"], maintain_order=True
                ).agg(aggregations)
                after = _bounded_materialize(
                    frame, limits, f"{op} output {index}"
                )
                rows_aggregated += before - after.height
                frame = after.lazy()
                continue
            elif op == "lookup_join":
                lookup_artifact = _resolve_artifact(
                    self.artifact_resolver,
                    binding_id=operation["lookup_binding_id"],
                    binding_hash=operation["lookup_binding_hash"],
                    correlation_id=correlation,
                )
                try:
                    lookup = self.artifact_store.read(
                        lookup_artifact["artifact_ref"],
                        lookup_artifact["digest"],
                        expected_schema_fields=operation["lookup_fields"],
                        limits=limits,
                    )
                except ValueError as exc:
                    raise NodeExecutionError(
                        "published Polars lookup artifact is invalid"
                    ) from exc
                lookup = _bounded_materialize(
                    lookup, limits, f"lookup input {index}"
                ).lazy()
                duplicate_count = (
                    lookup.group_by(operation["right_on"])
                    .len()
                    .filter(pl.col("len") > 1)
                    .select(pl.len())
                    .collect(engine="streaming")
                    .item()
                )
                if duplicate_count:
                    raise NodeExecutionError(
                        "published Polars lookup keys are not unique"
                    )
                selected_sources = list(operation["select"].values())
                lookup = lookup.select(
                    list(dict.fromkeys(
                        [*operation["right_on"], *selected_sources]
                    ))
                ).rename(
                    {
                        source_name: target_name
                        for target_name, source_name in operation[
                            "select"
                        ].items()
                        if source_name not in operation["right_on"]
                    }
                )
                frame = frame.join(
                    lookup,
                    left_on=operation["left_on"],
                    right_on=operation["right_on"],
                    how=operation["how"],
                )
                after = _bounded_materialize(
                    frame, limits, f"{op} output {index}"
                )
                if after.height > before:
                    raise NodeExecutionError(
                        "published Polars lookup join expanded its input"
                    )
                rows_join_dropped += before - after.height
                frame = after.lazy()
                continue

            frame = _bounded_materialize(
                frame, limits, f"{op} output {index}"
            ).lazy()

        output_names = [field["name"] for field in normalized["output_fields"]]
        frame = frame.select(output_names)
        bounded = _bounded_materialize(frame, limits, "output")
        try:
            artifact = self.artifact_store.write(
                bounded.lazy(),
                correlation,
                self.artifact_ttl_seconds,
                schema_fields=normalized["output_fields"],
                limits=limits,
            )
        except ValueError as exc:
            raise NodeExecutionError(
                "published Polars output artifact write failed"
            ) from exc
        try:
            self.artifact_resolver.register(
                binding_id=normalized["output_binding_id"],
                correlation_id=correlation,
                artifact=artifact,
                kind="output",
                binding_hash=normalized["output_binding_hash"],
            )
        except Exception as exc:
            delete = getattr(self.artifact_store, "delete", None)
            if callable(delete):
                with suppress(Exception):
                    delete(artifact["artifact_ref"])
            raise NodeExecutionError(
                "published Polars output artifact registration failed"
            ) from exc
        violation_count = sum(item["count"] for item in violations)
        return {
            **artifact,
            "rows_in": rows_in,
            "rows_out": bounded.height,
            "rows_rejected": rows_rejected,
            "rows_filtered": rows_filtered,
            "rows_deduplicated": rows_deduplicated,
            "rows_join_dropped": rows_join_dropped,
            "rows_aggregated": rows_aggregated,
            "violation_count": violation_count,
            "violations": violations,
            "commit_outcome": "committed",
        }