Register custom code checks as user defined functions

[marrov]

I just discovered this library and I think it is really amazing!

I'm coming from pandera which is great for pandas dfs but is lacking in support for pyspark which is my main concern. In particular, the fact that you not only capture boolean pass/fail checks but also capture metrics (i.e. as factional pass rate) is a key requirement for my application.

However, the only thing I'm missing is custom checks. I'm aware of the satisfies check that allows for a custom SQL expression, but this is quite limiting for building custom checks. Ideally, I'd like to have the ability to write a function that can be used as a check and register it to add it to the available built-in check, similar to how it is done in pandera.

Would this be a feature that could make it into the library? If so, I'd gladly contribute to this but I have the feeling I'd need the support of the core devs for this.

Let me know your thoughs!

[canimus]

Hi @marrov thanks for opening a discussion, great to have new fresh ideas!
I understand the concept of the custom check, and thanks for sharing the example, what I am curious about is to actually know what is the problem you are trying to solve?
Can you elaborate on an example where a custom check will be required? in that way it may be easier to follow along on the need.
Saludos!

[marrov]

So the background for this request is that I'm working on a metadata-driven data transformation framework for a couple of tens of users that develop data products in pyspark. These data products entail a range of workflows, from simple dashboarding to more complex ML solutions.

For all of these, I am developing a generic data quality solution, for which I originally went with pandera. The requirement is now to also include metrics for the checks and not only boolean pass/fail checks as currently implemented in pandera. On checking the landscape, cuallee seems like an ideal candidate for what I need.

However, another key requirement as I mentioned in my first post is to have user defined quality checks as functions. This would be similar to the generic satisfies method but with more flexibility and also we idiosyncratically want to move away from SQL syntax to enforce pythonic code standards. So, this code-based checks would allow for great flexibility when defining a custom concept of quality for a given dataframe. Let me give an of example (note that the function signature is for demonstration purposes, probably there should be a unifying interface for function io if this custom functionality were to be added to the library):

• Check if a date column of a given frequency has gaps. For example, for a daily frequency check if the diff in days between the sorted dates is different than exactly one. This check is typical for data that is going to be used in time series forecasting. An example function would be:

import pyspark.sql.functions as F
from pyspark.sql import SparkSession
from pyspark.sql.window import Window

spark = SparkSession.builder.appName("date_gaps").getOrCreate()


def has_gaps_daily(df, date_column):
    windowSpec = Window.orderBy(date_column)
    return (
        df.withColumn("day_diff", F.datediff(F.col(date_column), F.lag(F.col(date_column)).over(windowSpec)) - 1)
        .filter(F.col("day_diff") > 0)
        .count()
    )


data = [("2023-01-01",), ("2023-01-02",), ("2023-01-04",), ("2023-01-07",)]  # Note the gaps
columns = ["date"]
df = spark.createDataFrame(data, schema=columns).withColumn("date", F.to_date(F.col("date")))

gap_count = has_gaps_daily(df, "date")
print(f"Number of gaps: {gap_count}")

As far as I can see, something like the above cannot be done with the builtin-checks. Then the SQL-syntax could be an option, but the advantage of this approach is that users can define a unit test for their custom logic so that they can be sure that it is working as expected.

One could also think about more complex checks, like to ensure that the std of a dataset does not vary period by period by more than x%. This would be a way of incorporating data drift checks as quality metrics for something like a gold layer (i.e. aggregated level) data product. Similarly, complex outlier detection would be programatically checked this way with a custom function that could even use an external python library. This is clearly not possible with the current architecture.

I hope this writeup clarifies the ask and the motivation behind it. Let me know if there are any more questions. If this sounds compelling I'd be happy to brainstorm some plugin/extension/registry code concepts.

[canimus]

Thanks for the detailed write up @marrov.

The example you provided is very illustrative. I must say, cuallee has a check that performs exactly that precise use case, of missing dates in a time series, via the is_daily check. You can even customize the day schedule by passing the indices of the dates depending on the dataframe abstraction you are using, meaning it can potentially check lets say [Mon, Wed, Fri] or any combination.

