from pydeequ.checks import *
from pydeequ.verification import *

# примерные данные
df = spark.sparkContext.parallelize([
    Row(a="foo", b=1, c=5),
    Row(a="bar", b=2, c=6),
    Row(a="baz", b=3, c=None)]).toDF()

# инициализация теста
check = Check(spark, CheckLevel.Error, "Integrity checks")

# rdd обзор / тест данных
checkResult = VerificationSuite(spark) \
    .onData(df) \
    .addCheck(
        check.hasSize(lambda x: x >= 3) \
        .hasMin("b", lambda x: x == 0) \
        .isComplete("c")  \
        .isUnique("a")  \
        .isContainedIn("a", ["foo", "bar", "baz"]) \
        .isNonNegative("b")) \
    .run()

# запуск верификации
checkResult_df = VerificationResult.checkResultsAsDataFrame(spark, checkResult)
checkResult_df.show()

# описание результата
if checkResult.status == "Success":
    print('Тесты пройдены')

else:
    print('Найдены ошибки:')

    for check_json in checkResult.checkResults:
        if check_json['constraint_status'] != "Success":
            print(f"\t{check_json['constraint']} причина: {check_json['constraint_message']}")