Merge branch 'main' into integral_type_conversion

Author: mwojtyczka

demos/dqx_demo_library.py

@@ -75,7 +75,7 @@
 from databricks.labs.dqx.engine import DQEngine
 from databricks.sdk import WorkspaceClient
 
-input_df = spark.createDataFrame([[1, 3, 3, 2], [2, 3, None, 1]], schema)
+input_df = spark.createDataFrame([[1, 3, 3, 2], [3, 3, None, 1]], schema)
 
 # load checks
 dq_engine = DQEngine(WorkspaceClient())
@@ -263,7 +263,7 @@
 - check:
     function: sql_expression
     arguments:
-      expression: pickup_datetime > dropoff_datetime
+      expression: pickup_datetime <= dropoff_datetime
       msg: pickup time must not be greater than dropff time
       name: pickup_datetime_greater_than_dropoff_datetime
   criticality: error
@@ -321,7 +321,33 @@ def ends_with_foo(col_name: str) -> Column:
 # COMMAND ----------
 
 # MAGIC %md
-# MAGIC ### Applying custom check function
+# MAGIC ### Applying custom check function using DQX classes
+
+# COMMAND ----------
+
+from databricks.labs.dqx.engine import DQEngine
+from databricks.sdk import WorkspaceClient
+from databricks.labs.dqx.col_functions import *
+
+# use built-in, custom and sql expression checks
+checks = [
+    DQRule(criticality="error", check=is_not_null_and_not_empty("col1")),
+    DQRule(criticality="warn", check=ends_with_foo("col1")),
+    DQRule(criticality="warn", check=sql_expression("col1 not like 'str%'", msg="col1 starts with 'str'")),
+]
+
+schema = "col1: string, col2: string"
+input_df = spark.createDataFrame([[None, "foo"], ["foo", None], [None, None]], schema)
+
+dq_engine = DQEngine(WorkspaceClient())
+
+valid_and_quarantined_df = dq_engine.apply_checks(input_df, checks)
+display(valid_and_quarantined_df)
+
+# COMMAND ----------
+
+# MAGIC %md
+# MAGIC ### Applying custom check function using YAML definition
 
 # COMMAND ----------
 
@@ -347,7 +373,7 @@ def ends_with_foo(col_name: str) -> Column:
   check:
     function: sql_expression
     arguments:
-      expression: col1 LIKE 'str%'
+      expression: col1 not like 'str%'
       msg: col1 starts with 'str'
 """
 )
@@ -358,8 +384,11 @@ def ends_with_foo(col_name: str) -> Column:
 dq_engine = DQEngine(WorkspaceClient())
 
 custom_check_functions = {"ends_with_foo": ends_with_foo}
-# or include all functions with globals() for simplicity
-#custom_check_functions=globals()
+# alternatively, you can also use globals to include all available functions
+#custom_check_functions = globals()
+
+status = dq_engine.validate_checks(checks, custom_check_functions)
+assert not status.has_errors
 
 valid_and_quarantined_df = dq_engine.apply_checks_by_metadata(input_df, checks, custom_check_functions)
 display(valid_and_quarantined_df)

demos/dqx_demo_tool.py

@@ -143,7 +143,7 @@
 - check:
     function: sql_expression
     arguments:
-      expression: pickup_datetime > dropoff_datetime
+      expression: pickup_datetime <= dropoff_datetime
       msg: pickup time must not be greater than dropff time
       name: pickup_datetime_greater_than_dropoff_datetime
   criticality: error

docs/dqx/docs/reference/quality_rules.mdx

@@ -130,14 +130,15 @@ def sql_expression(expression: str, msg: str | None = None, name: str | None = N
     expression_msg = expression
 
     if negate:
-        expr_col = ~expr_col
         expression_msg = "~(" + expression + ")"
+        message = F.concat_ws("", F.lit(f"Value is matching expression: {expression_msg}"))
+    else:
+        expr_col = ~expr_col
+        message = F.concat_ws("", F.lit(f"Value is not matching expression: {expression_msg}"))
 
     name = name if name else re.sub(normalize_regex, "_", expression)
 
-    if msg:
-        return make_condition(expr_col, msg, name)
-    return make_condition(expr_col, F.concat_ws("", F.lit(f"Value matches expression: {expression_msg}")), name)
+    return make_condition(expr_col, msg or message, name)
 
 
 def is_older_than_col2_for_n_days(col_name1: str, col_name2: str, days: int = 0) -> Column:
@@ -256,7 +257,7 @@ def is_not_less_than(
     """Checks whether the values in the input column are not less than the provided limit.
 
     :param col_name: column name
-    :param limit: limit to use in the condition as number, date, timestamp, column name or expression
+    :param limit: limit to use in the condition as number, date, timestamp, column name or sql expression
     :return: new Column
     """
     limit_expr = _get_column_expr_limit(limit)
@@ -275,7 +276,7 @@ def is_not_greater_than(
     """Checks whether the values in the input column are not greater than the provided limit.
 
     :param col_name: column name
