DataOps-platform/app/runner/polars_worker.py

"""OS-isolated execution boundary for allocation-heavy Polars work."""

from __future__ import annotations

import multiprocessing
import os
import time
from contextlib import suppress
from decimal import Decimal
from typing import Any

try:
    import resource
except ImportError:  # pragma: no cover - fail-closed portability guard
    resource = None

import polars as pl
import psutil

from app.core.data_rules.compilers.polars import validate_bound_polars_plan


class PolarsWorkerError(RuntimeError):
    """A bounded worker failed without exposing internal data."""


class PolarsWorkerResourceError(PolarsWorkerError):
    """The worker crossed a hard process resource boundary."""


def _set_process_memory_limit(
    *,
    baseline_vms: int,
    memory_limit_bytes: int,
) -> None:
    if resource is None or not hasattr(resource, "RLIMIT_AS"):
        raise OSError("OS address-space limits are unavailable")
    hard_limit = baseline_vms + memory_limit_bytes
    resource.setrlimit(resource.RLIMIT_AS, (hard_limit, hard_limit))


def _memory_probe_entry(
    connection,
    allocate_bytes: int,
    memory_limit_bytes: int,
) -> None:
    process = psutil.Process()
    memory = process.memory_info()
    connection.send(
        {
            "kind": "ready",
            "pid": os.getpid(),
            "baseline_rss": memory.rss,
        }
    )
    try:
        _set_process_memory_limit(
            baseline_vms=memory.vms,
            memory_limit_bytes=memory_limit_bytes,
        )
        allocated = bytearray(allocate_bytes)
        for index in range(0, len(allocated), 4096):
            allocated[index] = 1
        connection.send(
            {
                "kind": "result",
                "worker_pid": os.getpid(),
                "allocated_bytes": len(allocated),
            }
        )
    except (MemoryError, OSError):
        with suppress(Exception):
            connection.send({"kind": "resource_error"})
    except BaseException:
        with suppress(Exception):
            connection.send({"kind": "worker_error"})
    finally:
        connection.close()


class _ExpressionCompiler:
    def __init__(self, timezone: str):
        self.timezone = timezone

    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":
            return arguments[0].str.contains(
                ast["arguments"][1]["value"], 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.String)
                .str.to_datetime(time_zone=self.timezone, strict=True)
                .dt.date()
            )
        if function == "timestamp":
            return (
                arguments[0]
                .cast(pl.String)
                .str.to_datetime(time_zone=self.timezone, strict=True)
                .dt.replace_time_zone(None)
            )
        if function == "abs":
            return arguments[0].abs()
        raise ValueError("published Polars expression is unsupported")


def _cast_type(field: dict[str, Any]) -> pl.DataType:
    field_type = field["type"]
    if field_type == "boolean":
        return pl.Boolean
    if field_type == "date":
        return pl.Date
    if field_type == "decimal":
        return pl.Decimal(
            precision=field["precision"],
            scale=field["scale"],
        )
    if field_type == "double":
        return pl.Float64
    if field_type == "float":
        return pl.Float32
    if field_type == "integer":
        return pl.Int64
    if field_type == "string":
        return pl.String
    if field_type == "timestamp":
        return pl.Datetime
    if field_type == "timestamptz":
        return pl.Datetime(time_zone=field["timezone"])
    raise ValueError("published Polars cast type is unsupported")


def _count(frame: pl.LazyFrame) -> int:
    return int(frame.select(pl.len()).collect(engine="streaming").item())


def _validate_lazy_schema(
    frame: pl.LazyFrame,
    fields: list[dict[str, Any]],
) -> None:
    schema = frame.collect_schema()
    if set(schema.names()) != {field["name"] for field in fields}:
        raise ValueError("artifact schema fields do not match")
    for field in fields:
        if schema[field["name"]] != _cast_type(field):
            raise ValueError("artifact schema type does not match")
    required = [
        field["name"] for field in fields if not field["nullable"]
    ]
    if required:
        counts = frame.select(
            [pl.col(name).null_count().alias(name) for name in required]
        ).collect(engine="streaming")
        if any(counts[name].item() for name in required):
            raise ValueError("artifact nullable contract does not match")


