[FEATURE]: is_in_list to validate against a dataframe

[sachinwadhwa]

Is there an existing issue for this?

• I have searched the existing issues

Problem statement

current is_in_list only supports static values, if we can optionally supply dataframe, it will two benefits:

1. the list can be dynamic
2. helps validate foreign keys.

Proposed Solution

as above

Additional Context

No response

[mwojtyczka]

we can add this, also the is_not_null_and_is_in_list will need to be updated. We need to ensure this is performant, not using things like collect().

sample implementation:

def test_referential_integrity(ws, spark):
    from pyspark.sql import functions as F

    data = [(1,), (2,), (3,), (4,), (5,)]
    df = spark.createDataFrame(data, ["col_name: int"])

    allowed_data = [(2,), (3,), (5,)]
    allowed_df = spark.createDataFrame(allowed_data, ["allowed_col: int"])

    def is_in_list(col_name: str | Column, allowed: DataFrame, allowed_col: str) -> Column:
        temp_view = "allowed_temp_view"
        allowed.createOrReplaceTempView(temp_view)
        col_expr_str = get_column_as_string(col_name)
        condition = F.expr(f"EXISTS (SELECT 1 FROM {temp_view} WHERE {allowed_col} = {col_expr_str})")
        return condition

    df = df.withColumn("is_in_allowed_list", is_in_list("col_name", allowed_df, "allowed_col"))
    df.show()

[mwojtyczka]

We have to think how we want to support this. We may need to create new type of check as we need to compare 2 dataframes. Providing dataframe as in my example above will only work for checks defined using classes. For checks in yaml the user won't be able to do it. We could accept a table name but that would require spark session which the current checks don't have access to. Also the core logic should be only dependent on dataframe and not persistence directly like UC table.

[mwojtyczka]

For simple cases you could also use something like below to check referential integrity (in the example i have 2 dimension tables and 1 fact table). You can use it with checks defined in yaml and using classes.

It won't work for big dimension tables as the allowed list will be collected on the driver so you may trip the driver.

dim1_df = spark.read(dim_table1)
pk_dim1_list = [row[0] for row in dim1_df.select("pk_dim1").distinct().collect()]

dim2_df = spark.read(dim_table2)
pk_dim2_list = [row[0] for row in dim2_df.select("pk_dim2").distinct().collect()]

checks = [
   DQColRule(
         check_func=is_in_list,
         col_name="dim1_id",
         check_func_kwargs={"allowed": pk_dim1_list}
      ),
   DQColRule(
         check_func=is_in_list,
         col_name="dim2_id",
         check_func_kwargs={"allowed": pk_dim2_list}
      )
]

fact_df = spark.read(fact_table)
df = dq_engine.apply_checks(fact_df, checks)