Now back to the original question, I also believe that having the quality of life of linting, checking, unit test, and complexity reporting of code bases in any programming language, are more favorable to plain SQL, so I agree 100% that it could be beneficial to have that feature.
My hesitation of the dependency injection pattern in cuallee is all about the threat of bad-actor, as the Callable that can be used, can contain literally anything.

Wouldn't be an option to compute whatever logic is required, outside in a new column, and then just make the test cases or checks

Other argument of bringing logic to the testing framework, is derived from ISTQB and DAMA-DMBOK in both there are certain patterns that help you preserve quality, one of them is to make the evaluation a white box, and do not embed testing logic, simply elaborate in testing cases. A test case, is a logical condition and an expected result.

From experience, and after automating lots of things for performance and functional testing, is always very relevant to prevent the propagation of change, that is easily said than done, but it fosters the idea that when something changes you don't have to investigate every layer, logic inside of test cases, somehow lead to this phenomenon.

In any case, I am excited to see where this lands, and potentially see a feature that allows or enables such flexibility, and likely as you said, if the use case is moving from other frameworks, ideally for an end user it should be seamless if the features are compatible, like the ones you mention in pandera.

Thinking about how to implement, I think that is relatively simple, checks like:

• has_entropy
• has_workflow
• is_daily
  made use of curried functions that allow internal computation of more elaborated logic and allowing the scenario you are describing, so let's make a next step and see if it can fly!

[marrov]

Well that shows that I haven't gotten up to date with the full API yet, but good catch. The example not being a good one notwithstanding, it nice to see that you also feel there is value in the proposal.

I have one question though: how would the security concern of executing a function with unknown contents be different to SQL injection over the satisfies statement? Maybe there is nuance here that I am missing.

I guess computing a complex column that can be used alongside more simple DQ checks could be an option, but this column would be persisted in an otherwise "clean" table, which feels like something product owners would push against.

Unfortunately, I believe the following may have gone over my head. Could you elaborate a bit on the following?

Other argument of bringing logic to the testing framework, is derived from ISTQB and DAMA-DMBOK in both there are certain patterns that help you preserve quality, one of them is to make the evaluation a white box, and do not embed testing logic, simply elaborate in testing cases. A test case, is a logical condition and an expected result.

From experience, and after automating lots of things for performance and functional testing, is always very relevant to prevent the propagation of change, that is easily said than done, but it fosters the idea that when something changes you don't have to investigate every layer, logic inside of test cases, somehow lead to this phenomenon.

I'll check out the internals of these functions to understand the pattern first. What could be the next steps?

[canimus]

SQL injection and code injection in general is prevented in all frameworks, by passing the code to a syntax tree AST. That will immediately throws an exception when trying to execute something not recognized as part of the grammar. Perhaps pandas is the one with the most flexibility as it offers something like df.query('age.str.contains("import requests")', engine="python") which allows the python engine to expand the sql-like query interface of the query method. Analogously for polars, pyspark, duckdb and snowpark the situation is similar, AST will be created on query executions, when transpiling code into queries and therefore eliminate the threat.

The institution references, are just remarks on how the International Software Testing Qualifications Board (ISTQB) or the Data Management Body of Knowledge (DAMA-DMBOK) will suggest to structure Data Quality. There is large and well structured body of knowledge on the theories behind testing and data management/quality overall, things that help reducing chances of failure. cuallee attempts to adhere to such references, not by being dogmatic, nor situationalist, but striking a balance towards pragmatism.

Last, I was just making a reference that at the point that cuallee gets the handle of a DataFrame, assuming that means that you are in the position to make any transformation into that dataframe as you wish, without having to transfer that logic, to your test cases, the beauty of spark lazy evaluation is that you can literally make all transformations in one-go, relying in the physical structure of the data, because all of them are just logic projections of the data, and therefore live in memory at all time, unless checkpointed or persisted in the middle. So my example was simply evoking to the idea, that pyspark and all its functions will always be more expressive, and more powerful, that any sort of logic embedded in a function, as you showed in your first example, by windowing and lagging dates for potential validations.