-    :param limit: limit to use in the condition as number, date, timestamp, column name or expression
+    :param limit: limit to use in the condition as number, date, timestamp, column name or sql expression
     :return: new Column
     """
     limit_expr = _get_column_expr_limit(limit)
@@ -296,8 +297,8 @@ def is_in_range(
     """Checks whether the values in the input column are in the provided limits (inclusive of both boundaries).
 
     :param col_name: column name
-    :param min_limit: min limit to use in the condition as number, date, timestamp, column name or expression
-    :param max_limit: max limit to use in the condition as number, date, timestamp, column name or expression
+    :param min_limit: min limit to use in the condition as number, date, timestamp, column name or sql expression
+    :param max_limit: max limit to use in the condition as number, date, timestamp, column name or sql expression
     :return: new Column
     """
     min_limit_expr = _get_column_expr_limit(min_limit)
@@ -329,8 +330,8 @@ def is_not_in_range(
     """Checks whether the values in the input column are outside the provided limits (inclusive of both boundaries).
 
     :param col_name: column name
-    :param min_limit: min limit to use in the condition as number, date, timestamp, column name or expression
-    :param max_limit: min limit to use in the condition as number, date, timestamp, column name or expression
+    :param min_limit: min limit to use in the condition as number, date, timestamp, column name or sql expression
+    :param max_limit: min limit to use in the condition as number, date, timestamp, column name or sql expression
     :return: new Column
     """
     min_limit_expr = _get_column_expr_limit(min_limit)
@@ -364,11 +365,9 @@ def regex_match(col_name: str, regex: str, negate: bool = False) -> Column:
     """
     if negate:
         condition = F.col(col_name).rlike(regex)
-
         return make_condition(condition, f"Column {col_name} is matching regex", f"{col_name}_matching_regex")
 
     condition = ~F.col(col_name).rlike(regex)
-
     return make_condition(condition, f"Column {col_name} is not matching regex", f"{col_name}_not_matching_regex")
 
 
@@ -460,7 +459,7 @@ def _get_column_expr_limit(
 ) -> Column:
     """Helper function to generate a column expression limit based on the provided limit value.
 
-    :param limit: limit to use in the condition (literal value or Column expression)
+    :param limit: limit to use in the condition (literal value or column expression)
     :return: column expression.
     :raises ValueError: if limit is not provided.
     """

src/databricks/labs/dqx/col_functions.py

@@ -32,12 +32,12 @@ def checks_yml_content():
   check:
     function: sql_expression
     arguments:
-      expression: col1 like "Team %"
+      expression: col1 not like "Team %"
 - criticality: error
   check:
     function: sql_expression
     arguments:
-      expression: col2 like 'Team %'
+      expression: col2 not like 'Team %'
     """
 
 
@@ -76,11 +76,11 @@ def checks_json_content():
     },
     {
         "criticality": "error",
-        "check": {"function": "sql_expression", "arguments": {"expression": "col1 like \\"Team %\\""}}
+        "check": {"function": "sql_expression", "arguments": {"expression": "col1 not like \\"Team %\\""}}
     },
     {
         "criticality": "error",
-        "check": {"function": "sql_expression", "arguments": {"expression": "col2 like 'Team %'"}}
+        "check": {"function": "sql_expression", "arguments": {"expression": "col2 not like 'Team %'"}}
     }
 ]
     """
@@ -134,11 +134,11 @@ def expected_checks():
         },
         {
             "criticality": "error",
-            "check": {"function": "sql_expression", "arguments": {"expression": "col1 like \"Team %\""}},
+            "check": {"function": "sql_expression", "arguments": {"expression": "col1 not like \"Team %\""}},
         },
         {
             "criticality": "error",
-            "check": {"function": "sql_expression", "arguments": {"expression": "col2 like 'Team %'"}},
+            "check": {"function": "sql_expression", "arguments": {"expression": "col2 not like 'Team %'"}},
         },
     ]
 

tests/conftest.py

@@ -1,5 +1,4 @@
 from datetime import datetime
-
 import yaml
 import pyspark.sql.functions as F
 import pytest
@@ -747,11 +746,11 @@ def test_apply_checks_with_sql_expression(ws, spark):
     checks = [
         {
             "criticality": "error",
-            "check": {"function": "sql_expression", "arguments": {"expression": "col1 like \"val%\""}},
+            "check": {"function": "sql_expression", "arguments": {"expression": "col1 not like \"val%\""}},
         },
         {
             "criticality": "error",
-            "check": {"function": "sql_expression", "arguments": {"expression": "col2 like 'val%'"}},
+            "check": {"function": "sql_expression", "arguments": {"expression": "col2 not like 'val%'"}},
         },
     ]
 