def _execute_polars_job(job: dict[str, Any]) -> dict[str, Any]:
    plan = validate_bound_polars_plan(job["plan"])
    frame = pl.scan_parquet(job["input_path"])
    _validate_lazy_schema(frame, plan["input_fields"])
    rows_in = _count(frame)
    expressions = _ExpressionCompiler(plan["timezone"])
    violations = []
    metrics = {
        "rows_rejected": 0,
        "rows_filtered": 0,
        "rows_deduplicated": 0,
        "rows_join_dropped": 0,
        "rows_aggregated": 0,
    }
    for operation in plan["operations"]:
        op = operation["op"]
        before = _count(frame)
        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
                )
            )
            metrics["rows_filtered"] += before - _count(frame)
        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}
            )
            metrics["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 = operation["target_field"]
            expression = pl.col(operation["column"])
            source_type = frame.collect_schema()[operation["column"]]
            if source_type == pl.String and target["type"] == "timestamp":
                expression = (
                    expression.str.to_datetime(
                        time_zone=plan["timezone"],
                        strict=True,
                    ).dt.replace_time_zone(None)
                )
            elif (
                source_type == pl.String
                and target["type"] == "timestamptz"
            ):
                expression = expression.str.to_datetime(
                    time_zone=target["timezone"], strict=True
                )
            else:
                expression = expression.cast(
                    _cast_type(target), strict=True
                )
            frame = frame.with_columns(
                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,
                )
            )
            metrics["rows_deduplicated"] += before - _count(frame)
        elif op == "mask":
            if job["masking_policies"].get(
                operation["policy_id"]
            ) != operation["policy_kind"]:
                raise ValueError("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("***")
                )
            else:
                masked = pl.when(column.is_null()).then(None).otherwise(
                    pl.lit("***") + column.str.slice(-4)
                )
            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)
            metrics["rows_aggregated"] += before - _count(frame)
        elif op == "lookup_join":
            lookup = pl.scan_parquet(
                job["lookup_paths"][operation["lookup_binding_id"]]
            )
            _validate_lazy_schema(lookup, operation["lookup_fields"])
            duplicates = (
                lookup.group_by(operation["right_on"])
                .len()
                .filter(pl.col("len") > 1)
                .select(pl.len())
                .collect(engine="streaming")
                .item()
            )
            if duplicates:
                raise ValueError("lookup keys are not unique")
            selected = list(operation["select"].values())
            lookup = lookup.select(
                list(
                    dict.fromkeys(
                        [*operation["right_on"], *selected]
                    )
                )
            ).rename(
                {
                    source: target
                    for target, source in operation["select"].items()
                    if source not in operation["right_on"]
                }
            )
            frame = frame.join(
                lookup,
                left_on=operation["left_on"],
                right_on=operation["right_on"],
                how=operation["how"],
            )
            after = _count(frame)
            if after > before:
                raise ValueError("lookup join expanded its input")
            metrics["rows_join_dropped"] += before - after

    frame = frame.select(
        [field["name"] for field in plan["output_fields"]]
    )
    _validate_lazy_schema(frame, plan["output_fields"])
    rows_out = _count(frame)
    if rows_out > plan["resource_limits"]["max_rows"]:
        raise ValueError("output row limit exceeded")
    frame.sink_parquet(job["output_path"], engine="streaming")
    return {
        "worker_pid": os.getpid(),
        "rows_in": rows_in,
        "rows_out": rows_out,
        **metrics,
        "violation_count": sum(item["count"] for item in violations),
        "violations": violations,
    }


