test_apply_checks.py

from pathlib import Path
import pyspark.sql.functions as F
import pytest
from pyspark.sql import Column
from chispa.dataframe_comparer import assert_df_equality  # type: ignore
from databricks.labs.dqx.col_functions import is_not_null_and_not_empty, make_condition
from databricks.labs.dqx.engine import (
    DQEngine,
    ExtraParams,
)
from databricks.labs.dqx.rule import DQRule

SCHEMA = "a: int, b: int, c: int"
EXPECTED_SCHEMA = SCHEMA + ", _errors: map<string,string>, _warnings: map<string,string>"
EXPECTED_SCHEMA_WITH_CUSTOM_NAMES = SCHEMA + ", ERROR: map<string,string>, WARN: map<string,string>"


def test_apply_checks_on_empty_checks(ws, spark):
    dq_engine = DQEngine(ws)
    test_df = spark.createDataFrame([[1, 3, None], [2, 4, None]], SCHEMA)

    good = dq_engine.apply_checks(test_df, [])

    expected_df = spark.createDataFrame([[1, 3, None, None, None], [2, 4, None, None, None]], EXPECTED_SCHEMA)
    assert_df_equality(good, expected_df)


def test_apply_checks_and_split_on_empty_checks(ws, spark):
    dq_engine = DQEngine(ws)
    test_df = spark.createDataFrame([[1, 3, None], [2, 4, None]], SCHEMA)

    good, bad = dq_engine.apply_checks_and_split(test_df, [])

    expected_df = spark.createDataFrame([], EXPECTED_SCHEMA)

    assert_df_equality(good, test_df)
    assert_df_equality(bad, expected_df)


def test_apply_checks_passed(ws, spark):
    dq_engine = DQEngine(ws)
    test_df = spark.createDataFrame([[1, 3, 3]], SCHEMA)

    checks = [
        DQRule(name="col_a_is_null_or_empty", criticality="warn", check=is_not_null_and_not_empty("a")),
        DQRule(name="col_b_is_null_or_empty", criticality="error", check=is_not_null_and_not_empty("b")),
    ]

    checked = dq_engine.apply_checks(test_df, checks)

    expected = spark.createDataFrame([[1, 3, 3, None, None]], EXPECTED_SCHEMA)

    assert_df_equality(checked, expected, ignore_nullable=True)


def test_apply_checks(ws, spark):
    dq_engine = DQEngine(ws)
    test_df = spark.createDataFrame([[1, 3, 3], [2, None, 4], [None, 4, None], [None, None, None]], SCHEMA)

    checks = [
        DQRule(name="col_a_is_null_or_empty", criticality="warn", check=is_not_null_and_not_empty("a")),
        DQRule(name="col_b_is_null_or_empty", criticality="error", check=is_not_null_and_not_empty("b")),
        DQRule(name="col_c_is_null_or_empty", criticality="error", check=is_not_null_and_not_empty("c")),
    ]

    checked = dq_engine.apply_checks(test_df, checks)

    expected = spark.createDataFrame(
        [
            [1, 3, 3, None, None],
            [2, None, 4, {"col_b_is_null_or_empty": "Column b is null or empty"}, None],
            [
                None,
                4,
                None,
                {"col_c_is_null_or_empty": "Column c is null or empty"},
                {"col_a_is_null_or_empty": "Column a is null or empty"},
            ],
            [
                None,
                None,
                None,
                {
                    "col_b_is_null_or_empty": "Column b is null or empty",
                    "col_c_is_null_or_empty": "Column c is null or empty",
                },
                {"col_a_is_null_or_empty": "Column a is null or empty"},
            ],
        ],
        EXPECTED_SCHEMA,
    )

    assert_df_equality(checked, expected, ignore_nullable=True)


def test_apply_checks_invalid_criticality(ws, spark):
    dq_engine = DQEngine(ws)
    test_df = spark.createDataFrame([[1, 3, 3], [2, None, 4], [None, 4, None], [None, None, None]], SCHEMA)

    checks = [
        DQRule(name="col_a_is_null_or_empty", criticality="warn", check=is_not_null_and_not_empty("a")),
        DQRule(name="col_b_is_null_or_empty", criticality="error", check=is_not_null_and_not_empty("b")),
        DQRule(name="col_c_is_null_or_empty", criticality="invalid", check=is_not_null_and_not_empty("c")),
    ]

    with pytest.raises(ValueError, match="Invalid criticality value: invalid"):
        dq_engine.apply_checks(test_df, checks)


