fixed fmt

Author: mwojtyczka

src/databricks/labs/dqx/dq_engine.py

@@ -1,15 +1,16 @@
 import functools as ft
 import itertools
 import re
+from collections.abc import Callable
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Any, Callable, Dict, List, Tuple
+from typing import Any
 
 import pyspark.sql.functions as F
 from pyspark.sql import Column, DataFrame
 
 from databricks.labs.dqx import dq_functions
-from databricks.labs.dqx.dq_functions import *  # noqa: F403
+from databricks.labs.dqx.dq_functions import *  # noqa: F403 # pylint: disable=wildcard-import,unused-wildcard-import
 from databricks.labs.dqx.utils import get_column_name
 
 
@@ -47,7 +48,7 @@ def rule_criticality(self) -> str:
         :return: string describing criticality - `warn` or `error`. Raises exception if it's something else
         """
         criticality = self.criticality
-        if criticality != Criticality.WARN.value and criticality != Criticality.ERROR.value:
+        if criticality not in (Criticality.WARN.value, Criticality.ERROR.value):
             criticality = Criticality.ERROR.value
 
         return criticality
@@ -79,13 +80,13 @@ class DQRuleColSet:
     * `check_func_kwargs` - keyword /named arguments for the check function after the col_name
     """
 
-    columns: List[str]
+    columns: list[str]
     check_func: Callable
     criticality: str = "error"
-    check_func_args: List[Any] = field(default_factory=list)
-    check_func_kwargs: Dict[str, Any] = field(default_factory=dict)
+    check_func_args: list[Any] = field(default_factory=list)
+    check_func_kwargs: dict[str, Any] = field(default_factory=dict)
 
-    def get_rules(self) -> List[DQRule]:
+    def get_rules(self) -> list[DQRule]:
         """Build a list of rules for a set of columns.
 
         :return: list of dq rules
@@ -99,7 +100,7 @@ def get_rules(self) -> List[DQRule]:
         ]
 
 
-def _perform_checks(df: DataFrame, checks: List[DQRule]) -> DataFrame:
+def _perform_checks(df: DataFrame, checks: list[DQRule]) -> DataFrame:
     """Applies a list of checks to a given dataframe and append results at the end of the dataframe.
 
     :param df: dataframe to check
@@ -110,7 +111,7 @@ def _perform_checks(df: DataFrame, checks: List[DQRule]) -> DataFrame:
     return df.select("*", *checks_cols)
 
 
-def _make_null_filter(cols: List[str]) -> Column:
+def _make_null_filter(cols: list[str]) -> Column:
     """Creates a filter condition that check if all specified columns are null.
 
     :param cols: names of the columns to check
@@ -122,14 +123,14 @@ def _make_null_filter(cols: List[str]) -> Column:
     def update_nullability_func(func, col):
         return func & F.col(col).isNull()
 
-    f1 = F.col(cols[0]).isNull()
-    return ft.reduce(update_nullability_func, cols[1:], f1)
+    initial = F.col(cols[0]).isNull()
+    return ft.reduce(update_nullability_func, cols[1:], initial)
 
 
 remove_criticality_re = re.compile("^(.*)_(error|warn)$")
 
 
-def _with_checks_as_map(df: DataFrame, dest_col: str, cols: List[str]) -> DataFrame:
+def _with_checks_as_map(df: DataFrame, dest_col: str, cols: list[str]) -> DataFrame:
     """Collect individual check columns into corresponding map<string, string> errors or warnings columns.
 
     :param df: dataframe with added check columns of type map<string, string>
@@ -168,7 +169,7 @@ def _with_checks_as_map(df: DataFrame, dest_col: str, cols: List[str]) -> DataFr
     return ndf
 
 
-def _get_check_columns(checks: List[DQRule], criticality: str) -> List[str]:
+def _get_check_columns(checks: list[DQRule], criticality: str) -> list[str]:
     """Get check columns based on criticality.
 
     :param checks: list of checks to apply to the dataframe
@@ -191,7 +192,7 @@ def _append_empty_checks(df: DataFrame) -> DataFrame:
     )
 
 
-def apply_checks(df: DataFrame, checks: List[DQRule]) -> DataFrame:
+def apply_checks(df: DataFrame, checks: list[DQRule]) -> DataFrame:
     """Applies data quality checks to a given dataframe.
 
     :param df: dataframe to check