def _polars_job_entry(
    connection,
    job: dict[str, Any],
    memory_limit_bytes: int,
) -> None:
    # Initialize the Rust engine/thread pool before measuring the worker's
    # immutable runtime baseline. The plan allowance applies only afterwards.
    pl.DataFrame({"_warmup": [1]}).lazy().collect(engine="streaming")
    process = psutil.Process()
    memory = process.memory_info()
    connection.send(
        {
            "kind": "ready",
            "pid": os.getpid(),
            "baseline_rss": memory.rss,
        }
    )
    try:
        _set_process_memory_limit(
            baseline_vms=memory.vms,
            memory_limit_bytes=memory_limit_bytes,
        )
        result = _execute_polars_job(job)
        connection.send({"kind": "result", **result})
    except (MemoryError, OSError):
        with suppress(Exception):
            connection.send({"kind": "resource_error"})
    except BaseException as exc:
        message = str(exc).lower()
        kind = (
            "resource_error"
            if any(
                token in message
                for token in ("memory", "allocate", "allocation")
            )
            else "worker_error"
        )
        with suppress(Exception):
            connection.send({"kind": kind})
    finally:
        connection.close()


def _run_process(
    *,
    target,
    args: tuple[Any, ...],
    memory_limit_bytes: int,
    timeout_seconds: float = 15.0,
) -> dict[str, Any]:
    parent = None
    child = None
    process = None
    started = False
    try:
        context = multiprocessing.get_context("spawn")
        parent, child = context.Pipe(duplex=False)
        process = context.Process(
            target=target,
            args=(child, *args, memory_limit_bytes),
            daemon=True,
        )
        process.start()
        started = True
        child.close()
        deadline = time.monotonic() + timeout_seconds
        baseline_rss = None
        while time.monotonic() < deadline:
            if parent.poll(0.01):
                try:
                    message = parent.recv()
                except EOFError as exc:
                    raise PolarsWorkerResourceError(
                        "Polars worker crossed its hard memory limit"
                    ) from exc
                kind = message.get("kind")
                if kind == "ready":
                    baseline_rss = int(message["baseline_rss"])
                elif kind == "result":
                    process.join(timeout=1)
                    return message
                elif kind == "resource_error":
                    raise PolarsWorkerResourceError(
                        "Polars worker exceeded its hard memory limit"
                    )
                else:
                    raise PolarsWorkerError(
                        "Polars worker execution failed"
                    )
            if baseline_rss is not None and process.is_alive():
                try:
                    rss = psutil.Process(process.pid).memory_info().rss
                except (psutil.Error, ProcessLookupError):
                    rss = baseline_rss
                if rss - baseline_rss > memory_limit_bytes:
                    process.kill()
                    process.join(timeout=1)
                    raise PolarsWorkerResourceError(
                        "Polars worker exceeded its hard memory limit"
                    )
            if not process.is_alive():
                process.join(timeout=1)
                raise PolarsWorkerResourceError(
                    "Polars worker crossed its hard memory limit"
                )
        process.kill()
        process.join(timeout=1)
        raise PolarsWorkerResourceError(
            "Polars worker exceeded its bounded execution time"
        )
    except (PolarsWorkerError, PolarsWorkerResourceError):
        raise
    except Exception as exc:
        raise PolarsWorkerError(
            "Polars worker failed to start safely"
        ) from exc
    finally:
        if child is not None:
            with suppress(Exception):
                child.close()
        if parent is not None:
            with suppress(Exception):
                parent.close()
        if process is not None:
            if started:
                with suppress(Exception):
                    if process.is_alive():
                        process.kill()
                with suppress(Exception):
                    process.join(timeout=1)
            close = getattr(process, "close", None)
            if callable(close):
                with suppress(Exception):
                    close()


def run_isolated_memory_probe(
    *,
    allocate_bytes: int,
    memory_limit_bytes: int,
) -> dict[str, Any]:
    """Exercise the same process boundary used by plan execution."""

    return _run_process(
        target=_memory_probe_entry,
        args=(int(allocate_bytes),),
        memory_limit_bytes=int(memory_limit_bytes),
    )


def execute_isolated_polars_plan(
    job: dict[str, Any],
    *,
    memory_limit_bytes: int,
) -> dict[str, Any]:
    return _run_process(
        target=_polars_job_entry,
        args=(job,),
        memory_limit_bytes=memory_limit_bytes,
        timeout_seconds=60.0,
    )