DataOps-platform/tests/runner/test_polars_worker.py

from __future__ import annotations

import os

import polars as pl
import pytest

from tests.core.data_rules.test_polars_compiler import (
    _backend,
    _binding,
    _published_rule,
    _schema,
)


def _compile_worker_plan(
    *,
    input_schema,
    output_schema,
    steps,
    lookup_context=None,
):
    from app.core.data_rules.compilers.polars import PolarsRuleCompiler

    input_binding = _binding(input_schema, access_mode="read")
    output_binding = _binding(output_schema, access_mode="write")
    backend = _backend(
        max_rows=1_000_000,
        max_artifact_bytes=256 * 1024 * 1024,
        memory_limit_bytes=16 * 1024 * 1024,
        lookup_bindings=lookup_context or {},
    )
    rule = _published_rule(
        input_schema,
        output_schema,
        steps,
    )
    return PolarsRuleCompiler().compile(
        rule_version=rule,
        input_schema=input_schema,
        output_schema=output_schema,
        input_binding=input_binding,
        output_binding=output_binding,
        backend=backend,
    )["plan"]


def _assert_worker_resource_failure(plan, tmp_path, *, lookup_paths=None):
    from app.runner.polars_worker import (
        PolarsWorkerResourceError,
        execute_isolated_polars_plan,
    )

    with pytest.raises(
        PolarsWorkerResourceError,
        match="hard memory limit",
    ):
        execute_isolated_polars_plan(
            {
                "plan": plan,
                "input_path": str(tmp_path / "input.parquet"),
                "lookup_paths": lookup_paths or {},
                "output_path": str(tmp_path / "output.parquet"),
                "masking_policies": {},
            },
            memory_limit_bytes=plan["resource_limits"][
                "memory_limit_bytes"
            ],
        )


def test_isolated_polars_worker_runs_outside_runner_process():
    from app.runner.polars_worker import run_isolated_memory_probe

    result = run_isolated_memory_probe(
        allocate_bytes=1_024,
        memory_limit_bytes=16 * 1024 * 1024,
    )

    assert result["worker_pid"] != os.getpid()
    assert result["allocated_bytes"] == 1_024


def test_isolated_polars_worker_fails_deterministically_at_hard_memory_limit():
    from app.runner.polars_worker import (
        PolarsWorkerResourceError,
        run_isolated_memory_probe,
    )

    with pytest.raises(
        PolarsWorkerResourceError,
        match="hard memory limit",
    ):
        run_isolated_memory_probe(
            allocate_bytes=128 * 1024 * 1024,
            memory_limit_bytes=8 * 1024 * 1024,
        )


def test_regex_peak_allocation_fails_inside_isolated_worker(tmp_path):
    schema = _schema(
        "bd:regex:raw",
        [("text", "string", False)],
    )
    plan = _compile_worker_plan(
        input_schema=schema,
        output_schema={
            **schema,
            "id": _schema(
                "bd:regex:clean",
                [("text", "string", False)],
            )["id"],
            "schema_ref": "bd:regex:clean",
        },
        steps=[
            {
                "id": "expand",
                "op": "regex_replace",
                "column": "text",
                "pattern": "a",
                "replacement": "x" * 64,
            }
        ],
    )
    pl.DataFrame({"text": ["a" * 1_000_000]}).write_parquet(
        tmp_path / "input.parquet"
    )

    _assert_worker_resource_failure(plan, tmp_path)


def test_group_peak_allocation_fails_inside_isolated_worker(tmp_path):
    input_schema = _schema(
        "bd:group:raw",
        [("group_key", "string", False), ("amount", "integer", False)],
    )
    output_schema = _schema(
        "bd:group:clean",
        [
            ("group_key", "string", False),
            ("total_amount", "integer", False),
        ],
    )
    plan = _compile_worker_plan(
        input_schema=input_schema,
        output_schema=output_schema,
        steps=[
            {
                "id": "sum_by_key",
                "op": "aggregate",
                "group_by": ["group_key"],
                "aggregations": {
                    "total_amount": {
                        "function": "sum",
                        "column": "amount",
                    }
                },
            }
        ],
    )
    row_count = 300_000
    pl.DataFrame(
        {
            "group_key": [
                f"group-{index:040d}" for index in range(row_count)
            ],
            "amount": [1] * row_count,
        }
    ).write_parquet(tmp_path / "input.parquet")

    _assert_worker_resource_failure(plan, tmp_path)


def test_join_peak_allocation_fails_inside_isolated_worker(tmp_path):
    input_schema = _schema(
        "bd:join:raw",
        [("id", "integer", False)],
    )
    lookup_schema = _schema(
        "bd:join:lookup",
        [("lookup_id", "integer", False), ("label", "string", False)],
    )
    output_schema = _schema(
        "bd:join:clean",
        [("id", "integer", False), ("label", "string", False)],
    )
    lookup_binding = _binding(
        lookup_schema,
        access_mode="read",
        object_ref="join-lookup",
    )
    plan = _compile_worker_plan(
        input_schema=input_schema,
        output_schema=output_schema,
        steps=[
            {
                "id": "join_label",
                "op": "lookup_join",
                "lookup": {
                    "binding_id": lookup_binding["id"],
                    "left_on": ["id"],
                    "right_on": ["lookup_id"],
                    "select": {"label": "label"},
                    "how": "left",
                },
            }
        ],
        lookup_context={
            lookup_binding["id"]: {
                "binding": lookup_binding,
                "schema": lookup_schema,
            }
        },
    )
    row_count = 250_000
    pl.DataFrame({"id": range(row_count)}).write_parquet(
        tmp_path / "input.parquet"
    )
    lookup_path = tmp_path / "lookup.parquet"
    pl.DataFrame(
        {
            "lookup_id": range(row_count),
            "label": [f"label-{index:048d}" for index in range(row_count)],
        }
    ).write_parquet(lookup_path)

    _assert_worker_resource_failure(
        plan,
        tmp_path,
        lookup_paths={lookup_binding["id"]: str(lookup_path)},
    )