DataOps-platform/app/runner/rule_polars.py

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

from __future__ import annotations

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,
    "timestamptz": 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


def _dtype_name(dtype: pl.DataType) -> str:
    if dtype == pl.Boolean:
        return "boolean"
    if dtype == pl.Date:
        return "date"
    if dtype == pl.String:
        return "string"
    if dtype.is_integer():
        return "integer"
    if dtype == pl.Float32:
        return "float"
    if dtype == pl.Float64:
        return "double"
    if dtype.is_decimal():
        return "decimal"
    if isinstance(dtype, pl.Datetime):
        return "timestamptz" if dtype.time_zone else "timestamp"
    return str(dtype).lower()


def _validate_schema(
    frame: pl.LazyFrame,
    expected_fields: list[dict[str, Any]],
    label: str,
) -> None:
    schema = frame.collect_schema()
    expected = {field["name"]: field["type"] for field in expected_fields}
    actual = {name: _dtype_name(dtype) for name, dtype in schema.items()}
    if actual != expected:
        raise NodeExecutionError(
            f"{label} artifact schema does not match the published plan"
        )


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(name: str) -> pl.DataType:
    if name == "decimal":
        return pl.Decimal(38, 9)
    dtype = _TYPE_MAP.get(name)
    if dtype is None:
        raise NodeExecutionError("published Polars cast type is unsupported")
    return dtype


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")
        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"]
            )
        except ValueError as exc:
            raise NodeExecutionError(
                "published Polars input artifact is invalid"
            ) from exc
        _validate_schema(frame, normalized["input_fields"], "input")
        limits = normalized["resource_limits"]
        initial = frame.head(limits["max_rows"] + 1).collect(engine="streaming")
        if initial.height > limits["max_rows"]:
            raise NodeExecutionError(
                "published Polars input exceeds its row limit"
            )
        if initial.estimated_size() > limits["memory_limit_bytes"]:
            raise NodeExecutionError(
                "published Polars input exceeds its memory limit"
            )
        rows_in = initial.height
        frame = initial.lazy()
        expressions = _ExpressionCompiler()
        violations = []

        for operation in normalized["operations"]:
            op = operation["op"]
            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
                    )
                )
            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}
                )
                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":
                frame = frame.with_columns(
                    pl.col(operation["column"])
                    .cast(_cast_type(operation["to"]), strict=True)
                    .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,
                    )
                )
            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)
            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"],
                    )
                except ValueError as exc:
                    raise NodeExecutionError(
                        "published Polars lookup artifact is invalid"
                    ) from exc
                _validate_schema(
                    lookup, operation["lookup_fields"], "lookup"
                )
                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"],
                )

        output_names = [field["name"] for field in normalized["output_fields"]]
        frame = frame.select(output_names)
        _validate_schema(frame, normalized["output_fields"], "output")
        bounded = frame.head(limits["max_rows"] + 1).collect(engine="streaming")
        if bounded.height > limits["max_rows"]:
            raise NodeExecutionError(
                "published Polars output exceeds its row limit"
            )
        if bounded.estimated_size() > limits["memory_limit_bytes"]:
            raise NodeExecutionError(
                "published Polars output exceeds its memory limit"
            )
        try:
            artifact = self.artifact_store.write(
                bounded.lazy(),
                correlation,
                self.artifact_ttl_seconds,
            )
        except ValueError as exc:
            raise NodeExecutionError(
                "published Polars output artifact write failed"
            ) from exc
        violation_count = sum(item["count"] for item in violations)
        return {
            **artifact,
            "rows_in": rows_in,
            "rows_out": bounded.height,
            "rows_rejected": max(0, rows_in - bounded.height),
            "violation_count": violation_count,
            "violations": violations,
            "commit_outcome": "committed",
        }