Updated Readme and improved support for DLT

pyproject.toml

@@ -14,7 +14,7 @@ classifiers = [
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: Implementation :: CPython",
 ]
-dependencies = ["databricks-labs-blueprint[yaml]~=0.4.4"]
+dependencies = ["databricks-labs-blueprint[yaml]>=0.8,<0.9"]
 
 [project.urls]
 Issues = "https://github.com/databrickslabs/dqx/issues"
@@ -74,6 +74,9 @@ verify      = ["black --check .",
 [tool.isort]
 profile = "black"
 
+[tool.mypy]
+exclude = ['venv', '.venv']
+
 [tool.pytest.ini_options]
 addopts = "--no-header"
 cache_dir = ".venv/pytest-cache"

src/databricks/labs/dqx/col_functions.py

@@ -21,7 +21,7 @@ def make_condition(condition: Column, message: Column | str, alias: str) -> Colu
     return (F.when(condition, msg_col).otherwise(F.lit(None).cast("string"))).alias(_cleanup_alias_name(alias))
 
 
-def _cleanup_alias_name(col_name: str):
+def _cleanup_alias_name(col_name: str) -> str:
     # avoid issues with structs
     return col_name.replace(".", "_")
 

src/databricks/labs/dqx/engine.py

@@ -169,12 +169,30 @@ def apply_checks_and_split(df: DataFrame, checks: list[DQRule]) -> tuple[DataFra
 
     checked_df = apply_checks(df, checks)
 
-    good_df = checked_df.where(F.col(Columns.ERRORS.value).isNull()).drop(Columns.ERRORS.value, Columns.WARNINGS.value)
-    bad_df = checked_df.where(F.col(Columns.ERRORS.value).isNotNull() | F.col(Columns.WARNINGS.value).isNotNull())
+    good_df = get_valid(checked_df)
+    bad_df = get_invalid(checked_df)
 
     return good_df, bad_df
 
 
+def get_invalid(df: DataFrame) -> DataFrame:
+    """
+    Get invalid records only (errors and warnings).
+    @param df: input DataFrame
+    @return:
+    """
+    return df.where(F.col(Columns.ERRORS.value).isNotNull() | F.col(Columns.WARNINGS.value).isNotNull())
+
+
+def get_valid(df: DataFrame) -> DataFrame:
+    """
+    Get valid records only (errors only)
+    @param df: input DataFrame.
+    @return:
+    """
+    return df.where(F.col(Columns.ERRORS.value).isNull()).drop(Columns.ERRORS.value, Columns.WARNINGS.value)
+
+
 def build_checks_by_metadata(checks: list[dict], glbs: dict[str, Any] | None = None) -> list[DQRule]:
     """Build checks based on check specification, i.e. function name plus arguments.
 

src/databricks/labs/dqx/profiler/dlt_generator.py

@@ -2,7 +2,7 @@
 import re
 
 from databricks.labs.dqx.profiler.common import val_to_str
-from databricks.labs.dqx.profiler.profiler import DQProfile
+from databricks.labs.dqx.profiler.engine import DQProfile
 
 __name_sanitize_re__ = re.compile(r"[^a-zA-Z0-9]+")
 

src/databricks/labs/dqx/profiler/profiler.py → src/databricks/labs/dqx/profiler/engine.py

@@ -32,7 +32,7 @@ def do_cast(value: str | None, typ: T.DataType) -> Any | None:
 def get_df_summary_as_dict(df: DataFrame) -> dict[str, Any]:
     """Generate summary for Dataframe & return it as dictionary with column name as a key, and dict of metric/value
 
-    :param df: dataframe to profile
+    :param df: dataframe to _profile
     :return: dict with metrics per column
     """
     sm_dict: dict[str, dict] = {}
@@ -233,7 +233,7 @@ def get_columns_or_fields(cols: list[T.StructField]) -> list[T.StructField]:
 # TODO: split into managebale chunks
 # TODO: how to handle maps, arrays & structs?
 # TODO: return not only DQ rules, but also the profiling results - use named tuple?
-def profile_dataframe(
+def profile(
     df: DataFrame, cols: list[str] | None = None, opts: dict[str, Any] | None = None
 ) -> tuple[dict[str, Any], list[DQProfile]]:
     if opts is None:
@@ -254,12 +254,12 @@ def profile_dataframe(
     max_nulls = opts.get("max_null_ratio", 0)
     trim_strings = opts.get("trim_strings", True)
 
-    profile(df, df_cols, dq_rules, max_nulls, opts, summary_stats, total_count, trim_strings)
+    _profile(df, df_cols, dq_rules, max_nulls, opts, summary_stats, total_count, trim_strings)
 
     return summary_stats, dq_rules
 
 
-def profile(df, df_cols, dq_rules, max_nulls, opts, summary_stats, total_count, trim_strings):
+def _profile(df, df_cols, dq_rules, max_nulls, opts, summary_stats, total_count, trim_strings):
     # TODO: think, how we can do it in fewer passes. Maybe only for specific things, like, min_max, etc.
     for field in get_columns_or_fields(df_cols):
         field_name = field.name

src/databricks/labs/dqx/profiler/generator.py

@@ -1,5 +1,5 @@
 from databricks.labs.dqx.profiler.common import val_maybe_to_str
-from databricks.labs.dqx.profiler.profiler import DQProfile
+from databricks.labs.dqx.profiler.engine import DQProfile
 
 
 def dq_generate_is_in(col_name: str, level: str = "error", **params: dict):

tests/unit/test_dlt_rules_generator.py

@@ -2,7 +2,7 @@
 from typing import List
 
 from databricks.labs.dqx.profiler.dlt_generator import generate_dlt_rules
-from databricks.labs.dqx.profiler.profiler import DQProfile
+from databricks.labs.dqx.profiler.engine import DQProfile
 
 test_empty_rules: List[DQProfile] = []
 

tests/unit/test_profiler.py

@@ -2,11 +2,11 @@
 
 from pyspark.sql import SparkSession
 
-from databricks.labs.dqx.profiler.profiler import (
+from databricks.labs.dqx.profiler.engine import (
     DQProfile,
     T,
     get_columns_or_fields,
-    profile_dataframe,
+    profile,
 )
 
 
@@ -82,7 +82,7 @@ def test_profiler(spark_session: SparkSession):
         ],
         schema=inp_schema,
     )
-    stats, rules = profile_dataframe(inp_df)
+    stats, rules = profile(inp_df)
     # pprint.pprint(stats)
     # pprint.pprint(rules)
     expected_rules = [
@@ -114,6 +114,6 @@ def test_profiler(spark_session: SparkSession):
 def test_profiler_empty_df(spark_session: SparkSession):
     test_df = spark_session.createDataFrame([], "data: string")
 
-    actual_summary_stats, actual_dq_rule = profile_dataframe(test_df)
+    actual_summary_stats, actual_dq_rule = profile(test_df)
 
     assert len(actual_dq_rule) == 0

tests/unit/test_rules_generator.py

@@ -1,7 +1,7 @@
 import datetime
 
+from databricks.labs.dqx.profiler.engine import DQProfile
 from databricks.labs.dqx.profiler.generator import generate_dq_rules
-from databricks.labs.dqx.profiler.profiler import DQProfile
 
 test_rules = [
     DQProfile(