@@ -765,8 +764,8 @@ def test_apply_checks_with_sql_expression(ws, spark):
                 "val1",
                 "val2",
                 {
-                    "col_col1_like_val_": "Value matches expression: col1 like \"val%\"",
-                    "col_col2_like_val_": "Value matches expression: col2 like 'val%'",
+                    "col_col1_not_like_val_": "Value is not matching expression: col1 not like \"val%\"",
+                    "col_col2_not_like_val_": "Value is not matching expression: col2 not like 'val%'",
                 },
                 None,
             ],
@@ -1081,8 +1080,8 @@ def test_apply_checks_all_checks_as_yaml(ws, spark):
       check:
         function: sql_expression
         arguments:
-          expression: col3 > col2 and col3 < 10
-          msg: col3 is greater than col2 and col3 less than 10
+          expression: col3 >= col2 and col3 <= 10
+          msg: ol3 is less than col2 and col3 is greater than 10
           name: custom_output_name
           negate: false
     """
@@ -1186,8 +1185,8 @@ def test_apply_checks_all_checks_using_classes(ws, spark):
         DQRule(
             criticality="error",
             check=sql_expression(
-                expression="col3 > col2 and col3 < 10",
-                msg="col3 is greater than col2 and col3 less than 10",
+                expression="col3 >= col2 and col3 <= 10",
+                msg="col3 is less than col2 and col3 is greater than 10",
                 name="custom_output_name",
                 negate=False,
             ),

tests/integration/test_apply_checks.py

@@ -2,8 +2,6 @@
 from decimal import Decimal
 import pyspark.sql.functions as F
 from chispa.dataframe_comparer import assert_df_equality  # type: ignore
-
-
 from databricks.labs.dqx.col_functions import (
     is_in_range,
     is_not_empty,
@@ -104,21 +102,21 @@ def test_col_is_not_in_list(spark):
 
 
 def test_col_sql_expression(spark):
-    test_df = spark.createDataFrame([["str1", 1, 1], ["str2", None, None], ["", 3, 2]], SCHEMA + ", c: string")
+    test_df = spark.createDataFrame([["str1", 1, 1], ["str2", None, None], ["", 2, 3]], SCHEMA + ", c: string")
 
     actual = test_df.select(
-        sql_expression("a = 'str1'"),
-        sql_expression("b is not null", name="test", negate=True),
-        sql_expression("c is not null", msg="failed validation", negate=True),
-        sql_expression("b > c", msg="b is not greater than c", negate=True),
+        sql_expression("a = 'str2'"),
+        sql_expression("b is null", name="test", negate=True),
+        sql_expression("c is null", msg="failed validation", negate=True),
+        sql_expression("b < c", msg="b is greater or equal c", negate=False),
     )
 
-    checked_schema = "a_str1_: string, test: string, c_is_not_null: string, b_c: string"
+    checked_schema = "a_str2_: string, test: string, c_is_null: string, b_c: string"
     expected = spark.createDataFrame(
         [
-            ["Value matches expression: a = 'str1'", None, None, "b is not greater than c"],
-            [None, "Value matches expression: ~(b is not null)", "failed validation", None],
-            [None, None, None, None],
+            ["Value is not matching expression: a = 'str2'", None, None, "b is greater or equal c"],
+            [None, "Value is matching expression: ~(b is null)", "failed validation", None],
+            ["Value is not matching expression: a = 'str2'", None, None, None],
         ],
         checked_schema,
     )
@@ -497,7 +495,10 @@ def test_col_is_not_in_near_future_cur(spark):
 
 
 def test_col_is_not_null_and_not_empty_array(spark):
-    schema_array = "str_col: array<string>, int_col: array<int> , timestamp_col: array<timestamp>, date_col: array<string>, struct_col: array<struct<a: string, b: int>>"
+    schema_array = (
+        "str_col: array<string>, int_col: array<int> , timestamp_col: array<timestamp>, "
+        "date_col: array<string>, struct_col: array<struct<a: string, b: int>>"
+    )
     data = [
         (
             ["a", "b", None],
@@ -527,7 +528,11 @@ def test_col_is_not_null_and_not_empty_array(spark):
         is_not_null_and_not_empty_array("struct_col"),
     )
 
-    checked_schema = "str_col_is_null_or_empty_array: string, int_col_is_null_or_empty_array: string, timestamp_col_is_null_or_empty_array: string, date_col_is_null_or_empty_array: string, struct_col_is_null_or_empty_array: string"
+    checked_schema = (
+        "str_col_is_null_or_empty_array: string, int_col_is_null_or_empty_array: string, "
+        "timestamp_col_is_null_or_empty_array: string, date_col_is_null_or_empty_array: string, "
+        "struct_col_is_null_or_empty_array: string"
+    )
     # Create the data
     checked_data = [
         (None, None, None, None, None),