I will suggest the following:

• Let's start thinking how the new Check.custom implementation will look like and probably start prototyping something
• From reading again the details, there are certain things that I see as more complex, for instance the use of additional libraries. This probably be the most challenging one, as running in a pyspark cluster, means that additional libraries will have to be available or installed across the cluster nodes

I am happy to continue the discussion over email or a call, you can reach me at [email protected]
Have a great weekend!
Regards, H.

[canimus]

@marrov v0.11.1 has now a is_custom check.
The following test case show how to use it, in general, is just a function that receives a dataframe, and by default it will use the last column of the dataframe to compute the evaluation summary:

def test_positive(spark):
    df = spark.range(10)
    check = Check(CheckLevel.WARNING, "pytest")
    check.is_custom("id", lambda x: x.withColumn("test", F.col("id") >= 0))
    rs = check.validate(df)
    assert rs.first().status == "PASS"
    assert rs.first().violations == 0
    assert rs.first().pass_threshold == 1.0

At the moment is only covered for pyspark at a later stage, we could see if its as straight forward in the other dataframe libraries.
Let us know if this is helpful.
Regards, H.

[marrov]

This is quite amazing @canimus, thanks for the fast implementation. I had some time to play around with this today and it looks like what I was going for! I have a couple of questions but let me introduce an example:

import inspect
import pyspark.sql.functions as F
import pyspark.sql.types as T
from cuallee import Check
from pyspark.sql import DataFrame, SparkSession
from toolz import curry

spark = SparkSession.builder.getOrCreate()
data = [("A", 1), ("B", -1), ("B", 0), ("C", 2)]
schema = T.StructType([T.StructField("id", T.StringType(), True), T.StructField("quantity", T.IntegerType(), True)])
orders = spark.createDataFrame(data, schema=schema)
orders.show()

check = Check(name="orders_checks")
check = check.add_rule("is_unique", "id", 1)
check = check.add_rule("is_greater_than", "quantity", 0, 0.5)


# Define and add a custom check
@curry
def mean_above_threshold(df: DataFrame, column_name: str, threshold: float) -> DataFrame:
    mean_value = df.select(F.mean(column_name).alias("mean")).collect()[0]["mean"]
    is_above_threshold = mean_value > threshold
    return df.withColumn("mean_above_threshold", F.lit(is_above_threshold))

col_name = "quantity"
check = check.add_rule("is_custom", col_name, mean_above_threshold(column_name=col_name, threshold=0), 1)


# Define a custom check function for data type validation
@curry
def is_correct_dtype(df: DataFrame, column_name: str, expected_dtype: T.DataType) -> DataFrame:
    actual_dtype = [field.dataType for field in df.schema.fields if field.name == column_name][0]
    is_dtype_correct = actual_dtype == expected_dtype
    return df.withColumn(f"{column_name}_is_dtype_correct", F.lit(is_dtype_correct))
check = check.add_rule("is_custom", "id", is_correct_dtype(column_name="id", expected_dtype=T.StringType()), 1)
check = check.add_rule("is_custom", "quantity", is_correct_dtype(column_name="quantity", expected_dtype=T.IntegerType()), 1)

# Run the checks
output = check.validate(orders)
output.show()

# Verbose alternative to `f(x)?`
func = check.rules[-1].value
print(f"{func.__name__}{inspect.signature(func)}")

So from this example I have a couple of questions. Some could probably be a different issue, but let's at least pose them here:

1. For this case, the output.show() is as follows:

+---+-------------------+-------------+-------+--------+---------------+-----+----+----------+---------+--------------+------+
| id|          timestamp|        check|  level|  column|           rule|value|rows|violations|pass_rate|pass_threshold|status|
+---+-------------------+-------------+-------+--------+---------------+-----+----+----------+---------+--------------+------+
|  1|2024-07-03 09:54:29|orders_checks|WARNING|      id|      is_unique|  N/A|   4|         1|     0.75|           1.0|  FAIL|
|  2|2024-07-03 09:54:29|orders_checks|WARNING|quantity|is_greater_than|    0|   4|         2|      0.5|           0.5|  PASS|
|  3|2024-07-03 09:54:29|orders_checks|WARNING|quantity|      is_custom| f(x)|   4|         0|      1.0|           1.0|  PASS|
|  4|2024-07-03 09:54:29|orders_checks|WARNING|      id|      is_custom| f(x)|   4|         0|      1.0|           1.0|  PASS|
|  5|2024-07-03 09:54:29|orders_checks|WARNING|quantity|      is_custom| f(x)|   4|         0|      1.0|           1.0|  PASS|
+---+-------------------+-------------+-------+--------+---------------+-----+----+----------+---------+--------------+------+

As you can see when I start adding a couple of custom functions it is very hard to track which is the logic due to value being f(x). I suggest adding the function name or signature although this is not strictly speaking a value. This verbose example could look something like:

is_correct_dtype(df: pyspark.sql.dataframe.DataFrame = '__no__default__', *, column_name: str = 'quantity', expected_dtype: pyspark.sql.types.DataType = IntegerType()) -> pyspark.sql.dataframe.DataFrame