@@ -214,7 +215,7 @@ def apply_checks(df: DataFrame, checks: List[DQRule]) -> DataFrame:
     return checked_df_map
 
 
-def apply_checks_and_split(df: DataFrame, checks: List[DQRule]) -> Tuple[DataFrame, DataFrame]:
+def apply_checks_and_split(df: DataFrame, checks: list[DQRule]) -> tuple[DataFrame, DataFrame]:
     """Applies data quality checks to a given dataframe and split it into two ("good" and "bad"),
     according to the data quality checks.
 
@@ -234,7 +235,7 @@ def apply_checks_and_split(df: DataFrame, checks: List[DQRule]) -> Tuple[DataFra
     return good_df, bad_df
 
 
-def build_checks_by_metadata(checks: List[dict], glbs=None) -> List[DQRule]:
+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.
 
     :param checks: list of dictionaries describing checks. Each check is a dictionary consisting of following fields:
@@ -248,25 +249,25 @@ def build_checks_by_metadata(checks: List[dict], glbs=None) -> List[DQRule]:
     :return: list of data quality check rules
     """
     dq_rule_checks = []
-    for ch in checks:
-        check = ch.get("check")
+    for check_def in checks:
+        check = check_def.get("check")
         if not check:
-            raise Exception(f"'check' block should be provided in the check: {ch}")
+            raise ValueError(f"'check' block should be provided in the check: {check}")
 
         func_name = check.get("function")
         if not func_name:
-            raise Exception(f"'function' argument should be provided in the check: {ch}")
+            raise ValueError(f"'function' argument should be provided in the check: {check}")
 
         if glbs:
             func = glbs.get(func_name)
         else:
             func = getattr(dq_functions, func_name)
 
         if not func or not callable(func):
-            raise Exception(f"function {func_name} is not defined")
+            raise ValueError(f"function {func_name} is not defined")
 
         func_args = check.get("arguments", {})
-        criticality = ch.get("criticality", "error")
+        criticality = check_def.get("criticality", "error")
 
         if "col_names" in func_args:
             dq_rule_checks += DQRuleColSet(
@@ -277,14 +278,16 @@ def build_checks_by_metadata(checks: List[dict], glbs=None) -> List[DQRule]:
                 check_func_kwargs={k: func_args[k] for k in func_args.keys() - {"col_names"}},
             ).get_rules()
         else:
-            name = ch.get("name", None)
+            name = check_def.get("name", None)
             check_func = func(**func_args)
             dq_rule_checks.append(DQRule(check=check_func, name=name, criticality=criticality))
 
     return dq_rule_checks
 
 
-def apply_checks_by_metadata_and_split(df: DataFrame, checks: List[dict], glbs=None) -> Tuple[DataFrame, DataFrame]:
+def apply_checks_by_metadata_and_split(
+    df: DataFrame, checks: list[dict], glbs: dict[str, Any] | None = None
+) -> tuple[DataFrame, DataFrame]:
     """Wrapper around `apply_checks_and_split` for use in the metadata-driven pipelines. The main difference
     is how the checks are specified - instead of using functions directly, they are described as function name plus
     arguments.
@@ -308,7 +311,7 @@ def apply_checks_by_metadata_and_split(df: DataFrame, checks: List[dict], glbs=N
     return good_df, bad_df
 
 
-def apply_checks_by_metadata(df: DataFrame, checks: List[dict], glbs=None) -> DataFrame:
+def apply_checks_by_metadata(df: DataFrame, checks: list[dict], glbs: dict[str, Any] | None = None) -> DataFrame:
     """Wrapper around `apply_checks` for use in the metadata-driven pipelines. The main difference
     is how the checks are specified - instead of using functions directly, they are described as function name plus
     arguments.