def test_apply_checks_with_autogenerated_col_names(ws, spark):
    dq_engine = DQEngine(ws)
    test_df = spark.createDataFrame([[1, 3, 3], [2, None, 4], [None, 4, None], [None, None, None]], SCHEMA)

    checks = [
        DQRule(criticality="warn", check=is_not_null_and_not_empty("a")),
        DQRule(criticality="error", check=is_not_null_and_not_empty("b")),
        DQRule(criticality="error", check=is_not_null_and_not_empty("c")),
    ]

    checked = dq_engine.apply_checks(test_df, checks)

    expected = spark.createDataFrame(
        [
            [1, 3, 3, None, None],
            [2, None, 4, {"col_b_is_null_or_empty": "Column b is null or empty"}, None],
            [
                None,
                4,
                None,
                {"col_c_is_null_or_empty": "Column c is null or empty"},
                {"col_a_is_null_or_empty": "Column a is null or empty"},
            ],
            [
                None,
                None,
                None,
                {
                    "col_b_is_null_or_empty": "Column b is null or empty",
                    "col_c_is_null_or_empty": "Column c is null or empty",
                },
                {"col_a_is_null_or_empty": "Column a is null or empty"},
            ],
        ],
        EXPECTED_SCHEMA,
    )

    assert_df_equality(checked, expected, ignore_nullable=True)


def test_apply_checks_and_split(ws, spark):
    dq_engine = DQEngine(ws)
    test_df = spark.createDataFrame([[1, 3, 3], [2, None, 4], [None, 4, None], [None, None, None]], SCHEMA)

    checks = [
        DQRule(name="col_a_is_null_or_empty", criticality="warn", check=is_not_null_and_not_empty("a")),
        DQRule(name="col_b_is_null_or_empty", criticality="error", check=is_not_null_and_not_empty("b")),
        DQRule(name="col_c_is_null_or_empty", criticality="warn", check=is_not_null_and_not_empty("c")),
    ]

    good, bad = dq_engine.apply_checks_and_split(test_df, checks)

    expected_good = spark.createDataFrame([[1, 3, 3], [None, 4, None]], SCHEMA)
    assert_df_equality(good, expected_good)

    expected_bad = spark.createDataFrame(
        [
            [2, None, 4, {"col_b_is_null_or_empty": "Column b is null or empty"}, None],
            [
                None,
                4,
                None,
                None,
                {
                    "col_a_is_null_or_empty": "Column a is null or empty",
                    "col_c_is_null_or_empty": "Column c is null or empty",
                },
            ],
            [
                None,
                None,
                None,
                {"col_b_is_null_or_empty": "Column b is null or empty"},
                {
                    "col_a_is_null_or_empty": "Column a is null or empty",
                    "col_c_is_null_or_empty": "Column c is null or empty",
                },
            ],
        ],
        EXPECTED_SCHEMA,
    )

    assert_df_equality(bad, expected_bad, ignore_nullable=True)


def test_apply_checks_and_split_by_metadata(ws, spark):
    dq_engine = DQEngine(ws)
    test_df = spark.createDataFrame([[1, 3, 3], [2, None, 4], [None, 4, None], [None, None, None]], SCHEMA)

    checks = [
        {
            "name": "col_a_is_null_or_empty",
            "criticality": "warn",
            "check": {"function": "is_not_null_and_not_empty", "arguments": {"col_name": "a"}},
        },
        {
            "name": "col_b_is_null_or_empty",
            "criticality": "error",
            "check": {"function": "is_not_null_and_not_empty", "arguments": {"col_name": "b"}},
        },
        {
            "name": "col_c_is_null_or_empty",
            "criticality": "warn",
            "check": {"function": "is_not_null_and_not_empty", "arguments": {"col_name": "c"}},
        },
        {
            "name": "col_a_value_is_not_in_the_list",
            "criticality": "warn",
            "check": {"function": "value_is_in_list", "arguments": {"col_name": "a", "allowed": [1, 3, 4]}},
        },
        {
            "name": "col_c_value_is_not_in_the_list",
            "criticality": "warn",
            "check": {"function": "value_is_in_list", "arguments": {"col_name": "c", "allowed": [1, 3, 4]}},
        },
    ]

    good, bad = dq_engine.apply_checks_by_metadata_and_split(test_df, checks)

    expected_good = spark.createDataFrame([[1, 3, 3], [None, 4, None]], SCHEMA)
    assert_df_equality(good, expected_good)

    expected_bad = spark.createDataFrame(
        [
            [
                2,
                None,
                4,
                {"col_b_is_null_or_empty": "Column b is null or empty"},
                {"col_a_value_is_not_in_the_list": "Value 2 is not in the allowed list: [1, 3, 4]"},
            ],
            [
                None,
                4,
                None,
                None,
                {
                    "col_a_is_null_or_empty": "Column a is null or empty",
                    "col_c_is_null_or_empty": "Column c is null or empty",
                },
            ],
            [
                None,
                None,
                None,
                {"col_b_is_null_or_empty": "Column b is null or empty"},
                {
                    "col_a_is_null_or_empty": "Column a is null or empty",
                    "col_c_is_null_or_empty": "Column c is null or empty",
                },
            ],
        ],
        EXPECTED_SCHEMA,
    )

    assert_df_equality(bad, expected_bad, ignore_nullable=True)