But maybe just the name or at least without the return types. Ideally, there would maybe be something like an optional description column? I'm sure you will have an even better suggestion than these.
2. Why does add_rule support only *args instead of also **kwargs for passing the function arguments programatically? This could become an issue for me due to automated pickup dict params from a yaml file which could in the wrong order order due to user input.
3. In relation to my last two "custom" checks here, I have not seen a way to check dtypes in cuallee. This could be relevant, even in this example where I would like to enforce that quantity is an int, but without the custom check the is_greater_than check passes of course. The result in the case that the input df is has the wrong schema (i.e. quantity is a float I would like to get something like:

+---+-------------------+-------------+-------+--------+---------------+-----+----+----------+---------+--------------+------+
| id|          timestamp|        check|  level|  column|           rule|value|rows|violations|pass_rate|pass_threshold|status|
+---+-------------------+-------------+-------+--------+---------------+-----+----+----------+---------+--------------+------+
|  1|2024-07-03 09:54:29|orders_checks|WARNING|      id|      is_unique|  N/A|   4|         1|     0.75|           1.0|  FAIL|
|  2|2024-07-03 09:54:29|orders_checks|WARNING|quantity|is_greater_than|    0|   4|         2|      0.5|           0.5|  PASS|
|  3|2024-07-03 09:54:29|orders_checks|WARNING|quantity|      is_custom| f(x)|   4|         0|      1.0|           1.0|  PASS|
|  4|2024-07-03 09:54:29|orders_checks|WARNING|      id|      is_custom| f(x)|   4|         0|      1.0|           1.0|  PASS|
|  5|2024-07-03 09:54:29|orders_checks|WARNING|quantity|      is_custom| f(x)|   4|         4|      0.0|           1.0|  FAIL|
+---+-------------------+-------------+-------+--------+---------------+-----+----+----------+---------+--------------+------+

Am I missing something or this cannot be done in this framework?

[canimus]

Thanks @marrov for the detailed example, very illustrative 💯
Here some responses:

1. I think indeed adding a rename parameter as an option to the rule, could be beneficial as that will give the flexibility to give any desired name of the function,
2. This one is an easy one, I added it just now, ran the tests and all looks good, thanks for pointing out, feel free to open PRs if you spot any inconsistency or something that could affect overall functionality.
3. Last but not least, I can write literally a book about the consequences of type-checking and the difficulties to make inferences on types, malformed records, quarantine records, etc. Long story short, it was purposely that schema validation was left aside, as that is a rather complex subject and lots of side effects, so to answer your question, no schema validation or type-checking is available on cuallee at this moment. One reason to give an example, is that you could read data in pyspark, assign a schema to that dataset, and believe that it could enforce it, but it will not, specially when working with text data, not with binary formats like .parquet, .delta or .orc is because of that, that type validation becomes more tricky, because the schema does not guarantee that values in the columns are indeed of such types, and a coerce strategy is required to do so. Hope the long answer gives some context of the reason why.

[canimus]

Hi @marrov

Here a new example with the rename functionality:

import inspect
import pyspark.sql.functions as F
import pyspark.sql.types as T
from cuallee import Check
from pyspark.sql import DataFrame, SparkSession
from toolz import curry

spark = SparkSession.builder.getOrCreate()
data = [("A", 1), ("B", -1), ("B", 0), ("C", 2)]
schema = T.StructType([T.StructField("id", T.StringType(), True), T.StructField("quantity", T.IntegerType(), True)])
orders = spark.createDataFrame(data, schema=schema)
orders.show()

check = Check(name="orders_checks")
check = check.add_rule("is_unique", "id", 1)
check = check.add_rule("is_greater_than", "quantity", 0, 0.5)


# Define and add a custom check
@curry
def mean_above_threshold(df: DataFrame, column_name: str, threshold: float) -> DataFrame:
    mean_value = df.select(F.mean(column_name).alias("mean")).collect()[0]["mean"]
    is_above_threshold = mean_value > threshold
    return df.withColumn("mean_above_threshold", F.lit(is_above_threshold))

col_name = "quantity"
check = check.add_rule("is_custom", col_name, mean_above_threshold(column_name=col_name, threshold=0), 1, options={"name" : "mean_above_threshold", "custom_value": f"{col_name}>0"})


# Define a custom check function for data type validation
@curry
def is_correct_dtype(df: DataFrame, column_name: str, expected_dtype: T.DataType) -> DataFrame:
    actual_dtype = [field.dataType for field in df.schema.fields if field.name == column_name][0]
    is_dtype_correct = actual_dtype == expected_dtype
    return df.withColumn(f"{column_name}_is_dtype_correct", F.lit(is_dtype_correct))
check = check.add_rule("is_custom", "id", is_correct_dtype(column_name="id", expected_dtype=T.StringType()), 1, options={"name" : "is_correct_dtype", "custom_value": "string"})
check = check.add_rule("is_custom", "quantity", is_correct_dtype(column_name="quantity", expected_dtype=T.IntegerType()), 1, options={"name" : "is_correct_dtype", "custom_value": "integer"})

# Run the checks
output = check.validate(orders)
output.show()

# Verbose alternative to `f(x)?`
#func = check.rules[-1].value
#print(f"{func.__name__}{inspect.signature(func)}")

producing the following outcome

+---+-------------------+-------------+-------+--------+--------------------+----------+----+----------+---------+--------------+------+
| id|          timestamp|        check|  level|  column|                rule|     value|rows|violations|pass_rate|pass_threshold|status|
+---+-------------------+-------------+-------+--------+--------------------+----------+----+----------+---------+--------------+------+
|  1|2024-07-13 20:54:43|orders_checks|WARNING|      id|           is_unique|      f(x)|   4|         1|     0.75|           1.0|  FAIL|
|  2|2024-07-13 20:54:43|orders_checks|WARNING|quantity|     is_greater_than|         0|   4|         2|      0.5|           0.5|  PASS|
|  3|2024-07-13 20:54:43|orders_checks|WARNING|quantity|mean_above_threshold|quantity>0|   4|         0|      1.0|           1.0|  PASS|
|  4|2024-07-13 20:54:43|orders_checks|WARNING|      id|    is_correct_dtype|    string|   4|         0|      1.0|           1.0|  PASS|
|  5|2024-07-13 20:54:43|orders_checks|WARNING|quantity|    is_correct_dtype|   integer|   4|         0|      1.0|           1.0|  PASS|
+---+-------------------+-------------+-------+--------+--------------------+----------+----+----------+---------+--------------+------+

[marrov]

@canimus - This looks amazing, exactly what I was looking for, thanks!

One small question: before this update, with built-in checks like is_unique the value defaulted to N/A but now it seems that you have favoured f(x) here. I don't mind either way as it is equally clear to me, but just wanted to hear your thoughts on this if any.

Regarding you point above (no. 3) on type checking, the arguments you raise are quite interesting. I did not anticipate this being that complex and have not encountered before the. From my side my only use case is the definition of an expectation of what the types of the different columns of an unknown data asset that you have not prepare are on read. This would basically work as a simple initial check so that for example ints and dates are the correct types and not a strings as it sometimes happens due to human error/lack of experience.

I understand your not wanting to include data type checking this is the library, but coming from an assessment of the different tools available in the ecosystem as a user (note: user as in "not an expert") the omission is a head scratcher. I am ok with this because, as you have shown above, this is something that can be built by oneself if needed with custom checks. Actually, this is 100% what I will do as type checking is a must in my use case. I was even considering using pandera just for the schema validation as in their library you can turn on/off data and type validation individually. However, building this with custom checks is much cleaner.

As a final question, and excuse me if this is in the docs as I am still in the process of learning the API: is there an easy way to generalise a check (built-in or custom) to all columns? I understand controls plays this role to some extent but not all built-in checks are covered. AFAIK, custom check allows for a single column and multiple columns but not necessarily an all columns at the same time. Is this input str bring validated against the columns? If not, could one just pass a custom all str that is only used in the final report table and as the custom functions gets past the full df just do the check on all columns? Maybe there is something missing here but just wanted to check with you on this.

[canimus]

Hi @marrov well spotted! thanks for pointing out, that shows that we miss some test cases in the validation of the outputs after validation.
Now the correct and backward compatible outputs show correctly:

+---+-------------------+-------------+-------+--------+--------------------+----------+----+----------+---------+--------------+------+
| id|          timestamp|        check|  level|  column|                rule|     value|rows|violations|pass_rate|pass_threshold|status|
+---+-------------------+-------------+-------+--------+--------------------+----------+----+----------+---------+--------------+------+
|  1|2024-07-14 10:36:05|orders_checks|WARNING|      id|           is_unique|       N/A|   4|         1|     0.75|           1.0|  FAIL|
|  2|2024-07-14 10:36:05|orders_checks|WARNING|quantity|     is_greater_than|         0|   4|         2|      0.5|           0.5|  PASS|
|  3|2024-07-14 10:36:05|orders_checks|WARNING|quantity|mean_above_threshold|quantity>0|   4|         0|      1.0|           1.0|  PASS|
|  4|2024-07-14 10:36:05|orders_checks|WARNING|      id|    is_correct_dtype|    string|   4|         0|      1.0|           1.0|  PASS|
|  5|2024-07-14 10:36:05|orders_checks|WARNING|quantity|    is_correct_dtype|   integer|   4|         0|      1.0|           1.0|  PASS|
+---+-------------------+-------------+-------+--------+--------------------+----------+----+----------+---------+--------------+------+

With regards to the application of a check to all columns, apart of the Control.completeness there is no feature that does that, however you are 1-line away with a list comprehension, for example using the definitions above:

check = Check()

# Parameters of the check
options = {"name" : "is_correct_dtype", "custom_value": "integer"}
_fn = lambda x: is_correct_dtype(column_name=x, expected_dtype=T.IntegerType())

# Add same check to all columns of dataframe
[check.is_custom(col, _fn(col), 1, options=options) for col in df.columns] 

This approach is in fact part of the inception of cuallee to be sincere, as I had to apply many checks to dataframes with thousands of columns, and doing it in the citizen way with yaml or configurations was painful, and I think departing from imperatively expressing checks ends up as a limitation factor on how to express and evaluate quality at any scale.

This has been a great discussion @marrov, thanks for your suggestions, if you consider that your initial question has been answered, would you mind marking it as answered just to help future participants that feel motivated to engage in other discussions?

Thanks, and have a great one!