@@ -329,7 +332,7 @@ def apply_checks_by_metadata(df: DataFrame, checks: List[dict], glbs=None) -> Da
     return apply_checks(df, dq_rule_checks)
 
 
-def build_checks(*rules_col_set: DQRuleColSet) -> List[DQRule]:
+def build_checks(*rules_col_set: DQRuleColSet) -> list[DQRule]:
     """
     Build rules from dq rules and rule sets.
 

src/databricks/labs/dqx/dq_functions.py

@@ -33,10 +33,10 @@ def col_is_not_null_and_not_empty(col_name: str, trim_strings: bool = False) ->
     :param trim_strings: boolean flag to trim spaces from strings
     :return: Column object for condition
     """
-    cl = F.col(col_name)
+    column = F.col(col_name)
     if trim_strings:
-        cl = F.trim(cl).alias(col_name)
-    condition = cl.isNull() | (cl == "") | (cl == "null")
+        column = F.trim(column).alias(col_name)
+    condition = column.isNull() | (column == "") | (column == "null")
     return make_condition_col(condition, f"Column {col_name} is null or empty", f"{col_name}_is_null_or_empty")
 
 
@@ -46,8 +46,8 @@ def col_is_not_empty(col_name: str) -> Column:
     :param col_name: column name to check
     :return: Column object for condition
     """
-    cl = F.col(col_name)
-    return make_condition_col((cl == ""), f"Column {col_name} is empty", f"{col_name}_is_empty")
+    column = F.col(col_name)
+    return make_condition_col((column == ""), f"Column {col_name} is empty", f"{col_name}_is_empty")
 
 
 def col_is_not_null(col_name: str) -> Column:
@@ -56,8 +56,8 @@ def col_is_not_null(col_name: str) -> Column:
     :param col_name: column name to check
     :return: Column object for condition
     """
-    cl = F.col(col_name)
-    return make_condition_col(cl.isNull(), f"Column {col_name} is null", f"{col_name}_is_null")
+    column = F.col(col_name)
+    return make_condition_col(column.isNull(), f"Column {col_name} is null", f"{col_name}_is_null")
 
 
 def col_value_is_not_null_and_is_in_list(col_name: str, allowed: list) -> Column:
@@ -68,14 +68,14 @@ def col_value_is_not_null_and_is_in_list(col_name: str, allowed: list) -> Column
     :return: Column object for condition
     """
     allowed_cols = [item if isinstance(item, Column) else F.lit(item) for item in allowed]
-    cl = F.col(col_name)
-    condition = cl.isNull() | ~cl.isin(*allowed_cols)
+    column = F.col(col_name)
+    condition = column.isNull() | ~column.isin(*allowed_cols)
     return make_condition_col(
         condition,
         F.concat_ws(
             "",
             F.lit("Value "),
-            F.when(cl.isNull(), F.lit("null")).otherwise(cl),
+            F.when(column.isNull(), F.lit("null")).otherwise(column),
             F.lit(" is not in the allowed list: ["),
             F.concat_ws(", ", *allowed_cols),
             F.lit("]"),
@@ -92,14 +92,14 @@ def col_value_is_in_list(col_name: str, allowed: list) -> Column:
     :return: Column object for condition
     """
     allowed_cols = [item if isinstance(item, Column) else F.lit(item) for item in allowed]
