Creating Complex Custom Checks

[ahosk]

I'm looking to implement a DQX check that validates a column's value is less than a certain number if another column is equal to a certain value.

I've tried using the sql_expression function as well as creating a custom check, but I've had no luck. Is there a way to do this currently? I saw that some checks allow for a filter, but none of the ones that seem to make sense for this check. I've added in some code below so hopefully i could get pointed in the right direction.

sql_expression:

    function: sql_expression
    arguments:
      expression: col_1 < 1500 and col_2= col_2_value
      msg: col_1 must be less than 1500 for col_2_value
      name: col_1_lt_1500_for_col_2_value 

custom check:

from pyspark.sql import Column
from databricks.labs.dqx.col_functions import make_condition
import yaml

def custom_check_func(col_1_name: str, col_1_value: int, col_2_name: str, col_2_value: str) -> Column:
    col_1_name = F.col(col_1_name)
    col_2_name = F.col(col_2_name)
    return make_condition((col_1_name < col_1_value & col_2_name == col_2_value), f"{col_1_name} is less than {col_1_value} and {col_2_name} is {col_2_value}")


ws = WorkspaceClient()


dq_engine = DQEngine(ws)

check = yaml.safe_load("""
- criticality: warn
  check:
    function: custom_check_func
    arguments:
      col_1_name: foo
      col_2_name: bar
      col_1_value: foo
      col_2_value: bar
""")

valid_df, quarantined_df = dq_engine.apply_checks_by_metadata_and_split(df, checks, globals())

[ahosk]

Update:

Currently implementing this check by the following code and using .unionAll to the quarantined dataframe that is split out. This could probably be more efficient, but this was my first go at separating out this check.

def append_complex_check(df):
    logger.info('Appending complex check')
    invalid_records_filter = ((F.col('foo') == 'string') & (F.col('bar') > val))
    valid_records_filter = ~invalid_records_filter
    
    warning_msg = "foo = 'string' and bar > val"
    
    if df.filter(invalid_records_filter).count() > 0:
        logger.warning('Found quantity invalid records')
        quarantine = df.filter(invalid_records_filter).withColumn(
            "_warnings",
            F.create_map(F.lit("quantity_threshold"), F.lit(warning_msg))
        ).withColumn(
            "_errors",
            F.create_map() 
        )
        valid_df = df.filter(valid_records_filter)
        return valid_df, quarantine
    else:
        return df, None


df, df_custom = append_complex_check(df)
valid_df, quarantined_df = dq_engine.apply_checks_by_metadata_and_split(df, checks)
if df_custom is not None:
    logger.info('Appending complex check to quarantined df')
    quarantined_df = quarantined_df.unionAll(df_custom)
    

[rishujuneja1298]

Can't we write subquery in the sql expression. I am trying to do a reference check with another table.

How to write joins, exists queries?

[mwojtyczka]

I've tried using the sql_expression function as well as creating a custom check, but I've had no luck. Is there a way to do this currently? I saw that some checks allow for a filter, but none of the ones that seem to make sense for this check. I've added in some code below so hopefully i could get pointed in the right direction.

The filter is not limited to certain checks. It works for all checks.

I'm looking to implement a DQX check that validates a column's value is less than a certain number if another column is equal to a certain value.

You can leverage filter and a check not_greater_than for this:

- criticality: error
  filter: col1 < 3
  check:
    function: not_greater_than
    arguments:
      col_name: col2
      limit: 4

[mwojtyczka]

We have also introduced ability to define checks in the complex column types like struct, map, array. Check out the latest release: https://github.com/databrickslabs/dqx/releases/tag/v0.3.0

Docs: https://databrickslabs.github.io/dqx/docs/reference/quality_rules/#usage-examples-of-row-level-checks-for-complex-column-types-struct-map-array

[mwojtyczka]

You can now create custom dataset level checks: https://databrickslabs.github.io/dqx/docs/reference/quality_rules/#creating-custom-dataset-level-checks

[ahosk]

Thank you @mwojtyczka!