def test_apply_checks_and_split_by_metadata_with_autogenerated_col_names(ws, spark):
    dq_engine = DQEngine(ws)
    test_df = spark.createDataFrame([[1, 3, 3], [2, None, 4], [None, 4, None], [None, None, None]], SCHEMA)

    checks = [
        {
            "criticality": "warn",
            "check": {"function": "is_not_null_and_not_empty", "arguments": {"col_names": ["a", "c"]}},
        },
        {
            "criticality": "error",
            "check": {"function": "is_not_null_and_not_empty", "arguments": {"col_name": "b"}},
        },
        {
            "criticality": "warn",
            "check": {"function": "value_is_in_list", "arguments": {"col_names": ["a", "c"], "allowed": [1, 3, 4]}},
        },
    ]

    good, bad = dq_engine.apply_checks_by_metadata_and_split(test_df, checks)

    expected_good = spark.createDataFrame([[1, 3, 3], [None, 4, None]], SCHEMA)
    assert_df_equality(good, expected_good)

    expected_bad = spark.createDataFrame(
        [
            [
                2,
                None,
                4,
                {"col_b_is_null_or_empty": "Column b is null or empty"},
                {"col_a_value_is_not_in_the_list": "Value 2 is not in the allowed list: [1, 3, 4]"},
            ],
            [
                None,
                4,
                None,
                None,
                {
                    "col_a_is_null_or_empty": "Column a is null or empty",
                    "col_c_is_null_or_empty": "Column c is null or empty",
                },
            ],
            [
                None,
                None,
                None,
                {"col_b_is_null_or_empty": "Column b is null or empty"},
                {
                    "col_a_is_null_or_empty": "Column a is null or empty",
                    "col_c_is_null_or_empty": "Column c is null or empty",
                },
            ],
        ],
        EXPECTED_SCHEMA,
    )

    assert_df_equality(bad, expected_bad, ignore_nullable=True)


def test_apply_checks_by_metadata(ws, spark):
    dq_engine = DQEngine(ws)
    test_df = spark.createDataFrame([[1, 3, 3], [2, None, 4], [None, 4, None], [None, None, None]], SCHEMA)

    checks = [
        {
            "criticality": "warn",
            "check": {"function": "is_not_null_and_not_empty", "arguments": {"col_names": ["a", "c"]}},
        },
        {
            "criticality": "error",
            "check": {"function": "is_not_null_and_not_empty", "arguments": {"col_name": "b"}},
        },
        {
            "criticality": "warn",
            "check": {"function": "value_is_in_list", "arguments": {"col_names": ["a", "c"], "allowed": [1, 3, 4]}},
        },
    ]

    checked = dq_engine.apply_checks_by_metadata(test_df, checks)

    expected = spark.createDataFrame(
        [
            [1, 3, 3, None, None],
            [
                2,
                None,
                4,
                {"col_b_is_null_or_empty": "Column b is null or empty"},
                {"col_a_value_is_not_in_the_list": "Value 2 is not in the allowed list: [1, 3, 4]"},
            ],
            [
                None,
                4,
                None,
                None,
                {
                    "col_a_is_null_or_empty": "Column a is null or empty",
                    "col_c_is_null_or_empty": "Column c is null or empty",
                },
            ],
            [
                None,
                None,
                None,
                {"col_b_is_null_or_empty": "Column b is null or empty"},
                {
                    "col_a_is_null_or_empty": "Column a is null or empty",
                    "col_c_is_null_or_empty": "Column c is null or empty",
                },
            ],
        ],
        EXPECTED_SCHEMA,
    )

    assert_df_equality(checked, expected, ignore_nullable=True)


def test_apply_checks_from_json_file_by_metadata(ws, spark):
    dq_engine = DQEngine(ws)
    schema = "col1: int, col2: int, col3: int, col4 int"
    test_df = spark.createDataFrame([[1, 3, 3, 1], [2, None, 4, 1]], schema)

    base_path = str(Path(__file__).resolve().parent.parent)
    checks = DQEngine.load_checks_from_local_file(base_path + "/test_data/checks.json")

    actual = dq_engine.apply_checks_by_metadata(test_df, checks)

    expected = spark.createDataFrame(
        [[1, 3, 3, 1, None, None], [2, None, 4, 1, {"col_col2_is_null": "Column col2 is null"}, None]],
        schema + ", _errors: map<string,string>, _warnings: map<string,string>",
    )

    assert_df_equality(actual, expected, ignore_nullable=True)