-    cl = F.col(col_name)
-    condition = ~cl.isin(*allowed_cols)
+    column = F.col(col_name)
+    condition = ~column.isin(*allowed_cols)
     return make_condition_col(
         condition,
         F.concat_ws(
             "",
             F.lit("Value "),
-            F.when(cl.isNull(), F.lit("null")).otherwise(cl),
+            F.when(column.isNull(), F.lit("null")).otherwise(column),
             F.lit(" is not in the allowed list: ["),
             F.concat_ws(", ", *allowed_cols),
             F.lit("]"),
@@ -137,7 +137,7 @@ def col_sql_expression(
     return make_condition_col(expr_col, F.concat_ws("", F.lit(f"Value matches expression: {expression_msg}")), name)
 
 
-def is_col_older_than_col2_for_N_days(col_name1: str, col_name2: str, days: int) -> Column:
+def is_col_older_than_col2_for_N_days(col_name1: str, col_name2: str, days: int) -> Column:  # pylint: disable=invalid-name
     """Creates a condition column for case when one date or timestamp column is older than another column by N days.
 
     :param col_name1: first column
@@ -163,7 +163,7 @@ def is_col_older_than_col2_for_N_days(col_name1: str, col_name2: str, days: int)
     )
 
 
-def is_col_older_than_N_days(col_name: str, days: int, curr_date: Column | None = None) -> Column:
+def is_col_older_than_N_days(col_name: str, days: int, curr_date: Column | None = None) -> Column:  # pylint: disable=invalid-name
     """Creates a condition column for case when specified date or timestamp column is older (compared to current date)
     than N days.
 
@@ -222,6 +222,7 @@ def col_not_in_near_future(col_name: str, offset: int = 0, curr_timestamp: Colum
 
     :param col_name: column name
     :param offset: offset (in seconds) to add to the current timestamp at time of execution
+    :param curr_timestamp: (optional) set current timestamp
     :return: new Column
     """
     if curr_timestamp is None:
@@ -246,11 +247,11 @@ def col_not_in_near_future(col_name: str, offset: int = 0, curr_timestamp: Colum
     )
 
 
-def col_not_less_than(col_name: str, limit) -> Column:
+def col_not_less_than(col_name: str, limit: int) -> Column:
     """Creates a condition column that checks if a value is less than specified limit.
 
     :param col_name: column name
-    :param limit: limit
+    :param limit: limit to use in the condition
     :return: new Column
     """
     condition = F.col(col_name) < limit
@@ -262,11 +263,11 @@ def col_not_less_than(col_name: str, limit) -> Column:
     )
 
 
-def col_not_greater_than(col_name: str, limit) -> Column:
+def col_not_greater_than(col_name: str, limit: int) -> Column:
     """Creates a condition column that checks if a value is greater than specified limit.
 
     :param col_name: column name
-    :param limit: limit
+    :param limit: limit to use in the condition
     :return: new Column
     """
     condition = F.col(col_name) > limit
@@ -278,12 +279,12 @@ def col_not_greater_than(col_name: str, limit) -> Column:
     )
 
 
-def col_is_in_range(col_name: str, min_limit, max_limit) -> Column:
+def col_is_in_range(col_name: str, min_limit: int, max_limit: int) -> Column:
     """Creates a condition column that checks if a value is smaller than min limit or greater than max limit.
 
     :param col_name: column name
     :param min_limit: min limit
-    :param min_limit: max limit
+    :param max_limit: max limit
     :return: new Column
     """
     condition = (F.col(col_name) < min_limit) | (F.col(col_name) > max_limit)
@@ -304,12 +305,12 @@ def col_is_in_range(col_name: str, min_limit, max_limit) -> Column:
     )
 
 
-def col_is_not_in_range(col_name: str, min_limit, max_limit) -> Column:
+def col_is_not_in_range(col_name: str, min_limit: int, max_limit: int) -> Column:
     """Creates a condition column that checks if a value is within min and max limits.
 
     :param col_name: column name
     :param min_limit: min limit
-    :param min_limit: max limit
+    :param max_limit: max limit
     :return: new Column
     """
     condition = (F.col(col_name) > min_limit) & (F.col(col_name) < max_limit)

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

@@ -4,75 +4,75 @@
 from databricks.labs.dqx.profiler.profiler import DQRule
 
 
-def dq_generate_is_in(cl: str, level: str = "error", **params: dict):
+def dq_generate_is_in(col_name: str, level: str = "error", **params: dict):
     return {
-        "check": {"function": "col_value_is_in_list", "arguments": {"col_name": cl, "allowed": params["in"]}},
-        "name": f"{cl}_other_value",
+        "check": {"function": "col_value_is_in_list", "arguments": {"col_name": col_name, "allowed": params["in"]}},
+        "name": f"{col_name}_other_value",
         "criticality": level,
     }
 
 
 # TODO: rewrite it
-def dq_generate_min_max(cl: str, level: str = "error", **params: dict):
-    mn = params.get("min")
-    mx = params.get("max")
+def dq_generate_min_max(col_name: str, level: str = "error", **params: dict):
+    min = params.get("min")
+    max = params.get("max")
 
-    if mn is not None and mx is not None:
+    if min is not None and max is not None:
         return {
             "check": {
                 "function": "col_is_in_range",
                 "arguments": {
-                    "col_name": cl,
-                    "min_limit": val_maybe_to_str(mn, include_sql_quotes=False),
-                    "max_limit": val_maybe_to_str(mx, include_sql_quotes=False),
+                    "col_name": col_name,
+                    "min_limit": val_maybe_to_str(min, include_sql_quotes=False),
+                    "max_limit": val_maybe_to_str(max, include_sql_quotes=False),
                 },
             },
-            "name": f"{cl}_isnt_in_range",
+            "name": f"{col_name}_isnt_in_range",
             "criticality": level,
         }
-    elif mx is not None:
+    elif max is not None:
         return {
             "check": {
                 "function": "col_not_greater_than",
                 "arguments": {
-                    "col_name": cl,
-                    "val": val_maybe_to_str(mx, include_sql_quotes=False),
+                    "col_name": col_name,
+                    "val": val_maybe_to_str(max, include_sql_quotes=False),
                 },
             },
-            "name": f"{cl}_not_greater_than",
+            "name": f"{col_name}_not_greater_than",
             "criticality": level,
         }
-    elif mn is not None:
+    elif min is not None:
         return {
             "check": {
                 "function": "col_not_less_than",
                 "arguments": {
-                    "col_name": cl,
-                    "val": val_maybe_to_str(mn, include_sql_quotes=False),
+                    "col_name": col_name,
+                    "val": val_maybe_to_str(min, include_sql_quotes=False),
                 },
             },
-            "name": f"{cl}_not_less_than",
+            "name": f"{col_name}_not_less_than",
             "criticality": level,
         }
 
     return None
 
 
-def dq_generate_is_not_null(cl: str, level: str = "error", **params: dict):
+def dq_generate_is_not_null(col_name: str, level: str = "error", **params: dict):
     return {
-        "check": {"function": "col_is_not_null", "arguments": {"col_name": cl}},
-        "name": f"{cl}_is_null",
+        "check": {"function": "col_is_not_null", "arguments": {"col_name": col_name}},
+        "name": f"{col_name}_is_null",
         "criticality": level,
     }
 
 
-def dq_generate_is_not_null_or_empty(cl: str, level: str = "error", **params: dict):
+def dq_generate_is_not_null_or_empty(col_name: str, level: str = "error", **params: dict):
     return {
         "check": {
             "function": "col_is_not_null_and_not_empty",
-            "arguments": {"col_name": cl, "trim_strings": params.get("trim_strings", True)},
+            "arguments": {"col_name": col_name, "trim_strings": params.get("trim_strings", True)},
         },
-        "name": f"{cl}_is_null_or_empty",
+        "name": f"{col_name}_is_null_or_empty",
         "criticality": level,
     }
 
@@ -91,12 +91,12 @@ def generate_dq_rules(rules: Optional[list[DQRule]] = None, level: str = "error"
     dq_rules = []
     for rule in rules:
         nm = rule.name
-        cl = rule.column
+        col_name = rule.column
         params = rule.parameters or {}
         if nm not in dq_mapping:
-            print(f"No rule '{nm}' for column '{cl}'. skipping...")
+            print(f"No rule '{nm}' for column '{col_name}'. skipping...")
             continue
-        expr = dq_mapping[nm](cl, level, **params)
+        expr = dq_mapping[nm](col_name, level, **params)
         if expr:
             dq_rules.append(expr)
 

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

@@ -119,14 +119,14 @@ def extract_min_max(dst: DataFrame, nm: str, typ, metrics, opts: dict[str, Any]
     mx = None
 
     outlier_cols = opts.get("outlier_columns", [])
-    cl = dst.columns[0]
+    column = dst.columns[0]
     if opts.get("remove_outliers", True) and (len(outlier_cols) == 0 or nm in outlier_cols):  # detect outliers
         if typ == T.DateType():
-            dst = dst.select(F.col(cl).cast("timestamp").cast("bigint").alias(cl))
+            dst = dst.select(F.col(column).cast("timestamp").cast("bigint").alias(column))
         elif typ == T.TimestampType():
-            dst = dst.select(F.col(cl).cast("bigint").alias(cl))
+            dst = dst.select(F.col(column).cast("bigint").alias(column))
         # TODO: do summary instead? to get percentiles, etc.?
-        mn_mx = dst.agg(F.min(cl), F.max(cl), F.mean(cl), F.stddev(cl)).collect()
+        mn_mx = dst.agg(F.min(column), F.max(column), F.mean(column), F.stddev(column)).collect()
         if mn_mx and len(mn_mx) > 0:
             metrics["min"] = mn_mx[0][0]
             metrics["max"] = mn_mx[0][1]
@@ -175,7 +175,7 @@ def extract_min_max(dst: DataFrame, nm: str, typ, metrics, opts: dict[str, Any]
         else:
             print(f"Can't get min/max for field {nm}")
     else:
-        mn_mx = dst.agg(F.min(cl), F.max(cl)).collect()
+        mn_mx = dst.agg(F.min(column), F.max(column)).collect()
         if mn_mx and len(mn_mx) > 0:
             metrics["min"] = mn_mx[0][0]
             metrics["max"] = mn_mx[0][1]
@@ -239,14 +239,14 @@ def profile_dataframe(
 
     # 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):
-        nm = field.name
+        field_name = field.name
         typ = field.dataType
-        if nm not in summary_stats:
-            summary_stats[nm] = {}
-        metrics = summary_stats[nm]
+        if field_name not in summary_stats:
+            summary_stats[field_name] = {}
+        metrics = summary_stats[field_name]
 
         # calculate metrics
-        dst = df.select(nm).dropna()
+        dst = df.select(field_name).dropna()
         if typ == T.StringType() and trim_strings:
             cl = dst.columns[0]
             dst = dst.select(F.trim(F.col(cl)).alias(cl))
@@ -263,30 +263,32 @@ def profile_dataframe(
                 dq_rules.append(
                     DQRule(
                         name="is_not_null",
-                        column=nm,
-                        description=f"Column {nm} has {null_percentage * 100:.1f}% of null values "
+                        column=field_name,
+                        description=f"Column {field_name} has {null_percentage * 100:.1f}% of null values "
                         f"(allowed {max_nulls * 100:.1f}%)",
                     )
                 )
             else:
-                dq_rules.append(DQRule(name="is_not_null", column=nm))
+                dq_rules.append(DQRule(name="is_not_null", column=field_name))
 
         if type_supports_distinct(typ):
             dst2 = dst.dropDuplicates()
             cnt = dst2.count()
             if 0 < cnt < total_count * opts["distinct_ratio"] and cnt < opts["max_in_count"]:
-                dq_rules.append(DQRule(name="is_in", column=nm, parameters={"in": [row[0] for row in dst2.collect()]}))
+                dq_rules.append(
+                    DQRule(name="is_in", column=field_name, parameters={"in": [row[0] for row in dst2.collect()]})
+                )
 
         if typ == T.StringType():
-            dst2 = dst.filter(F.col(nm) == "")
+            dst2 = dst.filter(F.col(field_name) == "")
             cnt = dst2.count()
             if cnt <= (metrics["count"] * opts.get("max_empty_ratio", 0)):
                 dq_rules.append(
-                    DQRule(name="is_not_null_or_empty", column=nm, parameters={"trim_strings": trim_strings})
+                    DQRule(name="is_not_null_or_empty", column=field_name, parameters={"trim_strings": trim_strings})
                 )
 
         if metrics["count_non_null"] > 0 and type_supports_min_max(typ):
-            rule = extract_min_max(dst, nm, typ, metrics, opts)
+            rule = extract_min_max(dst, field_name, typ, metrics, opts)
             if rule:
                 dq_rules.append(rule)