Update demo and docs for custom checks (#362)

## Changes
* Improved demo and docs in various places for custom checks

### Tests

- [x] manually tested
- [x] added unit tests
- [ ] added integration tests

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: mwojtyczka

demos/dqx_demo_library.py

@@ -458,9 +458,9 @@
 from pyspark.sql import Column
 from databricks.labs.dqx.check_funcs import make_condition
 
-def ends_with_foo(column: str) -> Column:
+def not_ends_with(column: str, suffix: str) -> Column:
     col_expr = F.col(column)
-    return make_condition(col_expr.endswith("foo"), f"Column {column} ends with foo", f"{column}_ends_with_foo")
+    return make_condition(col_expr.endswith(suffix), f"Column {column} ends with {suffix}", f"{column}_ends_with_{suffix}")
 
 # COMMAND ----------
 
@@ -471,16 +471,20 @@ def ends_with_foo(column: str) -> Column:
 
 from databricks.labs.dqx.engine import DQEngine
 from databricks.sdk import WorkspaceClient
+from databricks.labs.dqx.rule import DQRowRule
 from databricks.labs.dqx.check_funcs import is_not_null_and_not_empty, sql_expression
 
-# use built-in, custom and sql expression checks
+
 checks = [
-    DQRowRule(criticality="error", check_func=is_not_null_and_not_empty, column="col1"),
-    DQRowRule(criticality="warn", check_func=ends_with_foo, column="col1"),
+    # custom check
+    DQRowRule(criticality="warn", check_func=not_ends_with, column="col1", check_func_kwargs={"suffix": "foo"}),
+    # sql expression check
     DQRowRule(criticality="warn", check_func=sql_expression, check_func_kwargs={
             "expression": "col1 like 'str%'", "msg": "col1 not starting with 'str'"
         }
     ),
+    # built-in check
+    DQRowRule(criticality="error", check_func=is_not_null_and_not_empty, column="col1"),
 ]
 
 schema = "col1: string, col2: string"
@@ -502,25 +506,26 @@ def ends_with_foo(column: str) -> Column:
 from databricks.labs.dqx.engine import DQEngine
 from databricks.sdk import WorkspaceClient
 
-# use built-in, custom and sql expression checks
+
 checks = yaml.safe_load(
 """
-- criticality: error
-  check:
-    function: is_not_null_and_not_empty
-    arguments:
-      column: col1
 - criticality: warn
   check:
-    function: ends_with_foo
+    function: not_ends_with
     arguments:
       column: col1
+      suffix: foo
 - criticality: warn
   check:
     function: sql_expression
     arguments:
       expression: col1 like 'str%'
       msg: col1 not starting with 'str'
+- criticality: error
+  check:
+    function: is_not_null_and_not_empty
+    arguments:
+      column: col1
 """
 )
 
@@ -529,7 +534,7 @@ def ends_with_foo(column: str) -> Column:
 
 dq_engine = DQEngine(WorkspaceClient())
 
-custom_check_functions = {"ends_with_foo": ends_with_foo}
+custom_check_functions = {"not_ends_with": not_ends_with}
 # alternatively, you can also use globals to include all available functions
 #custom_check_functions = globals()
 

docs/dqx/docs/reference/quality_rules.mdx

@@ -10,7 +10,7 @@ This page provides a reference for the quality checks (rule functions) available
 
 The following built-in row-level checks are currently available in DQX.
 These checks are applied to each row of a PySpark DataFrame and report quality issues as additional columns.
-You can also define your own custom checks (see [Creating custom user-defined checks](#creating-custom-user-defined-checks)).
+You can also define your own custom checks (see [Creating custom checks](#creating-custom-checks)).
 
 Some of the checks such as `is_unique`, `is_aggr_not_greater_than`, or `is_aggr_not_less_than` are performed over a group of rows using window functions, while others operate on individual rows.
 In all cases, the quality check results are always reported for individual rows in the reporting columns.
@@ -1021,7 +1021,7 @@ For example, to check that a column `a` is not null only when a column `b` is po
       column: col1
 ```
 
-## Creating Custom User-Defined Checks
+## Creating Custom Checks
 
 ### Using SQL Expression
 
@@ -1048,7 +1048,7 @@ SQL Expressions are also useful if you need to make cross-column validation, for
 
 ### Using Python function
 
-If you need a reusable check or want to implement more complex logic which is challenging to implement using SQL expression, you can define your own custom check function.
+If you need a reusable check or want to implement more complex logic which is challenging to implement using SQL expression, you can define your own custom check function in Python.
 A check function is a callable that uses `make_condition` to return `pyspark.sql.Column`.
 
 #### Custom check example
@@ -1060,9 +1060,10 @@ A check function is a callable that uses `make_condition` to return `pyspark.sql
     from pyspark.sql import Column
     from databricks.labs.dqx.check_funcs import make_condition
 
-    def ends_with_foo(column: str) -> Column:
+    def not_ends_with(column: str, suffix: str) -> Column:
       col_expr = F.col(column)
-      return make_condition(col_expr.endswith("foo"), f"Column {column} ends with foo", f"{column}_ends_with_foo")
+      return make_condition(col_expr.endswith(suffix), f"Column {column} ends with {suffix}", f"{column}_ends_with_{suffix}")
+
     ```
   </TabItem>
 </Tabs>
@@ -1079,8 +1080,10 @@ A check function is a callable that uses `make_condition` to return `pyspark.sql
     from databricks.labs.dqx.check_funcs import is_not_null
 
     checks = [
+      # custom check
+      DQRowRule(criticality="error", check_func=not_ends_with, column="col1", check_func_kwargs={"suffix": "foo"}),
+      # built-in check
       DQRowRule(criticality="error", check_func=is_not_null, column="col1"),
-      DQRowRule(criticality="error", check_func=ends_with_foo, column="col1"),
     ]
 
     dq_engine = DQEngine(WorkspaceClient())
@@ -1089,7 +1092,7 @@ A check function is a callable that uses `make_condition` to return `pyspark.sql
     valid_df, quarantine_df = dq_engine.apply_checks_and_split(input_df, checks)
 
     # Option 2: apply quality rules on the dataframe and report issues as additional columns
-    valid_and_quarantine_df = dq_engine.apply_checks_by_metadata(input_df, checks)
+    valid_and_quarantine_df = dq_engine.apply_checks(input_df, checks)
     ```
   </TabItem>
 </Tabs>
@@ -1106,9 +1109,10 @@ A check function is a callable that uses `make_condition` to return `pyspark.sql
     checks = yaml.safe_load("""
     - criticality: error
       check:
-        function: ends_with_foo
+        function: not_ends_with
         arguments:
           column: col1
+          suffix: foo
     - criticality: error
       check:
         function: is_not_null
@@ -1118,7 +1122,7 @@ A check function is a callable that uses `make_condition` to return `pyspark.sql
 
     dq_engine = DQEngine(WorkspaceClient())
 
-    custom_check_functions = {"ends_with_foo": ends_with_foo}  # list of custom check functions
+    custom_check_functions = {"not_ends_with": not_ends_with}  # list of custom check functions
     # or include all functions with globals() for simplicity
     #custom_check_functions=globals()
 

src/databricks/labs/dqx/engine.py

@@ -86,15 +86,15 @@ def apply_checks_and_split(self, df: DataFrame, checks: list[DQRule]) -> tuple[D
     def apply_checks_by_metadata_and_split(
         self, df: DataFrame, checks: list[dict], custom_check_functions: dict[str, Any] | None = None
     ) -> tuple[DataFrame, DataFrame]:
-        dq_rule_checks = self.build_checks_by_metadata(checks, custom_check_functions)
+        dq_rule_checks = self.build_quality_rules_by_metadata(checks, custom_check_functions)
 
         good_df, bad_df = self.apply_checks_and_split(df, dq_rule_checks)
         return good_df, bad_df
 
     def apply_checks_by_metadata(
         self, df: DataFrame, checks: list[dict], custom_check_functions: dict[str, Any] | None = None
     ) -> DataFrame:
-        dq_rule_checks = self.build_checks_by_metadata(checks, custom_check_functions)
+        dq_rule_checks = self.build_quality_rules_by_metadata(checks, custom_check_functions)
 
         return self.apply_checks(df, dq_rule_checks)
 
@@ -219,7 +219,7 @@ def build_dataframe_from_quality_rules(
         if spark is None:
             spark = SparkSession.builder.getOrCreate()
 
-        dq_rule_checks = DQEngineCore.build_checks_by_metadata(checks)
+        dq_rule_checks = DQEngineCore.build_quality_rules_by_metadata(checks)
 
         dq_rule_rows = []
         for dq_rule_check in dq_rule_checks:
@@ -246,7 +246,9 @@ def build_dataframe_from_quality_rules(
         return spark.createDataFrame(dq_rule_rows, DQEngineCore.CHECKS_TABLE_SCHEMA)
 
     @staticmethod
-    def build_checks_by_metadata(checks: list[dict], custom_checks: dict[str, Any] | None = None) -> list[DQRule]:
+    def build_quality_rules_by_metadata(
+        checks: list[dict], custom_checks: dict[str, Any] | None = None
+    ) -> list[DQRule]:
         """Build checks based on check specification, i.e. function name plus arguments.
 
         :param checks: list of dictionaries describing checks. Each check is a dictionary
@@ -327,9 +329,9 @@ def build_checks_by_metadata(checks: list[dict], custom_checks: dict[str, Any] |
         return dq_rule_checks
 
     @staticmethod
-    def build_checks(*rules_col_set: DQRowRuleForEachCol) -> list[DQRule]:
+    def build_quality_rules_foreach_col(*rules_col_set: DQRowRuleForEachCol) -> list[DQRule]:
         """
-        Build rules from dq rules and rule sets.
+        Build rules for each column from DQRowRuleForEachCol sets.
 
         :param rules_col_set: list of dq rules which define multiple columns for the same check function
         :return: list of dq rules

tests/integration/test_apply_checks.py

@@ -1257,6 +1257,13 @@ def custom_check_func_global(column: str) -> Column:
     return check_funcs.make_condition(col_expr.isNull(), "custom check failed", f"{column}_is_null_custom")
 
 
+def custom_check_func_custom_args(column_custom_arg: str) -> Column:
+    col_expr = F.col(column_custom_arg)
+    return check_funcs.make_condition(
+        col_expr.isNull(), "custom check with custom args failed", f"{column_custom_arg}_is_null_custom_args"
+    )
+
+
 def custom_check_func_global_a_column_no_args() -> Column:
     col_expr = F.col("a")
     return check_funcs.make_condition(col_expr.isNull(), "custom check without args failed", "a_is_null_custom")
@@ -1275,8 +1282,13 @@ def test_apply_checks_with_custom_check(ws, spark):
     checks = [
         DQRowRule(criticality="warn", check_func=check_funcs.is_not_null_and_not_empty, column="a"),
         DQRowRule(criticality="warn", check_func=custom_check_func_global, column="a"),
+        DQRowRule(criticality="warn", check_func=custom_check_func_global, check_func_kwargs={"column": "a"}),
         DQRowRule(criticality="warn", check_func=custom_check_func_global_annotated, column="a"),
+        DQRowRule(criticality="warn", check_func=custom_check_func_global_annotated, check_func_kwargs={"column": "a"}),
         DQRowRule(criticality="warn", check_func=custom_check_func_global_a_column_no_args),
+        DQRowRule(
+            criticality="warn", check_func=custom_check_func_custom_args, check_func_kwargs={"column_custom_arg": "a"}
+        ),
     ]
 
     checked = dq_engine.apply_checks(test_df, checks)
@@ -1309,6 +1321,15 @@ def test_apply_checks_with_custom_check(ws, spark):
                         "run_time": RUN_TIME,
                         "user_metadata": {},
                     },
+                    {
+                        "name": "a_is_null_custom",
+                        "message": "custom check failed",
+                        "columns": None,
+                        "filter": None,
+                        "function": "custom_check_func_global",
+                        "run_time": RUN_TIME,
+                        "user_metadata": {},
+                    },
                     {
                         "name": "a_is_null_custom",
                         "message": "custom check annotated failed",
@@ -1318,6 +1339,15 @@ def test_apply_checks_with_custom_check(ws, spark):
                         "run_time": RUN_TIME,
                         "user_metadata": {},
                     },
+                    {
+                        "name": "a_is_null_custom",
+                        "message": "custom check annotated failed",
+                        "columns": None,
+                        "filter": None,
+                        "function": "custom_check_func_global_annotated",
+                        "run_time": RUN_TIME,
+                        "user_metadata": {},
+                    },
                     {
                         "name": "a_is_null_custom",
                         "message": "custom check without args failed",
@@ -1327,6 +1357,15 @@ def test_apply_checks_with_custom_check(ws, spark):
                         "run_time": RUN_TIME,
                         "user_metadata": {},
                     },
+                    {
+                        "name": "a_is_null_custom_args",
+                        "message": "custom check with custom args failed",
+                        "columns": None,
+                        "filter": None,
+                        "function": "custom_check_func_custom_args",
+                        "run_time": RUN_TIME,
+                        "user_metadata": {},
+                    },
                 ],
             ],
             [
@@ -1353,6 +1392,15 @@ def test_apply_checks_with_custom_check(ws, spark):
                         "run_time": RUN_TIME,
                         "user_metadata": {},
                     },
+                    {
+                        "name": "a_is_null_custom",
+                        "message": "custom check failed",
+                        "columns": None,
+                        "filter": None,
+                        "function": "custom_check_func_global",
+                        "run_time": RUN_TIME,
+                        "user_metadata": {},
+                    },
                     {
                         "name": "a_is_null_custom",
                         "message": "custom check annotated failed",
@@ -1362,6 +1410,15 @@ def test_apply_checks_with_custom_check(ws, spark):
                         "run_time": RUN_TIME,
                         "user_metadata": {},
                     },
+                    {
+                        "name": "a_is_null_custom",
+                        "message": "custom check annotated failed",
+                        "columns": None,
+                        "filter": None,
+                        "function": "custom_check_func_global_annotated",
+                        "run_time": RUN_TIME,
+                        "user_metadata": {},
+                    },
                     {
                         "name": "a_is_null_custom",
                         "message": "custom check without args failed",
@@ -1371,6 +1428,15 @@ def test_apply_checks_with_custom_check(ws, spark):
                         "run_time": RUN_TIME,
                         "user_metadata": {},
                     },
+                    {
+                        "name": "a_is_null_custom_args",
+                        "message": "custom check with custom args failed",
+                        "columns": None,
+                        "filter": None,
+                        "function": "custom_check_func_custom_args",
+                        "run_time": RUN_TIME,
+                        "user_metadata": {},
+                    },
                 ],
             ],
         ],
@@ -1391,6 +1457,10 @@ def test_apply_checks_by_metadata_with_custom_check(ws, spark):
             "criticality": "warn",
             "check": {"function": "custom_check_func_global_annotated", "arguments": {"column": "a"}},
         },
+        {
+            "criticality": "warn",
+            "check": {"function": "custom_check_func_custom_args", "arguments": {"column_custom_arg": "a"}},
+        },
     ]
 
     checked = dq_engine.apply_checks_by_metadata(
@@ -1399,6 +1469,7 @@ def test_apply_checks_by_metadata_with_custom_check(ws, spark):
         {
             "custom_check_func_global": custom_check_func_global,
             "custom_check_func_global_annotated": custom_check_func_global_annotated,
+            "custom_check_func_custom_args": custom_check_func_custom_args,
         },
     )
     # or for simplicity use globals
@@ -1441,6 +1512,15 @@ def test_apply_checks_by_metadata_with_custom_check(ws, spark):
                         "run_time": RUN_TIME,
                         "user_metadata": {},
                     },
+                    {
+                        "name": "a_is_null_custom_args",
+                        "message": "custom check with custom args failed",
+                        "columns": None,
+                        "filter": None,
+                        "function": "custom_check_func_custom_args",
+                        "run_time": RUN_TIME,
+                        "user_metadata": {},
+                    },
                 ],
             ],
             [
@@ -1476,6 +1556,15 @@ def test_apply_checks_by_metadata_with_custom_check(ws, spark):
                         "run_time": RUN_TIME,
                         "user_metadata": {},
                     },
+                    {
+                        "name": "a_is_null_custom_args",
+                        "message": "custom check with custom args failed",
+                        "columns": None,
+                        "filter": None,
+                        "function": "custom_check_func_custom_args",
+                        "run_time": RUN_TIME,
+                        "user_metadata": {},
+                    },
                 ],
             ],
         ],