def test_apply_checks_from_yml_file_by_metadata(ws, spark):
    dq_engine = DQEngine(ws)
    schema = "col1: int, col2: int, col3: int, col4 int"
    test_df = spark.createDataFrame([[1, 3, 3, 1], [2, None, 4, 1]], schema)

    base_path = str(Path(__file__).resolve().parent.parent)
    checks = DQEngine.load_checks_from_local_file(base_path + "/test_data/checks.yml")

    actual = dq_engine.apply_checks_by_metadata(test_df, checks)

    expected = spark.createDataFrame(
        [[1, 3, 3, 1, None, None], [2, None, 4, 1, {"col_col2_is_null": "Column col2 is null"}, None]],
        schema + ", _errors: map<string,string>, _warnings: map<string,string>",
    )

    assert_df_equality(actual, expected, ignore_nullable=True)


def test_apply_checks_by_metadata_with_func_defined_outside_framework(ws, spark):
    dq_engine = DQEngine(ws)
    test_df = spark.createDataFrame([[1, 3, 3], [2, None, 4], [None, 4, None], [None, None, None]], SCHEMA)

    checks = [{"criticality": "warn", "check": {"function": "col_test_check_func", "arguments": {"col_name": "a"}}}]

    checked = dq_engine.apply_checks_by_metadata(test_df, checks, globals())

    expected = spark.createDataFrame(
        [
            [1, 3, 3, None, None],
            [2, None, 4, None, None],
            [None, 4, None, None, {"col_a_is_null_or_empty": "new check failed"}],
            [None, None, None, None, {"col_a_is_null_or_empty": "new check failed"}],
        ],
        EXPECTED_SCHEMA,
    )

    assert_df_equality(checked, expected, ignore_nullable=True)


def col_test_check_func(col_name: str) -> Column:
    check_col = F.col(col_name)
    condition = check_col.isNull() | (check_col == "") | (check_col == "null")
    return make_condition(condition, "new check failed", f"{col_name}_is_null_or_empty")


def test_get_valid_records(ws, spark):
    dq_engine = DQEngine(ws)

    test_df = spark.createDataFrame(
        [
            [1, 1, 1, None, None],
            [None, 2, 2, None, {"col_a_is_null_or_empty": "check failed"}],
            [None, 2, 2, {"col_b_is_null_or_empty": "check failed"}, None],
        ],
        EXPECTED_SCHEMA,
    )

    valid_df = dq_engine.get_valid(test_df)

    expected_valid_df = spark.createDataFrame(
        [
            [1, 1, 1],
            [None, 2, 2],
        ],
        SCHEMA,
    )

    assert_df_equality(valid_df, expected_valid_df)


def test_get_invalid_records(ws, spark):
    dq_engine = DQEngine(ws)

    test_df = spark.createDataFrame(
        [
            [1, 1, 1, None, None],
            [None, 2, 2, None, {"col_a_is_null_or_empty": "check failed"}],
            [None, 2, 2, {"col_b_is_null_or_empty": "check failed"}, None],
        ],
        EXPECTED_SCHEMA,
    )

    invalid_df = dq_engine.get_invalid(test_df)

    expected_invalid_df = spark.createDataFrame(
        [
            [None, 2, 2, None, {"col_a_is_null_or_empty": "check failed"}],
            [None, 2, 2, {"col_b_is_null_or_empty": "check failed"}, None],
        ],
        EXPECTED_SCHEMA,
    )

    assert_df_equality(invalid_df, expected_invalid_df)


def test_apply_checks_with_custom_column_naming(ws, spark):
    dq_engine = DQEngine(ws, extra_params=ExtraParams(column_names={'errors': 'ERROR', 'warnings': 'WARN'}))
    test_df = spark.createDataFrame([[1, 3, 3], [2, None, 4], [None, 4, None], [None, None, None]], SCHEMA)

    checks = [{"criticality": "warn", "check": {"function": "col_test_check_func", "arguments": {"col_name": "a"}}}]
    checked = dq_engine.apply_checks_by_metadata(test_df, checks, globals())

    assert 'ERROR' in checked.columns
    assert 'WARN' in checked.columns

    expected = spark.createDataFrame(
        [
            [1, 3, 3, None, None],
            [2, None, 4, None, None],
            [None, 4, None, None, {"col_a_is_null_or_empty": "new check failed"}],
            [None, None, None, None, {"col_a_is_null_or_empty": "new check failed"}],
        ],
        EXPECTED_SCHEMA_WITH_CUSTOM_NAMES,
    )

    assert_df_equality(checked, expected, ignore_nullable=True)