From 4ffb641a84ce2621ca2390234cae67c8ae1c0f47 Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Tue, 9 Sep 2025 19:45:07 -0300 Subject: [PATCH 01/34] First commit: Installing DQX from feature branch --- src/databricks/labs/dqx/profiler/profiler.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 916fdbd33..3997c6380 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -53,6 +53,8 @@ def __init__(self, workspace_client: WorkspaceClient, spark: SparkSession | None "sample_fraction": 0.3, # fraction of data to sample (30%) "sample_seed": None, # seed for sampling "limit": 1000, # limit the number of samples + "filter":{}, # filter to apply before profiling + } @staticmethod From 17eb4ac1d6ff81a0a36a0c93f1b61a07511962d5 Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Wed, 10 Sep 2025 23:45:47 -0300 Subject: [PATCH 02/34] Implement methods in the DQProfiler class to filter dataframes before submited to profile. --- .gitignore | 1 + src/databricks/labs/dqx/profiler/profiler.py | 66 +++++++++++++++++++- 2 files changed, 65 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 16a704bde..196a43220 100644 --- a/.gitignore +++ b/.gitignore @@ -50,6 +50,7 @@ htmlcov/ .cache nosetests.xml coverage.xml +coverage-unit.xml *.cover *.py,cover .hypothesis/ diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 3997c6380..146439d3c 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -53,8 +53,7 @@ def __init__(self, workspace_client: WorkspaceClient, spark: SparkSession | None "sample_fraction": 0.3, # fraction of data to sample (30%) "sample_seed": None, # seed for sampling "limit": 1000, # limit the number of samples - "filter":{}, # filter to apply before profiling - + "dataset_filter": None, # filter to apply to the dataset } @staticmethod @@ -94,6 +93,7 @@ def profile( Returns: A tuple containing a dictionary of summary statistics and a list of data quality profiles. """ + columns = columns or df.columns df_columns = [f for f in df.schema.fields if f.name in columns] df = df.select(*[f.name for f in df_columns]) @@ -316,7 +316,10 @@ def _sample(df: DataFrame, opts: dict[str, Any]) -> DataFrame: sample_fraction = opts.get("sample_fraction", None) sample_seed = opts.get("sample_seed", None) limit = opts.get("limit", None) + filter = opts.get("dataset_filter", None) + if filter: + df = DQProfiler._filter_dataframe(df, filter) if sample_fraction: df = df.sample(withReplacement=False, fraction=sample_fraction, seed=sample_seed) if limit: @@ -796,3 +799,62 @@ def _round_decimal(value: decimal.Decimal, direction: str) -> decimal.Decimal: if direction == "up": return value.to_integral_value(rounding=decimal.ROUND_CEILING) return value + + @staticmethod + def _filter_dataframe(df: DataFrame, filter: dict[str, Any]) -> DataFrame: + """ + Filters a DataFrame based on a dynamic filter. + + Args: + df: The input DataFrame to filter. + filter: A dictionary where keys are column names and values are the conditions to filter on. + + Returns: + A filtered DataFrame. + """ + + # Get the list of columns in the DataFrame + df_columns = df.columns + + for column, condition in filter.items(): + if column not in df_columns: + raise ValueError(f"Column '{column}' does not exist in the DataFrame.") + if isinstance(condition, (list, tuple)): + # If the condition is a list or tuple, use the `isin` filter + df = df.filter(F.col(column).isin(condition)) + elif isinstance(condition, dict): + # If the condition is a dictionary, handle operators like >, <, etc. + for operator, value in condition.items(): + df = DQProfiler._apply_filter_operator(df, column, operator, value) + else: + # Default equality filter + df = df.filter(F.col(column) == condition) + return df + + @staticmethod + def _apply_filter_operator(df: DataFrame, column: str, operator: str, value: Any) -> DataFrame: + """ + Applies a specific operator-based condition to a DataFrame column. + + Args: + df: The DataFrame to filter. + column: The column name to apply the operator on. + operator: The operator to use (e.g., "gt", "lt"). + value: The value to compare against. + + Returns: + A filtered DataFrame. + """ + operators = { + "gt": F.col(column) > value, + "lt": F.col(column) < value, + "gte": F.col(column) >= value, + "lte": F.col(column) <= value, + "eq": F.col(column) == value, + "neq": F.col(column) != value, + } + + if operator not in operators: + raise ValueError(f"Unsupported operator '{operator}' for column '{column}'.") + + return df.filter(operators[operator]) From d75b09d6f3a550986f7e73907e7bba2cd169427b Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Thu, 11 Sep 2025 22:36:48 -0300 Subject: [PATCH 03/34] Refactor _sample method to utilize filtered dataframe --- src/databricks/labs/dqx/profiler/profiler.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 146439d3c..a14b812d3 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -318,12 +318,12 @@ def _sample(df: DataFrame, opts: dict[str, Any]) -> DataFrame: limit = opts.get("limit", None) filter = opts.get("dataset_filter", None) - if filter: - df = DQProfiler._filter_dataframe(df, filter) + + df_filter = DQProfiler._filter_dataframe(df, filter) if sample_fraction: - df = df.sample(withReplacement=False, fraction=sample_fraction, seed=sample_seed) + df = df_filter.sample(withReplacement=False, fraction=sample_fraction, seed=sample_seed) if limit: - df = df.limit(limit) + df = df_filter.limit(limit) return df @@ -812,6 +812,8 @@ def _filter_dataframe(df: DataFrame, filter: dict[str, Any]) -> DataFrame: Returns: A filtered DataFrame. """ + if not filter: + return df # Get the list of columns in the DataFrame df_columns = df.columns @@ -827,8 +829,8 @@ def _filter_dataframe(df: DataFrame, filter: dict[str, Any]) -> DataFrame: for operator, value in condition.items(): df = DQProfiler._apply_filter_operator(df, column, operator, value) else: - # Default equality filter - df = df.filter(F.col(column) == condition) + # If the condition is not a list, tuple, or dictionary, raise an error + raise ValueError(f"Unsupported filter condition for column '{column}': {condition}") return df @staticmethod From 530b8fbb4144e99d9b4f8a3818e60ca582c8ebf1 Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Fri, 12 Sep 2025 23:40:13 -0300 Subject: [PATCH 04/34] Refactor _sample method to utilize return the filtered dataframe --- src/databricks/labs/dqx/profiler/profiler.py | 9 ++++----- tests/unit/test_profiler.py | 1 + 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index a14b812d3..ddc3bc7ca 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -318,12 +318,11 @@ def _sample(df: DataFrame, opts: dict[str, Any]) -> DataFrame: limit = opts.get("limit", None) filter = opts.get("dataset_filter", None) - - df_filter = DQProfiler._filter_dataframe(df, filter) + df = DQProfiler._filter_dataframe(df, filter) if sample_fraction: - df = df_filter.sample(withReplacement=False, fraction=sample_fraction, seed=sample_seed) + df = df.sample(withReplacement=False, fraction=sample_fraction, seed=sample_seed) if limit: - df = df_filter.limit(limit) + df = df.limit(limit) return df @@ -830,7 +829,7 @@ def _filter_dataframe(df: DataFrame, filter: dict[str, Any]) -> DataFrame: df = DQProfiler._apply_filter_operator(df, column, operator, value) else: # If the condition is not a list, tuple, or dictionary, raise an error - raise ValueError(f"Unsupported filter condition for column '{column}': {condition}") + raise ValueError(f"Unsupported filter condition for column '{column}': {condition}") return df @staticmethod diff --git a/tests/unit/test_profiler.py b/tests/unit/test_profiler.py index 62b0b14d6..bfd033502 100644 --- a/tests/unit/test_profiler.py +++ b/tests/unit/test_profiler.py @@ -2,6 +2,7 @@ from databricks.labs.dqx.profiler.profiler import DQProfiler + def test_get_columns_or_fields(): inp = T.StructType( [ From 2c21d2fadc79ae433cf4e973cfb7b7a0d5b23fa5 Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Sat, 13 Sep 2025 16:47:32 -0300 Subject: [PATCH 05/34] Implement print statement to test dataset_filter feature --- src/databricks/labs/dqx/profiler/profiler.py | 7 +- tests/integration/test_profiler.py | 115 +++++++++++++++++++ tests/unit/test_profiler.py | 1 + 3 files changed, 120 insertions(+), 3 deletions(-) diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index ddc3bc7ca..66db616dc 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -53,7 +53,7 @@ def __init__(self, workspace_client: WorkspaceClient, spark: SparkSession | None "sample_fraction": 0.3, # fraction of data to sample (30%) "sample_seed": None, # seed for sampling "limit": 1000, # limit the number of samples - "dataset_filter": None, # filter to apply to the dataset + # "dataset_filter": None, # filter to apply to the dataset } @staticmethod @@ -102,14 +102,15 @@ def profile( options = {} options = {**self.default_profile_options, **options} # merge default options with user-provided options + print(options) df = self._sample(df, options) - + print(df) dq_rules: list[DQProfile] = [] total_count = df.count() summary_stats = self._get_df_summary_as_dict(df) if total_count == 0: return summary_stats, dq_rules - + print(summary_stats) self._profile(df, df_columns, dq_rules, options, summary_stats, total_count) return summary_stats, dq_rules diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index c842dccc4..b578c8490 100644 --- a/tests/integration/test_profiler.py +++ b/tests/integration/test_profiler.py @@ -1254,3 +1254,118 @@ def test_profile_tables_no_tables_or_patterns(ws): with pytest.raises(ValueError, match="Either 'tables' or 'patterns' must be provided"): profiler.profile_tables() + +def test_profile_sample_dataset_filter(spark, ws): + schema = T.StructType( + [ + T.StructField("maintenance_id", T.StringType(), False), + T.StructField("machine_id", T.StringType(), False), + T.StructField("maintenance_type", T.StringType(), True), + T.StructField("maintenance_date", T.DateType(), True), + T.StructField("duration_minutes", T.IntegerType(), True), + T.StructField("cost", T.DecimalType(10, 2), True), + T.StructField("next_scheduled_date", T.DateType(), True), + T.StructField("work_order_id", T.StringType(), True), + T.StructField("safety_check_passed", T.BooleanType(), True), + T.StructField("parts_list", T.ArrayType(T.StringType()), True), + T.StructField("ingest_date", T.DateType(), True), + ] + ) + maintenance_data = [ + # Ingest date 2025-04-28 + ( + "MTN-001", + "MCH-001", + "preventive", + datetime.strptime("2025-04-01", "%Y-%m-%d").date(), + 120, + Decimal("450.00"), + datetime.strptime("2025-07-01", "%Y-%m-%d").date(), + "WO-001", + True, + ["filter", "gasket"], + datetime.strptime("2025-04-28", "%Y-%m-%d").date(), + ), + ( + "MTN-002", + "MCH-002", + "corrective", + datetime.strptime("2025-04-15", "%Y-%m-%d").date(), + 240, + Decimal("1200.50"), + datetime.strptime("2026-04-01", "%Y-%m-%d").date(), + "WO-002", + False, + ["motor"], + datetime.strptime("2025-04-28", "%Y-%m-%d").date(), + ), # Future date issue + # Ingest date 2025-04-29 + ( + "MTN-003", + "MCH-003", + None, + datetime.strptime("2025-04-20", "%Y-%m-%d").date(), + -30, + Decimal("-500.00"), + datetime.strptime("2024-04-20", "%Y-%m-%d").date(), + "INVALID", + None, + [], + datetime.strptime("2025-04-29", "%Y-%m-%d").date(), + ), # Multiple issues + ( + "MTN-004", + "MCH-001", + "predictive", + datetime.strptime("2025-04-25", "%Y-%m-%d").date(), + 180, + Decimal("800.00"), + datetime.strptime("2025-10-01", "%Y-%m-%d").date(), + "WO-003", + True, + ["sensor"], + datetime.strptime("2025-04-29", "%Y-%m-%d").date(), + ), + # Ingest date 2025-04-30 + ( + "MTN-005", + "MCH-002", + "preventive", + datetime.strptime("2025-04-29", "%Y-%m-%d").date(), + 90, + Decimal("300.00"), + datetime.strptime("2025-07-15", "%Y-%m-%d").date(), + "WO-004", + True, + ["valve"], + datetime.strptime("2025-04-30", "%Y-%m-%d").date(), + ), + ( + "MTN-006", + "MCH-003", + "corrective", + datetime.strptime("2025-04-30", "%Y-%m-%d").date(), + 60, + Decimal("150.00"), + datetime.strptime("2025-08-01", "%Y-%m-%d").date(), + "WO-005", + False, + ["pump"], + datetime.strptime("2025-04-30", "%Y-%m-%d").date(), + ), + ] + + + input_df = spark.createDataFrame(maintenance_data, schema=schema) + + custom_options = { + "sample_fraction": None, + "limit": None, + "dataset_filter": {"machine_id":['MCH-002', 'MCH-003'], "maintenance_type":{'eq': 'preventive'}} + } + + + profiler = DQProfiler(ws) + filtered_df = profiler._sample(input_df, opts=custom_options) + assert filtered_df.count() == 1 + \ No newline at end of file diff --git a/tests/unit/test_profiler.py b/tests/unit/test_profiler.py index bfd033502..76a355b2f 100644 --- a/tests/unit/test_profiler.py +++ b/tests/unit/test_profiler.py @@ -29,3 +29,4 @@ def test_get_columns_or_fields(): T.StructField("ss1.s2.ns3", T.DateType()), ] assert fields == expected + From e8ee1d2aff198c4ddbb626978a86f2e0b05ebded Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Sun, 14 Sep 2025 22:27:32 -0300 Subject: [PATCH 06/34] Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter --- src/databricks/labs/dqx/profiler/profiler.py | 88 ++-------- tests/integration/test_profiler.py | 163 +++++++++++-------- 2 files changed, 117 insertions(+), 134 deletions(-) diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 66db616dc..8f7e2b2c5 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -53,7 +53,7 @@ def __init__(self, workspace_client: WorkspaceClient, spark: SparkSession | None "sample_fraction": 0.3, # fraction of data to sample (30%) "sample_seed": None, # seed for sampling "limit": 1000, # limit the number of samples - # "dataset_filter": None, # filter to apply to the dataset + "dataset_filter_expression": None, # filter to apply to the dataset } @staticmethod @@ -102,15 +102,15 @@ def profile( options = {} options = {**self.default_profile_options, **options} # merge default options with user-provided options - print(options) + df = self._sample(df, options) - print(df) + dq_rules: list[DQProfile] = [] total_count = df.count() summary_stats = self._get_df_summary_as_dict(df) if total_count == 0: return summary_stats, dq_rules - print(summary_stats) + self._profile(df, df_columns, dq_rules, options, summary_stats, total_count) return summary_stats, dq_rules @@ -317,9 +317,10 @@ def _sample(df: DataFrame, opts: dict[str, Any]) -> DataFrame: sample_fraction = opts.get("sample_fraction", None) sample_seed = opts.get("sample_seed", None) limit = opts.get("limit", None) - filter = opts.get("dataset_filter", None) + filter = opts.get("dataset_filter_expression", None) - df = DQProfiler._filter_dataframe(df, filter) + if filter: + df = df.filter(filter) if sample_fraction: df = df.sample(withReplacement=False, fraction=sample_fraction, seed=sample_seed) if limit: @@ -370,6 +371,7 @@ def _calculate_metrics( """ max_nulls = opts.get("max_null_ratio", 0) trim_strings = opts.get("trim_strings", True) + filter = opts.get("dataset_filter_expression", None) dst = df.select(field_name).dropna() if typ == T.StringType() and trim_strings: @@ -389,16 +391,19 @@ def _calculate_metrics( column=field_name, description=f"Column {field_name} has {null_percentage * 100:.1f}% of null values " f"(allowed {max_nulls * 100:.1f}%)", + parameters={"dataset_filter_expression": filter}, ) ) else: - dq_rules.append(DQProfile(name="is_not_null", column=field_name)) + dq_rules.append(DQProfile(name="is_not_null", column=field_name, parameters={"dataset_filter_expression": filter})) if self._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( - DQProfile(name="is_in", column=field_name, parameters={"in": [row[0] for row in dst2.collect()]}) + DQProfile(name="is_in", + column=field_name, + parameters={"in": [row[0] for row in dst2.collect()], "dataset_filter_expression": filter}) ) if ( typ == T.StringType() @@ -410,7 +415,9 @@ def _calculate_metrics( cnt = dst2.count() if cnt <= (metrics["count"] * opts.get("max_empty_ratio", 0)): dq_rules.append( - DQProfile(name="is_not_null_or_empty", column=field_name, parameters={"trim_strings": trim_strings}) + DQProfile(name="is_not_null_or_empty", + column=field_name, + parameters={"trim_strings": trim_strings, "dataset_filter_expression": filter}) ) if metrics["count_non_null"] > 0 and self._type_supports_min_max(typ): rule = self._extract_min_max(dst, field_name, typ, metrics, opts) @@ -528,7 +535,7 @@ def _extract_min_max( if opts is None: opts = {} - + filter= opts.get("dataset_filter_expression", None) outlier_cols = opts.get("outlier_columns", []) column = dst.columns[0] if opts.get("remove_outliers", True) and ( @@ -564,7 +571,7 @@ def _extract_min_max( logger.info(f"Can't get min/max for field {col_name}") if descr and min_limit and max_limit: return DQProfile( - name="min_max", column=col_name, parameters={"min": min_limit, "max": max_limit}, description=descr + name="min_max", column=col_name, parameters={"min": min_limit, "max": max_limit,"dataset_filter_expression": filter}, description=descr ) return None @@ -800,63 +807,4 @@ def _round_decimal(value: decimal.Decimal, direction: str) -> decimal.Decimal: return value.to_integral_value(rounding=decimal.ROUND_CEILING) return value - @staticmethod - def _filter_dataframe(df: DataFrame, filter: dict[str, Any]) -> DataFrame: - """ - Filters a DataFrame based on a dynamic filter. - - Args: - df: The input DataFrame to filter. - filter: A dictionary where keys are column names and values are the conditions to filter on. - - Returns: - A filtered DataFrame. - """ - if not filter: - return df - - # Get the list of columns in the DataFrame - df_columns = df.columns - - for column, condition in filter.items(): - if column not in df_columns: - raise ValueError(f"Column '{column}' does not exist in the DataFrame.") - if isinstance(condition, (list, tuple)): - # If the condition is a list or tuple, use the `isin` filter - df = df.filter(F.col(column).isin(condition)) - elif isinstance(condition, dict): - # If the condition is a dictionary, handle operators like >, <, etc. - for operator, value in condition.items(): - df = DQProfiler._apply_filter_operator(df, column, operator, value) - else: - # If the condition is not a list, tuple, or dictionary, raise an error - raise ValueError(f"Unsupported filter condition for column '{column}': {condition}") - return df - - @staticmethod - def _apply_filter_operator(df: DataFrame, column: str, operator: str, value: Any) -> DataFrame: - """ - Applies a specific operator-based condition to a DataFrame column. - - Args: - df: The DataFrame to filter. - column: The column name to apply the operator on. - operator: The operator to use (e.g., "gt", "lt"). - value: The value to compare against. - - Returns: - A filtered DataFrame. - """ - operators = { - "gt": F.col(column) > value, - "lt": F.col(column) < value, - "gte": F.col(column) >= value, - "lte": F.col(column) <= value, - "eq": F.col(column) == value, - "neq": F.col(column) != value, - } - - if operator not in operators: - raise ValueError(f"Unsupported operator '{operator}' for column '{column}'.") - return df.filter(operators[operator]) diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index b578c8490..ce580dc8a 100644 --- a/tests/integration/test_profiler.py +++ b/tests/integration/test_profiler.py @@ -1255,104 +1255,85 @@ def test_profile_tables_no_tables_or_patterns(ws): with pytest.raises(ValueError, match="Either 'tables' or 'patterns' must be provided"): profiler.profile_tables() -def test_profile_sample_dataset_filter(spark, ws): - schema = T.StructType( - [ - T.StructField("maintenance_id", T.StringType(), False), +def test_profile_with_dataset_filter(spark, ws): + schema = T.StructType( [ + T.StructField("machine_id", T.StringType(), False), T.StructField("maintenance_type", T.StringType(), True), - T.StructField("maintenance_date", T.DateType(), True), - T.StructField("duration_minutes", T.IntegerType(), True), + T.StructField("maintenance_date", T.DateType(), True), T.StructField("cost", T.DecimalType(10, 2), True), T.StructField("next_scheduled_date", T.DateType(), True), - T.StructField("work_order_id", T.StringType(), True), T.StructField("safety_check_passed", T.BooleanType(), True), - T.StructField("parts_list", T.ArrayType(T.StringType()), True), - T.StructField("ingest_date", T.DateType(), True), + ] ) maintenance_data = [ - # Ingest date 2025-04-28 + ( - "MTN-001", "MCH-001", "preventive", - datetime.strptime("2025-04-01", "%Y-%m-%d").date(), - 120, + date(2025, 4,1), Decimal("450.00"), - datetime.strptime("2025-07-01", "%Y-%m-%d").date(), - "WO-001", - True, - ["filter", "gasket"], - datetime.strptime("2025-04-28", "%Y-%m-%d").date(), + date(2025,7,1), + True, ), ( - "MTN-002", "MCH-002", "corrective", - datetime.strptime("2025-04-15", "%Y-%m-%d").date(), - 240, + date(2025, 4,15), Decimal("1200.50"), - datetime.strptime("2026-04-01", "%Y-%m-%d").date(), - "WO-002", - False, - ["motor"], - datetime.strptime("2025-04-28", "%Y-%m-%d").date(), - ), # Future date issue - # Ingest date 2025-04-29 + date(2026,4,1), + False, + ), ( - "MTN-003", "MCH-003", None, - datetime.strptime("2025-04-20", "%Y-%m-%d").date(), - -30, + date(2025,4,20), Decimal("-500.00"), - datetime.strptime("2024-04-20", "%Y-%m-%d").date(), - "INVALID", - None, - [], - datetime.strptime("2025-04-29", "%Y-%m-%d").date(), - ), # Multiple issues + date(2024,4,20), + None, + ), ( - "MTN-004", "MCH-001", "predictive", - datetime.strptime("2025-04-25", "%Y-%m-%d").date(), - 180, + date(2025,4,25), Decimal("800.00"), - datetime.strptime("2025-10-01", "%Y-%m-%d").date(), - "WO-003", - True, - ["sensor"], - datetime.strptime("2025-04-29", "%Y-%m-%d").date(), + date(2025,10,1), + True, ), - # Ingest date 2025-04-30 + ( - "MTN-005", "MCH-002", "preventive", - datetime.strptime("2025-04-29", "%Y-%m-%d").date(), - 90, + date(2025,4,29), Decimal("300.00"), - datetime.strptime("2025-07-15", "%Y-%m-%d").date(), - "WO-004", - True, - ["valve"], - datetime.strptime("2025-04-30", "%Y-%m-%d").date(), + date(2025,7,15), + True, ), ( - "MTN-006", "MCH-003", "corrective", - datetime.strptime("2025-04-30", "%Y-%m-%d").date(), - 60, + date(2025,4,30), + Decimal("150.00"), + date(2025,8,1), + False, + ), + ( + "MCH-002", + "preventive", + date(2025,5,30), Decimal("150.00"), - datetime.strptime("2025-08-01", "%Y-%m-%d").date(), - "WO-005", - False, - ["pump"], - datetime.strptime("2025-04-30", "%Y-%m-%d").date(), + date(2025,9,1), + True, ), + ( + "MCH-002", + "preventive", + date(2025,7,30), + Decimal("100.00"), + date(2025,12,1), + True, + ) ] @@ -1361,11 +1342,65 @@ def test_profile_sample_dataset_filter(spark, ws): custom_options = { "sample_fraction": None, "limit": None, - "dataset_filter": {"machine_id":['MCH-002', 'MCH-003'], "maintenance_type":{'eq': 'preventive'}} + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" } profiler = DQProfiler(ws) filtered_df = profiler._sample(input_df, opts=custom_options) - assert filtered_df.count() == 1 + + stats, rules = profiler.profile(input_df,options=custom_options) + + expected_rules = [ + DQProfile(name="is_not_null", + column="machine_id", + description=None, + parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + DQProfile(name="is_not_null", + column="maintenance_type", + description=None, + parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + DQProfile(name="is_not_null", + column="maintenance_date", + description=None, + parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + DQProfile( + name="min_max", + column="maintenance_date", + description="Real min/max values were used", + parameters={"min": date(2025, 4, 29), + "max": date(2025, 7, 30), + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + ), + DQProfile(name="is_not_null", + column="cost", + description=None, + parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + DQProfile( + name="min_max", + column="cost", + parameters={"min": Decimal('100.00') , "max": Decimal('300.00'),"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + description="Real min/max values were used" + ), + DQProfile(name="is_not_null", + column="next_scheduled_date", + description=None, + parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + DQProfile( + name="min_max", + column="next_scheduled_date", + description="Real min/max values were used", + parameters={"min": date(2025, 7, 15), "max": date(2025, 12, 1),"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + ), + DQProfile(name="is_not_null", + column="safety_check_passed", + description=None, + parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + + ] + + + assert filtered_df.count() == 3 + assert len(stats.keys()) > 0 + assert rules == expected_rules \ No newline at end of file From 0b6a267346626b0ccb401bd1af1292486f4cb777 Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Fri, 19 Sep 2025 09:53:00 -0300 Subject: [PATCH 07/34] Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file. --- .gitignore | 3 + databricks.yml | 19 ++ src/databricks/labs/dqx/profiler/generator.py | 9 + src/databricks/labs/dqx/profiler/profiler.py | 35 ++-- tests/integration/test_profiler.py | 189 ++++++++++-------- tests/integration/test_rules_generator.py | 110 ++++++++++ .../test_save_checks_to_workspace_file.py | 11 +- tests/unit/test_profiler.py | 2 - 8 files changed, 280 insertions(+), 98 deletions(-) create mode 100644 databricks.yml diff --git a/.gitignore b/.gitignore index 196a43220..770857d8a 100644 --- a/.gitignore +++ b/.gitignore @@ -169,3 +169,6 @@ dev/cleanup.py # docgen docs/dqx/docs/reference/api !docs/dqx/docs/reference/api/index.mdx + +.databricks + diff --git a/databricks.yml b/databricks.yml new file mode 100644 index 000000000..6bdb7731f --- /dev/null +++ b/databricks.yml @@ -0,0 +1,19 @@ +# This is a Databricks asset bundle definition for dqx. +# The Databricks extension requires databricks.yml configuration file. +# See https://docs.databricks.com/dev-tools/bundles/index.html for documentation. + +bundle: + name: dqx + +targets: + dev: + mode: development + default: true + workspace: + host: https://dbc-08fa3e7e-3bcc.cloud.databricks.com + + ## Optionally, there could be 'staging' or 'prod' targets here. + # + # prod: + # workspace: + # host: https://dbc-08fa3e7e-3bcc.cloud.databricks.com diff --git a/src/databricks/labs/dqx/profiler/generator.py b/src/databricks/labs/dqx/profiler/generator.py index 5ca43681d..64d5e3654 100644 --- a/src/databricks/labs/dqx/profiler/generator.py +++ b/src/databricks/labs/dqx/profiler/generator.py @@ -58,6 +58,7 @@ def dq_generate_is_in(column: str, level: str = "error", **params: dict): "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"]}}, "name": f"{column}_other_value", "criticality": level, + "filter": params.get("dataset_filter_expression", None), } @staticmethod @@ -75,6 +76,7 @@ def dq_generate_min_max(column: str, level: str = "error", **params: dict): """ min_limit = params.get("min") max_limit = params.get("max") + filter = params.get("dataset_filter_expression", None) if not isinstance(min_limit, int) or not isinstance(max_limit, int): return None # TODO handle timestamp and dates: https://github.com/databrickslabs/dqx/issues/71 @@ -91,6 +93,7 @@ def dq_generate_min_max(column: str, level: str = "error", **params: dict): }, "name": f"{column}_isnt_in_range", "criticality": level, + "filter": filter, } if max_limit is not None: @@ -104,6 +107,7 @@ def dq_generate_min_max(column: str, level: str = "error", **params: dict): }, "name": f"{column}_not_greater_than", "criticality": level, + "filter": filter, } if min_limit is not None: @@ -117,6 +121,7 @@ def dq_generate_min_max(column: str, level: str = "error", **params: dict): }, "name": f"{column}_not_less_than", "criticality": level, + "filter": filter, } return None @@ -135,10 +140,12 @@ def dq_generate_is_not_null(column: str, level: str = "error", **params: dict): A dictionary representing the data quality rule. """ params = params or {} + filter = params.get("dataset_filter_expression", None) return { "check": {"function": "is_not_null", "arguments": {"column": column}}, "name": f"{column}_is_null", "criticality": level, + "filter": filter, } @staticmethod @@ -154,6 +161,7 @@ def dq_generate_is_not_null_or_empty(column: str, level: str = "error", **params Returns: A dictionary representing the data quality rule. """ + filter = params.get("dataset_filter_expression", None) return { "check": { "function": "is_not_null_and_not_empty", @@ -161,6 +169,7 @@ def dq_generate_is_not_null_or_empty(column: str, level: str = "error", **params }, "name": f"{column}_is_null_or_empty", "criticality": level, + "filter": filter, } _checks_mapping = { diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 8f7e2b2c5..4598990c7 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -102,15 +102,15 @@ def profile( options = {} options = {**self.default_profile_options, **options} # merge default options with user-provided options - + df = self._sample(df, options) - + dq_rules: list[DQProfile] = [] total_count = df.count() summary_stats = self._get_df_summary_as_dict(df) if total_count == 0: return summary_stats, dq_rules - + self._profile(df, df_columns, dq_rules, options, summary_stats, total_count) return summary_stats, dq_rules @@ -395,15 +395,19 @@ def _calculate_metrics( ) ) else: - dq_rules.append(DQProfile(name="is_not_null", column=field_name, parameters={"dataset_filter_expression": filter})) + dq_rules.append( + DQProfile(name="is_not_null", column=field_name, parameters={"dataset_filter_expression": filter}) + ) if self._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( - DQProfile(name="is_in", - column=field_name, - parameters={"in": [row[0] for row in dst2.collect()], "dataset_filter_expression": filter}) + DQProfile( + name="is_in", + column=field_name, + parameters={"in": [row[0] for row in dst2.collect()], "dataset_filter_expression": filter}, + ) ) if ( typ == T.StringType() @@ -415,9 +419,11 @@ def _calculate_metrics( cnt = dst2.count() if cnt <= (metrics["count"] * opts.get("max_empty_ratio", 0)): dq_rules.append( - DQProfile(name="is_not_null_or_empty", - column=field_name, - parameters={"trim_strings": trim_strings, "dataset_filter_expression": filter}) + DQProfile( + name="is_not_null_or_empty", + column=field_name, + parameters={"trim_strings": trim_strings, "dataset_filter_expression": filter}, + ) ) if metrics["count_non_null"] > 0 and self._type_supports_min_max(typ): rule = self._extract_min_max(dst, field_name, typ, metrics, opts) @@ -535,7 +541,7 @@ def _extract_min_max( if opts is None: opts = {} - filter= opts.get("dataset_filter_expression", None) + filter = opts.get("dataset_filter_expression", None) outlier_cols = opts.get("outlier_columns", []) column = dst.columns[0] if opts.get("remove_outliers", True) and ( @@ -571,7 +577,10 @@ def _extract_min_max( logger.info(f"Can't get min/max for field {col_name}") if descr and min_limit and max_limit: return DQProfile( - name="min_max", column=col_name, parameters={"min": min_limit, "max": max_limit,"dataset_filter_expression": filter}, description=descr + name="min_max", + column=col_name, + parameters={"min": min_limit, "max": max_limit, "dataset_filter_expression": filter}, + description=descr, ) return None @@ -806,5 +815,3 @@ def _round_decimal(value: decimal.Decimal, direction: str) -> decimal.Decimal: if direction == "up": return value.to_integral_value(rounding=decimal.ROUND_CEILING) return value - - diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index ce580dc8a..bd3bab5ec 100644 --- a/tests/integration/test_profiler.py +++ b/tests/integration/test_profiler.py @@ -1255,152 +1255,179 @@ def test_profile_tables_no_tables_or_patterns(ws): with pytest.raises(ValueError, match="Either 'tables' or 'patterns' must be provided"): profiler.profile_tables() + def test_profile_with_dataset_filter(spark, ws): - schema = T.StructType( [ - + schema = T.StructType( + [ T.StructField("machine_id", T.StringType(), False), T.StructField("maintenance_type", T.StringType(), True), - T.StructField("maintenance_date", T.DateType(), True), + T.StructField("maintenance_date", T.DateType(), True), T.StructField("cost", T.DecimalType(10, 2), True), T.StructField("next_scheduled_date", T.DateType(), True), T.StructField("safety_check_passed", T.BooleanType(), True), - ] ) maintenance_data = [ - ( "MCH-001", "preventive", - date(2025, 4,1), + date(2025, 4, 1), Decimal("450.00"), - date(2025,7,1), - True, + date(2025, 7, 1), + True, ), ( "MCH-002", "corrective", - date(2025, 4,15), + date(2025, 4, 15), Decimal("1200.50"), - date(2026,4,1), - False, - ), + date(2026, 4, 1), + False, + ), ( "MCH-003", None, - date(2025,4,20), + date(2025, 4, 20), Decimal("-500.00"), - date(2024,4,20), - None, - ), + date(2024, 4, 20), + None, + ), ( "MCH-001", "predictive", - date(2025,4,25), + date(2025, 4, 25), Decimal("800.00"), - date(2025,10,1), - True, + date(2025, 10, 1), + True, ), - ( "MCH-002", "preventive", - date(2025,4,29), + date(2025, 4, 29), Decimal("300.00"), - date(2025,7,15), - True, + date(2025, 7, 15), + True, ), ( "MCH-003", "corrective", - date(2025,4,30), + date(2025, 4, 30), Decimal("150.00"), - date(2025,8,1), - False, + date(2025, 8, 1), + False, ), ( "MCH-002", "preventive", - date(2025,5,30), + date(2025, 5, 30), Decimal("150.00"), - date(2025,9,1), - True, + date(2025, 9, 1), + True, ), ( "MCH-002", "preventive", - date(2025,7,30), + date(2025, 7, 30), Decimal("100.00"), - date(2025,12,1), - True, - ) - ] - - + date(2025, 12, 1), + True, + ), + ] + input_df = spark.createDataFrame(maintenance_data, schema=schema) - + custom_options = { - "sample_fraction": None, - "limit": None, - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - } - + "sample_fraction": None, + "limit": None, + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + } profiler = DQProfiler(ws) filtered_df = profiler._sample(input_df, opts=custom_options) - stats, rules = profiler.profile(input_df,options=custom_options) + stats, rules = profiler.profile(input_df, options=custom_options) expected_rules = [ - DQProfile(name="is_not_null", - column="machine_id", - description=None, - parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), - DQProfile(name="is_not_null", - column="maintenance_type", - description=None, - parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), - DQProfile(name="is_not_null", - column="maintenance_date", - description=None, - parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + DQProfile( + name="is_not_null", + column="machine_id", + description=None, + parameters={ + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + }, + ), + DQProfile( + name="is_not_null", + column="maintenance_type", + description=None, + parameters={ + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + }, + ), + DQProfile( + name="is_not_null", + column="maintenance_date", + description=None, + parameters={ + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + }, + ), DQProfile( name="min_max", column="maintenance_date", description="Real min/max values were used", - parameters={"min": date(2025, 4, 29), - "max": date(2025, 7, 30), - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + parameters={ + "min": date(2025, 4, 29), + "max": date(2025, 7, 30), + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + }, ), - DQProfile(name="is_not_null", - column="cost", - description=None, - parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), DQProfile( - name="min_max", - column="cost", - parameters={"min": Decimal('100.00') , "max": Decimal('300.00'),"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, - description="Real min/max values were used" - ), - DQProfile(name="is_not_null", - column="next_scheduled_date", - description=None, - parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + name="is_not_null", + column="cost", + description=None, + parameters={ + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + }, + ), + DQProfile( + name="min_max", + column="cost", + parameters={ + "min": Decimal('100.00'), + "max": Decimal('300.00'), + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + }, + description="Real min/max values were used", + ), + DQProfile( + name="is_not_null", + column="next_scheduled_date", + description=None, + parameters={ + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + }, + ), DQProfile( name="min_max", column="next_scheduled_date", description="Real min/max values were used", - parameters={"min": date(2025, 7, 15), "max": date(2025, 12, 1),"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + parameters={ + "min": date(2025, 7, 15), + "max": date(2025, 12, 1), + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + }, ), - DQProfile(name="is_not_null", - column="safety_check_passed", - description=None, - parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), - - ] - - + DQProfile( + name="is_not_null", + column="safety_check_passed", + description=None, + parameters={ + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + }, + ), + ] + assert filtered_df.count() == 3 assert len(stats.keys()) > 0 assert rules == expected_rules - \ No newline at end of file diff --git a/tests/integration/test_rules_generator.py b/tests/integration/test_rules_generator.py index aec7b81d6..7b33aeefd 100644 --- a/tests/integration/test_rules_generator.py +++ b/tests/integration/test_rules_generator.py @@ -1,5 +1,6 @@ import datetime from decimal import Decimal +# from datetime import date, datetime, timezone from databricks.labs.dqx.profiler.generator import DQGenerator from databricks.labs.dqx.profiler.profiler import DQProfile @@ -122,3 +123,112 @@ def test_generate_dq_no_rules(ws): generator = DQGenerator(ws) expectations = generator.generate_dq_rules(None, level="warn") assert not expectations + +def test_generate_dq_rules_dataframe_filter(ws): + generator = DQGenerator(ws) + test_rules=[ + DQProfile(name="is_not_null", + column="machine_id", + description=None, + parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + + DQProfile( + name="min_max", + column="maintenance_date", + description="Real min/max values were used", + parameters={"min": datetime.date(2025, 4, 29), + "max": datetime.date(2025, 7, 30), + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + ), + DQProfile(name="is_not_null", + column="cost", + description=None, + parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + DQProfile( + name="min_max", + column="cost", + description="Real min/max values were used", + parameters={"min": Decimal('100.00') , "max": Decimal('300.00'),"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + ), + DQProfile(name="is_not_null", + column="next_scheduled_date", + description=None, + parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + DQProfile( + name="min_max", + column="next_scheduled_date", + description="Real min/max values were used", + parameters={"min": datetime.date(2025, 7, 15), "max": datetime.date(2025, 12, 1),"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + ), + DQProfile(name="is_not_null", + column="safety_check_passed", + description=None, + parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + + ] + expectations = generator.generate_dq_rules(test_rules) + + expected = [ + { + "check": {"function": "is_not_null", "arguments": {"column": "machine_id"}}, + "name": "machine_id_is_null", + "criticality": "Error", + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + }, + { + "check": { + "function": "min_max", + "arguments": {"column": "maintenance_date","min": datetime.date(2025, 4, 29), "max": datetime.date(2025, 7, 30)} + }, + "name": "maintenance_date_isnt_in_range", + "criticality": "Error", + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + }, + { + "check": { + "function": "is_not_null", + "arguments": {"column": "cost"}}, + + "name": "cost_is_null", + "criticality": "Error", + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + }, + { + "check": { + "function": "min_max", + "arguments": {"column": "cost","min": Decimal('100.00') , "max": Decimal('300.00')} + }, + "name": "cost_isnt_in_range", + "criticality": "Error", + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + }, + { + "check": { + "function": "is_not_null", + "arguments": {"column": "next_scheduled_date"}}, + + "name": "next_scheduled_date_is_null", + "criticality": "Error", + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + }, + { + "check": { + "function": "min_max", + "arguments": {"column": "next_scheduled_date","min": datetime.date(2025, 7, 15), "max": datetime.date(2025, 12, 1)} + }, + "name": "cost_isnt_in_range", + "criticality": "Error", + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + }, + { + "check": { + "function": "is_not_null", + "arguments": {"column": "safety_check_passed"}}, + + "name": "safety_check_passed_is_null", + "criticality": "Error", + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + } + + ] + assert expectations == expected diff --git a/tests/integration/test_save_checks_to_workspace_file.py b/tests/integration/test_save_checks_to_workspace_file.py index 93ced82b6..a01941e2c 100644 --- a/tests/integration/test_save_checks_to_workspace_file.py +++ b/tests/integration/test_save_checks_to_workspace_file.py @@ -18,10 +18,19 @@ { "criticality": "error", "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}}, + }, + { + "check": { + "function": "is_not_null", + "arguments": {"column": "next_scheduled_date"}}, + + "name": "next_scheduled_date_is_null", + "criticality": "Error", + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" } + ] - def test_save_checks_in_workspace_file_as_yaml(ws, spark, installation_ctx): installation_ctx.installation.save(installation_ctx.config) install_dir = installation_ctx.installation.install_folder() diff --git a/tests/unit/test_profiler.py b/tests/unit/test_profiler.py index 76a355b2f..62b0b14d6 100644 --- a/tests/unit/test_profiler.py +++ b/tests/unit/test_profiler.py @@ -2,7 +2,6 @@ from databricks.labs.dqx.profiler.profiler import DQProfiler - def test_get_columns_or_fields(): inp = T.StructType( [ @@ -29,4 +28,3 @@ def test_get_columns_or_fields(): T.StructField("ss1.s2.ns3", T.DateType()), ] assert fields == expected - From 7d9251db6f1abebf6b0df89b0535fc1b32e87110 Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Sat, 20 Sep 2025 20:14:41 -0300 Subject: [PATCH 08/34] Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter. --- src/databricks/labs/dqx/profiler/generator.py | 33 +++++------ src/databricks/labs/dqx/profiler/profiler.py | 29 ++++++---- tests/integration/test_profiler.py | 43 +++++--------- tests/integration/test_rules_generator.py | 58 +++++++++---------- 4 files changed, 76 insertions(+), 87 deletions(-) diff --git a/src/databricks/labs/dqx/profiler/generator.py b/src/databricks/labs/dqx/profiler/generator.py index 64d5e3654..1f38a50b4 100644 --- a/src/databricks/labs/dqx/profiler/generator.py +++ b/src/databricks/labs/dqx/profiler/generator.py @@ -55,10 +55,9 @@ def dq_generate_is_in(column: str, level: str = "error", **params: dict): A dictionary representing the data quality rule. """ return { - "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"]}}, + "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"],"filter": params.get("filter", None)}}, "name": f"{column}_other_value", - "criticality": level, - "filter": params.get("dataset_filter_expression", None), + "criticality": level } @staticmethod @@ -76,7 +75,7 @@ def dq_generate_min_max(column: str, level: str = "error", **params: dict): """ min_limit = params.get("min") max_limit = params.get("max") - filter = params.get("dataset_filter_expression", None) + filter = params.get("filter", None) if not isinstance(min_limit, int) or not isinstance(max_limit, int): return None # TODO handle timestamp and dates: https://github.com/databrickslabs/dqx/issues/71 @@ -89,11 +88,11 @@ def dq_generate_min_max(column: str, level: str = "error", **params: dict): "column": column, "min_limit": val_maybe_to_str(min_limit, include_sql_quotes=False), "max_limit": val_maybe_to_str(max_limit, include_sql_quotes=False), + "filter": params.get("filter", None) }, }, "name": f"{column}_isnt_in_range", - "criticality": level, - "filter": filter, + "criticality": level } if max_limit is not None: @@ -103,11 +102,11 @@ def dq_generate_min_max(column: str, level: str = "error", **params: dict): "arguments": { "column": column, "limit": val_maybe_to_str(max_limit, include_sql_quotes=False), + "filter": params.get("filter", None) }, }, "name": f"{column}_not_greater_than", - "criticality": level, - "filter": filter, + "criticality": level } if min_limit is not None: @@ -117,11 +116,11 @@ def dq_generate_min_max(column: str, level: str = "error", **params: dict): "arguments": { "column": column, "limit": val_maybe_to_str(min_limit, include_sql_quotes=False), + "filter": params.get("filter", None) }, }, "name": f"{column}_not_less_than", - "criticality": level, - "filter": filter, + "criticality": level } return None @@ -140,12 +139,11 @@ def dq_generate_is_not_null(column: str, level: str = "error", **params: dict): A dictionary representing the data quality rule. """ params = params or {} - filter = params.get("dataset_filter_expression", None) + filter = params.get("filter", None) return { - "check": {"function": "is_not_null", "arguments": {"column": column}}, + "check": {"function": "is_not_null", "arguments": {"column": column, "filter": params.get("filter", None)}}, "name": f"{column}_is_null", - "criticality": level, - "filter": filter, + "criticality": level } @staticmethod @@ -161,15 +159,14 @@ def dq_generate_is_not_null_or_empty(column: str, level: str = "error", **params Returns: A dictionary representing the data quality rule. """ - filter = params.get("dataset_filter_expression", None) + filter = params.get("filter", None) return { "check": { "function": "is_not_null_and_not_empty", - "arguments": {"column": column, "trim_strings": params.get("trim_strings", True)}, + "arguments": {"column": column, "trim_strings": params.get("trim_strings", True), "filter": params.get("filter", None)}, }, "name": f"{column}_is_null_or_empty", - "criticality": level, - "filter": filter, + "criticality": level } _checks_mapping = { diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 4598990c7..b89ab1bb9 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -31,7 +31,9 @@ class DQProfile: column: str description: str | None = None parameters: dict[str, Any] | None = None - + filter: str | None = None + + class DQProfiler(DQEngineBase): """Data Quality Profiler class to profile input data.""" @@ -53,7 +55,7 @@ def __init__(self, workspace_client: WorkspaceClient, spark: SparkSession | None "sample_fraction": 0.3, # fraction of data to sample (30%) "sample_seed": None, # seed for sampling "limit": 1000, # limit the number of samples - "dataset_filter_expression": None, # filter to apply to the dataset + "filter": None, # filter to apply to the dataset } @staticmethod @@ -112,6 +114,11 @@ def profile( return summary_stats, dq_rules self._profile(df, df_columns, dq_rules, options, summary_stats, total_count) + filter=options.get("filter",None) + if filter: + for rule in dq_rules: + rule.filter=filter + return summary_stats, dq_rules @@ -317,7 +324,7 @@ def _sample(df: DataFrame, opts: dict[str, Any]) -> DataFrame: sample_fraction = opts.get("sample_fraction", None) sample_seed = opts.get("sample_seed", None) limit = opts.get("limit", None) - filter = opts.get("dataset_filter_expression", None) + filter = opts.get("filter", None) if filter: df = df.filter(filter) @@ -371,8 +378,7 @@ def _calculate_metrics( """ max_nulls = opts.get("max_null_ratio", 0) trim_strings = opts.get("trim_strings", True) - filter = opts.get("dataset_filter_expression", None) - + dst = df.select(field_name).dropna() if typ == T.StringType() and trim_strings: col_name = dst.columns[0] @@ -390,13 +396,12 @@ def _calculate_metrics( name="is_not_null", column=field_name, description=f"Column {field_name} has {null_percentage * 100:.1f}% of null values " - f"(allowed {max_nulls * 100:.1f}%)", - parameters={"dataset_filter_expression": filter}, + f"(allowed {max_nulls * 100:.1f}%)", ) ) else: dq_rules.append( - DQProfile(name="is_not_null", column=field_name, parameters={"dataset_filter_expression": filter}) + DQProfile(name="is_not_null", column=field_name) ) if self._type_supports_distinct(typ): dst2 = dst.dropDuplicates() @@ -406,7 +411,7 @@ def _calculate_metrics( DQProfile( name="is_in", column=field_name, - parameters={"in": [row[0] for row in dst2.collect()], "dataset_filter_expression": filter}, + parameters={"in": [row[0] for row in dst2.collect()]}, ) ) if ( @@ -422,7 +427,7 @@ def _calculate_metrics( DQProfile( name="is_not_null_or_empty", column=field_name, - parameters={"trim_strings": trim_strings, "dataset_filter_expression": filter}, + parameters={"trim_strings": trim_strings}, ) ) if metrics["count_non_null"] > 0 and self._type_supports_min_max(typ): @@ -541,7 +546,7 @@ def _extract_min_max( if opts is None: opts = {} - filter = opts.get("dataset_filter_expression", None) + outlier_cols = opts.get("outlier_columns", []) column = dst.columns[0] if opts.get("remove_outliers", True) and ( @@ -579,7 +584,7 @@ def _extract_min_max( return DQProfile( name="min_max", column=col_name, - parameters={"min": min_limit, "max": max_limit, "dataset_filter_expression": filter}, + parameters={"min": min_limit, "max": max_limit}, description=descr, ) diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index bd3bab5ec..417e7c270 100644 --- a/tests/integration/test_profiler.py +++ b/tests/integration/test_profiler.py @@ -1339,7 +1339,7 @@ def test_profile_with_dataset_filter(spark, ws): custom_options = { "sample_fraction": None, "limit": None, - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", } profiler = DQProfiler(ws) @@ -1351,26 +1351,21 @@ def test_profile_with_dataset_filter(spark, ws): DQProfile( name="is_not_null", column="machine_id", - description=None, - parameters={ - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - }, + description=None, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + ), DQProfile( name="is_not_null", column="maintenance_type", description=None, - parameters={ - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - }, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), DQProfile( name="is_not_null", column="maintenance_date", description=None, - parameters={ - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - }, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), DQProfile( name="min_max", @@ -1378,17 +1373,14 @@ def test_profile_with_dataset_filter(spark, ws): description="Real min/max values were used", parameters={ "min": date(2025, 4, 29), - "max": date(2025, 7, 30), - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", - }, - ), + "max": date(2025, 7, 30)}, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + ), DQProfile( name="is_not_null", column="cost", description=None, - parameters={ - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - }, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), DQProfile( name="min_max", @@ -1396,17 +1388,15 @@ def test_profile_with_dataset_filter(spark, ws): parameters={ "min": Decimal('100.00'), "max": Decimal('300.00'), - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", }, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", description="Real min/max values were used", ), DQProfile( name="is_not_null", column="next_scheduled_date", description=None, - parameters={ - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - }, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), DQProfile( name="min_max", @@ -1414,17 +1404,14 @@ def test_profile_with_dataset_filter(spark, ws): description="Real min/max values were used", parameters={ "min": date(2025, 7, 15), - "max": date(2025, 12, 1), - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", - }, + "max": date(2025, 12, 1)}, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), DQProfile( name="is_not_null", column="safety_check_passed", description=None, - parameters={ - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - }, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", ), ] diff --git a/tests/integration/test_rules_generator.py b/tests/integration/test_rules_generator.py index 7b33aeefd..ec1c5c932 100644 --- a/tests/integration/test_rules_generator.py +++ b/tests/integration/test_rules_generator.py @@ -9,7 +9,7 @@ DQProfile( name="is_not_null", column="vendor_id", description="Column vendor_id has 0.3% of null values (allowed 1.0%)" ), - DQProfile(name="is_in", column="vendor_id", parameters={"in": ["1", "4", "2"]}), + DQProfile(name="is_in", column="vendor_id",parameters={"in": ["1", "4", "2"]}), DQProfile(name="is_not_null_or_empty", column="vendor_id", parameters={"trim_strings": True}), DQProfile( name="min_max", @@ -130,7 +130,7 @@ def test_generate_dq_rules_dataframe_filter(ws): DQProfile(name="is_not_null", column="machine_id", description=None, - parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + parameters={"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), DQProfile( name="min_max", @@ -138,96 +138,96 @@ def test_generate_dq_rules_dataframe_filter(ws): description="Real min/max values were used", parameters={"min": datetime.date(2025, 4, 29), "max": datetime.date(2025, 7, 30), - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, ), DQProfile(name="is_not_null", column="cost", description=None, - parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + parameters={"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), DQProfile( name="min_max", column="cost", description="Real min/max values were used", - parameters={"min": Decimal('100.00') , "max": Decimal('300.00'),"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + parameters={"min": Decimal('100.00') , "max": Decimal('300.00'),"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, ), DQProfile(name="is_not_null", column="next_scheduled_date", description=None, - parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + parameters={"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), DQProfile( name="min_max", column="next_scheduled_date", description="Real min/max values were used", - parameters={"min": datetime.date(2025, 7, 15), "max": datetime.date(2025, 12, 1),"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + parameters={"min": datetime.date(2025, 7, 15), "max": datetime.date(2025, 12, 1),"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, ), DQProfile(name="is_not_null", column="safety_check_passed", description=None, - parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + parameters={"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), ] expectations = generator.generate_dq_rules(test_rules) expected = [ { - "check": {"function": "is_not_null", "arguments": {"column": "machine_id"}}, + "check": {"function": "is_not_null", "arguments": {"column": "machine_id","filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}}, "name": "machine_id_is_null", - "criticality": "Error", - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + "criticality": "error" + }, { "check": { "function": "min_max", - "arguments": {"column": "maintenance_date","min": datetime.date(2025, 4, 29), "max": datetime.date(2025, 7, 30)} + "arguments": {"column": "maintenance_date","min": datetime.date(2025, 4, 29), "max": datetime.date(2025, 7, 30), "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} }, "name": "maintenance_date_isnt_in_range", - "criticality": "Error", - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + "criticality": "error", + }, { "check": { "function": "is_not_null", - "arguments": {"column": "cost"}}, + "arguments": {"column": "cost","filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}}, "name": "cost_is_null", - "criticality": "Error", - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + "criticality": "error", + }, { "check": { "function": "min_max", - "arguments": {"column": "cost","min": Decimal('100.00') , "max": Decimal('300.00')} + "arguments": {"column": "cost","min": Decimal('100.00') , "max": Decimal('300.00'),"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'",} }, "name": "cost_isnt_in_range", - "criticality": "Error", - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + "criticality": "error", + }, { "check": { "function": "is_not_null", - "arguments": {"column": "next_scheduled_date"}}, + "arguments": {"column": "next_scheduled_date","filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}}, "name": "next_scheduled_date_is_null", - "criticality": "Error", - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + "criticality": "error", + }, { "check": { "function": "min_max", - "arguments": {"column": "next_scheduled_date","min": datetime.date(2025, 7, 15), "max": datetime.date(2025, 12, 1)} + "arguments": {"column": "next_scheduled_date","min": datetime.date(2025, 7, 15), "max": datetime.date(2025, 12, 1), "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} }, "name": "cost_isnt_in_range", - "criticality": "Error", - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + "criticality": "error", + }, { "check": { "function": "is_not_null", - "arguments": {"column": "safety_check_passed"}}, + "arguments": {"column": "safety_check_passed", "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}}, "name": "safety_check_passed_is_null", - "criticality": "Error", - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + "criticality": "error", + } ] From 3b9a829f414115a0ef1589e1dcfed8fa4a4dd70b Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Sat, 20 Sep 2025 23:34:49 -0300 Subject: [PATCH 09/34] Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict. --- src/databricks/labs/dqx/profiler/generator.py | 56 ++++---- tests/integration/test_rules_generator.py | 126 +++++++++++------- 2 files changed, 104 insertions(+), 78 deletions(-) diff --git a/src/databricks/labs/dqx/profiler/generator.py b/src/databricks/labs/dqx/profiler/generator.py index 1f38a50b4..49cbda5df 100644 --- a/src/databricks/labs/dqx/profiler/generator.py +++ b/src/databricks/labs/dqx/profiler/generator.py @@ -25,14 +25,16 @@ def generate_dq_rules(self, profiles: list[DQProfile] | None = None, level: str if profiles is None: profiles = [] dq_rules = [] + for profile in profiles: rule_name = profile.name column = profile.column params = profile.parameters or {} + filter = profile.filter if rule_name not in self._checks_mapping: logger.info(f"No rule '{rule_name}' for column '{column}'. skipping...") continue - expr = self._checks_mapping[rule_name](column, level, **params) + expr = self._checks_mapping[rule_name](column, filter,level,**params) if expr: dq_rules.append(expr) @@ -42,7 +44,7 @@ def generate_dq_rules(self, profiles: list[DQProfile] | None = None, level: str return dq_rules @staticmethod - def dq_generate_is_in(column: str, level: str = "error", **params: dict): + def dq_generate_is_in(column: str, filter:str|None,level: str = "error",**params: dict ): """ Generates a data quality rule to check if a column's value is in a specified list. @@ -55,13 +57,13 @@ def dq_generate_is_in(column: str, level: str = "error", **params: dict): A dictionary representing the data quality rule. """ return { - "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"],"filter": params.get("filter", None)}}, + "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"]},"filter":filter}, "name": f"{column}_other_value", "criticality": level } @staticmethod - def dq_generate_min_max(column: str, level: str = "error", **params: dict): + def dq_generate_min_max(column: str, filter:str|None,level: str = "error", **params: dict): """ Generates a data quality rule to check if a column's value is within a specified range. @@ -75,7 +77,7 @@ def dq_generate_min_max(column: str, level: str = "error", **params: dict): """ min_limit = params.get("min") max_limit = params.get("max") - filter = params.get("filter", None) + if not isinstance(min_limit, int) or not isinstance(max_limit, int): return None # TODO handle timestamp and dates: https://github.com/databrickslabs/dqx/issues/71 @@ -87,12 +89,10 @@ def dq_generate_min_max(column: str, level: str = "error", **params: dict): "arguments": { "column": column, "min_limit": val_maybe_to_str(min_limit, include_sql_quotes=False), - "max_limit": val_maybe_to_str(max_limit, include_sql_quotes=False), - "filter": params.get("filter", None) - }, - }, - "name": f"{column}_isnt_in_range", - "criticality": level + "max_limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}, + "filter": filter}, + "name": f"{column}_isnt_in_range", + "criticality": level } if max_limit is not None: @@ -101,12 +101,10 @@ def dq_generate_min_max(column: str, level: str = "error", **params: dict): "function": "is_not_greater_than", "arguments": { "column": column, - "limit": val_maybe_to_str(max_limit, include_sql_quotes=False), - "filter": params.get("filter", None) - }, - }, - "name": f"{column}_not_greater_than", - "criticality": level + "limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}, + "filter": filter}, + "name": f"{column}_not_greater_than", + "criticality": level } if min_limit is not None: @@ -115,18 +113,16 @@ def dq_generate_min_max(column: str, level: str = "error", **params: dict): "function": "is_not_less_than", "arguments": { "column": column, - "limit": val_maybe_to_str(min_limit, include_sql_quotes=False), - "filter": params.get("filter", None) - }, - }, - "name": f"{column}_not_less_than", - "criticality": level + "limit": val_maybe_to_str(min_limit, include_sql_quotes=False)}, + "filter": filter}, + "name": f"{column}_not_less_than", + "criticality": level } return None @staticmethod - def dq_generate_is_not_null(column: str, level: str = "error", **params: dict): + def dq_generate_is_not_null(column: str, filter:str|None,level: str = "error", **params: dict): """ Generates a data quality rule to check if a column's value is not null. @@ -139,15 +135,15 @@ def dq_generate_is_not_null(column: str, level: str = "error", **params: dict): A dictionary representing the data quality rule. """ params = params or {} - filter = params.get("filter", None) + return { - "check": {"function": "is_not_null", "arguments": {"column": column, "filter": params.get("filter", None)}}, + "check": {"function": "is_not_null", "arguments": {"column": column},"filter":filter}, "name": f"{column}_is_null", "criticality": level } @staticmethod - def dq_generate_is_not_null_or_empty(column: str, level: str = "error", **params: dict): + def dq_generate_is_not_null_or_empty(column: str, filter:str|None,level: str = "error", **params: dict): """ Generates a data quality rule to check if a column's value is not null or empty. @@ -159,12 +155,12 @@ def dq_generate_is_not_null_or_empty(column: str, level: str = "error", **params Returns: A dictionary representing the data quality rule. """ - filter = params.get("filter", None) + return { "check": { "function": "is_not_null_and_not_empty", - "arguments": {"column": column, "trim_strings": params.get("trim_strings", True), "filter": params.get("filter", None)}, - }, + "arguments": {"column": column, "trim_strings": params.get("trim_strings", True)}, + "filter": filter}, "name": f"{column}_is_null_or_empty", "criticality": level } diff --git a/tests/integration/test_rules_generator.py b/tests/integration/test_rules_generator.py index ec1c5c932..d5ba230b6 100644 --- a/tests/integration/test_rules_generator.py +++ b/tests/integration/test_rules_generator.py @@ -44,14 +44,14 @@ def test_generate_dq_rules(ws): expectations = generator.generate_dq_rules(test_rules) expected = [ { - "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"}}, + "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"},"filter":None}, "name": "vendor_id_is_null", "criticality": "error", }, { "check": { "function": "is_in_list", - "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}, + "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]},"filter":None, }, "name": "vendor_id_other_value", "criticality": "error", @@ -59,7 +59,7 @@ def test_generate_dq_rules(ws): { "check": { "function": "is_not_null_and_not_empty", - "arguments": {"column": "vendor_id", "trim_strings": True}, + "arguments": {"column": "vendor_id", "trim_strings": True},"filter":None }, "name": "vendor_id_is_null_or_empty", "criticality": "error", @@ -68,6 +68,7 @@ def test_generate_dq_rules(ws): "check": { "function": "is_in_range", "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}, + "filter":None }, "name": "rate_code_id_isnt_in_range", "criticality": "error", @@ -81,14 +82,14 @@ def test_generate_dq_rules_warn(ws): expectations = generator.generate_dq_rules(test_rules, level="warn") expected = [ { - "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"}}, + "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"},"filter":None}, "name": "vendor_id_is_null", "criticality": "warn", }, { "check": { "function": "is_in_list", - "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}, + "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]},"filter":None, }, "name": "vendor_id_other_value", "criticality": "warn", @@ -96,7 +97,7 @@ def test_generate_dq_rules_warn(ws): { "check": { "function": "is_not_null_and_not_empty", - "arguments": {"column": "vendor_id", "trim_strings": True}, + "arguments": {"column": "vendor_id", "trim_strings": True},"filter":None, }, "name": "vendor_id_is_null_or_empty", "criticality": "warn", @@ -104,7 +105,7 @@ def test_generate_dq_rules_warn(ws): { "check": { "function": "is_in_range", - "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}, + "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265},"filter":None, }, "name": "rate_code_id_isnt_in_range", "criticality": "warn", @@ -130,105 +131,134 @@ def test_generate_dq_rules_dataframe_filter(ws): DQProfile(name="is_not_null", column="machine_id", description=None, - parameters={"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"), + DQProfile(name="is_in", + column="vendor_id", + parameters={"in": ["1", "4", "2"]}, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"), DQProfile( name="min_max", column="maintenance_date", description="Real min/max values were used", parameters={"min": datetime.date(2025, 4, 29), - "max": datetime.date(2025, 7, 30), - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + "max": datetime.date(2025, 7, 30)}, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", ), DQProfile(name="is_not_null", column="cost", description=None, - parameters={"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), DQProfile( name="min_max", column="cost", description="Real min/max values were used", - parameters={"min": Decimal('100.00') , "max": Decimal('300.00'),"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + parameters={"min": Decimal('100.00') , "max": Decimal('300.00')}, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), DQProfile(name="is_not_null", column="next_scheduled_date", description=None, - parameters={"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), - DQProfile( - name="min_max", - column="next_scheduled_date", - description="Real min/max values were used", - parameters={"min": datetime.date(2025, 7, 15), "max": datetime.date(2025, 12, 1),"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, - ), + filter= "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"), + DQProfile(name="is_not_null", column="safety_check_passed", description=None, - parameters={"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), + + DQProfile(name="is_not_null_or_empty", + column="vendor_id", + parameters={"trim_strings": True}, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), ] expectations = generator.generate_dq_rules(test_rules) - + expected = [ { - "check": {"function": "is_not_null", "arguments": {"column": "machine_id","filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}}, + "check": {"function": "is_not_null", + "arguments": {"column": "machine_id"}, + "filter":"machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, "name": "machine_id_is_null", "criticality": "error" }, { - "check": { - "function": "min_max", - "arguments": {"column": "maintenance_date","min": datetime.date(2025, 4, 29), "max": datetime.date(2025, 7, 30), "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} - }, - "name": "maintenance_date_isnt_in_range", + 'check': {"function": "is_in_list", + "arguments": {"allowed": ['1','4','2'],"column": "vendor_id"}, + "filter":"machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, "criticality": "error", - + "name": "vendor_id_other_value" }, + + # { + # "check": { + # "function": "min_max", + # "arguments": {"column": "maintenance_date","min": datetime.date(2025, 4, 29), "max": datetime.date(2025, 7, 30)}, + # "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + # }, + # "name": "maintenance_date_isnt_in_range", + # "criticality": "error", + + # }, { "check": { "function": "is_not_null", - "arguments": {"column": "cost","filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}}, + "arguments": {"column": "cost"}, + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, "name": "cost_is_null", "criticality": "error", }, - { - "check": { - "function": "min_max", - "arguments": {"column": "cost","min": Decimal('100.00') , "max": Decimal('300.00'),"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'",} - }, - "name": "cost_isnt_in_range", - "criticality": "error", + # { + # "check": { + # "function": "min_max", + # "arguments": {"column": "cost","min": Decimal('100.00') , "max": Decimal('300.00')}, + # "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + # }, + # "name": "cost_isnt_in_range", + # "criticality": "error", - }, + # }, { "check": { "function": "is_not_null", - "arguments": {"column": "next_scheduled_date","filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}}, + "arguments": {"column": "next_scheduled_date"}, + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, "name": "next_scheduled_date_is_null", "criticality": "error", }, - { - "check": { - "function": "min_max", - "arguments": {"column": "next_scheduled_date","min": datetime.date(2025, 7, 15), "max": datetime.date(2025, 12, 1), "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} - }, - "name": "cost_isnt_in_range", - "criticality": "error", + # { + # "check": { + # "function": "min_max", + # "arguments": {"column": "next_scheduled_date","min": datetime.date(2025, 7, 15), "max": datetime.date(2025, 12, 1)}, + # "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + # }, + # "name": "cost_isnt_in_range", + # "criticality": "error", - }, + # }, { "check": { "function": "is_not_null", - "arguments": {"column": "safety_check_passed", "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}}, + "arguments": {"column": "safety_check_passed"}, + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, "name": "safety_check_passed_is_null", "criticality": "error", - } + }, + { + "check": { + "function": "is_not_null_and_not_empty", + "arguments": {"column": "vendor_id", "trim_strings": True}, + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + "name": "vendor_id_is_null_or_empty", + "criticality": "error" + } ] assert expectations == expected From c79b9bc9276bcc972f0c54c17bca682064f456ce Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Mon, 22 Sep 2025 10:49:50 -0300 Subject: [PATCH 10/34] Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig. --- src/databricks/labs/dqx/profiler/generator.py | 22 ++-- .../test_load_checks_from_uc_volume.py | 24 +++- tests/integration/test_profiler.py | 22 ++-- tests/integration/test_rules_generator.py | 122 ++++++------------ .../test_save_and_load_checks_from_table.py | 8 +- .../test_save_checks_to_workspace_file.py | 15 ++- 6 files changed, 106 insertions(+), 107 deletions(-) diff --git a/src/databricks/labs/dqx/profiler/generator.py b/src/databricks/labs/dqx/profiler/generator.py index 49cbda5df..85562d80f 100644 --- a/src/databricks/labs/dqx/profiler/generator.py +++ b/src/databricks/labs/dqx/profiler/generator.py @@ -57,7 +57,8 @@ def dq_generate_is_in(column: str, filter:str|None,level: str = "error",**params A dictionary representing the data quality rule. """ return { - "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"]},"filter":filter}, + "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"]}}, + "filter":filter, "name": f"{column}_other_value", "criticality": level } @@ -89,8 +90,8 @@ def dq_generate_min_max(column: str, filter:str|None,level: str = "error", **par "arguments": { "column": column, "min_limit": val_maybe_to_str(min_limit, include_sql_quotes=False), - "max_limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}, - "filter": filter}, + "max_limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}}, + "filter": filter, "name": f"{column}_isnt_in_range", "criticality": level } @@ -101,8 +102,8 @@ def dq_generate_min_max(column: str, filter:str|None,level: str = "error", **par "function": "is_not_greater_than", "arguments": { "column": column, - "limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}, - "filter": filter}, + "limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}}, + "filter": filter, "name": f"{column}_not_greater_than", "criticality": level } @@ -113,8 +114,8 @@ def dq_generate_min_max(column: str, filter:str|None,level: str = "error", **par "function": "is_not_less_than", "arguments": { "column": column, - "limit": val_maybe_to_str(min_limit, include_sql_quotes=False)}, - "filter": filter}, + "limit": val_maybe_to_str(min_limit, include_sql_quotes=False)}}, + "filter": filter, "name": f"{column}_not_less_than", "criticality": level } @@ -137,7 +138,8 @@ def dq_generate_is_not_null(column: str, filter:str|None,level: str = "error", * params = params or {} return { - "check": {"function": "is_not_null", "arguments": {"column": column},"filter":filter}, + "check": {"function": "is_not_null", "arguments": {"column": column}}, + "filter":filter, "name": f"{column}_is_null", "criticality": level } @@ -159,8 +161,8 @@ def dq_generate_is_not_null_or_empty(column: str, filter:str|None,level: str = " return { "check": { "function": "is_not_null_and_not_empty", - "arguments": {"column": column, "trim_strings": params.get("trim_strings", True)}, - "filter": filter}, + "arguments": {"column": column, "trim_strings": params.get("trim_strings", True)}}, + "filter": filter, "name": f"{column}_is_null_or_empty", "criticality": level } diff --git a/tests/integration/test_load_checks_from_uc_volume.py b/tests/integration/test_load_checks_from_uc_volume.py index 81d50a5db..2eb7046cd 100644 --- a/tests/integration/test_load_checks_from_uc_volume.py +++ b/tests/integration/test_load_checks_from_uc_volume.py @@ -9,14 +9,34 @@ TEST_CHECKS = [ { "criticality": "error", - "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}}, + "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}, "filter":None}, + }, + { + "check": { + "function": "is_not_null", + "arguments": {"column": "next_scheduled_date"}, + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + + "name": "next_scheduled_date_is_null", + "criticality": "Error", + } ] EXPECTED_CHECKS = [ { "criticality": "error", - "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}}, + "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {},"filter":None}, + }, + { + "check": { + "function": "is_not_null", + "arguments": {"column": "next_scheduled_date"}, + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + + "name": "next_scheduled_date_is_null", + "criticality": "Error", + } ] diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index 417e7c270..848b342c4 100644 --- a/tests/integration/test_profiler.py +++ b/tests/integration/test_profiler.py @@ -67,18 +67,19 @@ def test_profiler(spark, ws): stats, rules = profiler.profile(inp_df, options={"sample_fraction": None}) expected_rules = [ - DQProfile(name="is_not_null", column="t1", description=None, parameters=None), + DQProfile(name="is_not_null", column="t1", description=None, parameters=None, filter=None), DQProfile( - name="min_max", column="t1", description="Real min/max values were used", parameters={"min": 1, "max": 3} + name="min_max", column="t1", description="Real min/max values were used", parameters={"min": 1, "max": 3}, filter=None ), - DQProfile(name='is_not_null', column='d1', description=None, parameters=None), + DQProfile(name='is_not_null', column='d1', description=None, parameters=None, filter=None), DQProfile( name='min_max', column='d1', description='Real min/max values were used', parameters={'max': Decimal('333323.00'), 'min': Decimal('1.23')}, + filter=None ), - DQProfile(name='is_not_null_or_empty', column='t2', description=None, parameters={'trim_strings': True}), + DQProfile(name='is_not_null_or_empty', column='t2', description=None, parameters={'trim_strings': True},filter=None), DQProfile(name="is_not_null", column="s1.ns1", description=None, parameters=None), DQProfile( name="min_max", @@ -88,16 +89,18 @@ def test_profiler(spark, ws): "min": datetime(2023, 1, 6, 0, 0, tzinfo=timezone.utc), "max": datetime(2023, 1, 9, 0, 0, tzinfo=timezone.utc), }, + filter=None ), - DQProfile(name="is_not_null", column="s1.s2.ns2", description=None, parameters=None), - DQProfile(name="is_not_null", column="s1.s2.ns3", description=None, parameters=None), + DQProfile(name="is_not_null", column="s1.s2.ns2", description=None, parameters=None,filter=None), + DQProfile(name="is_not_null", column="s1.s2.ns3", description=None, parameters=None,filter=None), DQProfile( name="min_max", column="s1.s2.ns3", description="Real min/max values were used", parameters={"min": date(2023, 1, 6), "max": date(2023, 1, 8)}, + filter=None ), - DQProfile(name="is_not_null", column="b1", description=None, parameters=None), + DQProfile(name="is_not_null", column="b1", description=None, parameters=None,filter=None), ] print(stats) assert len(stats.keys()) > 0 @@ -1304,7 +1307,7 @@ def test_profile_with_dataset_filter(spark, ws): "MCH-002", "preventive", date(2025, 4, 29), - Decimal("300.00"), + Decimal("300.50"), date(2025, 7, 15), True, ), @@ -1338,6 +1341,7 @@ def test_profile_with_dataset_filter(spark, ws): custom_options = { "sample_fraction": None, + "round":False, "limit": None, "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", } @@ -1387,7 +1391,7 @@ def test_profile_with_dataset_filter(spark, ws): column="cost", parameters={ "min": Decimal('100.00'), - "max": Decimal('300.00'), + "max": Decimal('300.50'), }, filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", description="Real min/max values were used", diff --git a/tests/integration/test_rules_generator.py b/tests/integration/test_rules_generator.py index d5ba230b6..93db395d2 100644 --- a/tests/integration/test_rules_generator.py +++ b/tests/integration/test_rules_generator.py @@ -44,32 +44,33 @@ def test_generate_dq_rules(ws): expectations = generator.generate_dq_rules(test_rules) expected = [ { - "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"},"filter":None}, + "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"}}, + "filter":None, "name": "vendor_id_is_null", "criticality": "error", }, { "check": { "function": "is_in_list", - "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]},"filter":None, - }, + "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}}, + "filter":None, "name": "vendor_id_other_value", "criticality": "error", }, { "check": { "function": "is_not_null_and_not_empty", - "arguments": {"column": "vendor_id", "trim_strings": True},"filter":None + "arguments": {"column": "vendor_id", "trim_strings": True} }, + "filter":None, "name": "vendor_id_is_null_or_empty", "criticality": "error", }, { "check": { "function": "is_in_range", - "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}, - "filter":None - }, + "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}}, + "filter":None, "name": "rate_code_id_isnt_in_range", "criticality": "error", }, @@ -82,31 +83,32 @@ def test_generate_dq_rules_warn(ws): expectations = generator.generate_dq_rules(test_rules, level="warn") expected = [ { - "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"},"filter":None}, + "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"}}, + "filter":None, "name": "vendor_id_is_null", "criticality": "warn", }, { "check": { "function": "is_in_list", - "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]},"filter":None, - }, + "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}}, + "filter":None, "name": "vendor_id_other_value", "criticality": "warn", }, { "check": { "function": "is_not_null_and_not_empty", - "arguments": {"column": "vendor_id", "trim_strings": True},"filter":None, - }, + "arguments": {"column": "vendor_id", "trim_strings": True}}, + "filter":None, "name": "vendor_id_is_null_or_empty", "criticality": "warn", }, { "check": { "function": "is_in_range", - "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265},"filter":None, - }, + "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}}, + "filter":None, "name": "rate_code_id_isnt_in_range", "criticality": "warn", }, @@ -137,39 +139,25 @@ def test_generate_dq_rules_dataframe_filter(ws): column="vendor_id", parameters={"in": ["1", "4", "2"]}, filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"), - DQProfile( - name="min_max", - column="maintenance_date", - description="Real min/max values were used", - parameters={"min": datetime.date(2025, 4, 29), - "max": datetime.date(2025, 7, 30)}, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", - ), + DQProfile(name="is_not_null", column="cost", description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), - DQProfile( - name="min_max", - column="cost", - description="Real min/max values were used", - parameters={"min": Decimal('100.00') , "max": Decimal('300.00')}, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - ), + ), + DQProfile(name="is_not_null", column="next_scheduled_date", description=None, - filter= "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"), + ), DQProfile(name="is_not_null", column="safety_check_passed", description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), + ), DQProfile(name="is_not_null_or_empty", column="vendor_id", - parameters={"trim_strings": True}, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), + parameters={"trim_strings": True}), ] expectations = generator.generate_dq_rules(test_rules) @@ -177,75 +165,47 @@ def test_generate_dq_rules_dataframe_filter(ws): expected = [ { "check": {"function": "is_not_null", - "arguments": {"column": "machine_id"}, - "filter":"machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + "arguments": {"column": "machine_id"}}, + "filter":"machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", "name": "machine_id_is_null", "criticality": "error" }, { 'check': {"function": "is_in_list", - "arguments": {"allowed": ['1','4','2'],"column": "vendor_id"}, - "filter":"machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, - "criticality": "error", + "arguments": {"allowed": ['1','4','2'],"column": "vendor_id"}}, + + "filter":"machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + "criticality": "error", "name": "vendor_id_other_value" - }, - - # { - # "check": { - # "function": "min_max", - # "arguments": {"column": "maintenance_date","min": datetime.date(2025, 4, 29), "max": datetime.date(2025, 7, 30)}, - # "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - # }, - # "name": "maintenance_date_isnt_in_range", - # "criticality": "error", - - # }, + }, { "check": { "function": "is_not_null", - "arguments": {"column": "cost"}, - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, - + "arguments": {"column": "cost"}}, + + "filter": None, "name": "cost_is_null", "criticality": "error", }, - # { - # "check": { - # "function": "min_max", - # "arguments": {"column": "cost","min": Decimal('100.00') , "max": Decimal('300.00')}, - # "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - # }, - # "name": "cost_isnt_in_range", - # "criticality": "error", - - # }, + { "check": { "function": "is_not_null", - "arguments": {"column": "next_scheduled_date"}, - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, - + "arguments": {"column": "next_scheduled_date"}}, + + "filter": None, "name": "next_scheduled_date_is_null", "criticality": "error", }, - # { - # "check": { - # "function": "min_max", - # "arguments": {"column": "next_scheduled_date","min": datetime.date(2025, 7, 15), "max": datetime.date(2025, 12, 1)}, - # "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - # }, - # "name": "cost_isnt_in_range", - # "criticality": "error", - - # }, + { "check": { "function": "is_not_null", - "arguments": {"column": "safety_check_passed"}, - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + "arguments": {"column": "safety_check_passed"}}, + "filter": None, "name": "safety_check_passed_is_null", "criticality": "error", @@ -254,8 +214,8 @@ def test_generate_dq_rules_dataframe_filter(ws): { "check": { "function": "is_not_null_and_not_empty", - "arguments": {"column": "vendor_id", "trim_strings": True}, - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + "arguments": {"column": "vendor_id", "trim_strings": True}}, + "filter": None, "name": "vendor_id_is_null_or_empty", "criticality": "error" } diff --git a/tests/integration/test_save_and_load_checks_from_table.py b/tests/integration/test_save_and_load_checks_from_table.py index 391c8c3ed..ac514a604 100644 --- a/tests/integration/test_save_and_load_checks_from_table.py +++ b/tests/integration/test_save_and_load_checks_from_table.py @@ -25,6 +25,7 @@ "name": "column_not_less_than", "criticality": "warn", "check": {"function": "is_not_less_than", "arguments": {"column": "col_2", "limit": 1}}, + "filter": "Col_3 >1", "user_metadata": {"check_type": "standardization", "check_owner": "someone_else@email.com"}, }, { @@ -52,8 +53,8 @@ "criticality": "warn", "check": { "function": "is_not_less_than", - "arguments": {"column": "col_2", "limit": 1}, - }, + "arguments": {"column": "col_2", "limit": 1}}, + "filter": "Col_3 >1", "user_metadata": {"check_type": "standardization", "check_owner": "someone_else@email.com"}, }, { @@ -64,6 +65,7 @@ ] + def test_load_checks_when_checks_table_does_not_exist(ws, make_schema, make_random, spark): catalog_name = "main" schema_name = make_schema(catalog_name=catalog_name).name @@ -128,7 +130,7 @@ def test_save_checks_to_table_with_unresolved_for_each_column(ws, make_schema, m "name": "column_not_less_than", "criticality": "warn", "check": {"function": "is_not_less_than", "arguments": {"limit": "1", "column": "\"col_2\""}}, - "filter": None, + "filter": "Col_3 >1", "run_config_name": "default", "user_metadata": {"check_type": "standardization", "check_owner": "someone_else@email.com"}, }, diff --git a/tests/integration/test_save_checks_to_workspace_file.py b/tests/integration/test_save_checks_to_workspace_file.py index a01941e2c..b27f47213 100644 --- a/tests/integration/test_save_checks_to_workspace_file.py +++ b/tests/integration/test_save_checks_to_workspace_file.py @@ -22,11 +22,22 @@ { "check": { "function": "is_not_null", - "arguments": {"column": "next_scheduled_date"}}, + "arguments": {"column": "next_scheduled_date"}, + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, "name": "next_scheduled_date_is_null", "criticality": "Error", - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + + }, + { + "check": { + "function": "is_not_null", + "arguments": {"column": "cost"}, + "filter": None}, + + "name": "cost_is_null", + "criticality": "error", + } ] From 287df6d78b52cf667374f32f7cf5d0bd7f405528 Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Tue, 23 Sep 2025 16:30:10 -0300 Subject: [PATCH 11/34] Formatting code to submit pull request --- src/databricks/labs/dqx/profiler/generator.py | 64 ++++---- src/databricks/labs/dqx/profiler/profiler.py | 20 +-- .../test_load_checks_from_uc_volume.py | 18 +-- tests/integration/test_profiler.py | 51 +++--- tests/integration/test_rules_generator.py | 150 ++++++++---------- .../test_save_and_load_checks_from_table.py | 5 +- .../test_save_checks_to_workspace_file.py | 22 +-- 7 files changed, 149 insertions(+), 181 deletions(-) diff --git a/src/databricks/labs/dqx/profiler/generator.py b/src/databricks/labs/dqx/profiler/generator.py index 85562d80f..7a3566b5f 100644 --- a/src/databricks/labs/dqx/profiler/generator.py +++ b/src/databricks/labs/dqx/profiler/generator.py @@ -25,7 +25,7 @@ def generate_dq_rules(self, profiles: list[DQProfile] | None = None, level: str if profiles is None: profiles = [] dq_rules = [] - + for profile in profiles: rule_name = profile.name column = profile.column @@ -34,7 +34,7 @@ def generate_dq_rules(self, profiles: list[DQProfile] | None = None, level: str if rule_name not in self._checks_mapping: logger.info(f"No rule '{rule_name}' for column '{column}'. skipping...") continue - expr = self._checks_mapping[rule_name](column, filter,level,**params) + expr = self._checks_mapping[rule_name](column, filter, level, **params) if expr: dq_rules.append(expr) @@ -44,7 +44,7 @@ def generate_dq_rules(self, profiles: list[DQProfile] | None = None, level: str return dq_rules @staticmethod - def dq_generate_is_in(column: str, filter:str|None,level: str = "error",**params: dict ): + def dq_generate_is_in(column: str, filter: str | None, level: str = "error", **params: dict): """ Generates a data quality rule to check if a column's value is in a specified list. @@ -58,13 +58,13 @@ def dq_generate_is_in(column: str, filter:str|None,level: str = "error",**params """ return { "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"]}}, - "filter":filter, + "filter": filter, "name": f"{column}_other_value", - "criticality": level + "criticality": level, } @staticmethod - def dq_generate_min_max(column: str, filter:str|None,level: str = "error", **params: dict): + def dq_generate_min_max(column: str, filter: str | None, level: str = "error", **params: dict): """ Generates a data quality rule to check if a column's value is within a specified range. @@ -78,7 +78,6 @@ def dq_generate_min_max(column: str, filter:str|None,level: str = "error", **par """ min_limit = params.get("min") max_limit = params.get("max") - if not isinstance(min_limit, int) or not isinstance(max_limit, int): return None # TODO handle timestamp and dates: https://github.com/databrickslabs/dqx/issues/71 @@ -90,40 +89,40 @@ def dq_generate_min_max(column: str, filter:str|None,level: str = "error", **par "arguments": { "column": column, "min_limit": val_maybe_to_str(min_limit, include_sql_quotes=False), - "max_limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}}, - "filter": filter, - "name": f"{column}_isnt_in_range", - "criticality": level + "max_limit": val_maybe_to_str(max_limit, include_sql_quotes=False), + }, + }, + "filter": filter, + "name": f"{column}_isnt_in_range", + "criticality": level, } if max_limit is not None: return { "check": { "function": "is_not_greater_than", - "arguments": { - "column": column, - "limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}}, - "filter": filter, - "name": f"{column}_not_greater_than", - "criticality": level + "arguments": {"column": column, "limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}, + }, + "filter": filter, + "name": f"{column}_not_greater_than", + "criticality": level, } if min_limit is not None: return { "check": { "function": "is_not_less_than", - "arguments": { - "column": column, - "limit": val_maybe_to_str(min_limit, include_sql_quotes=False)}}, - "filter": filter, - "name": f"{column}_not_less_than", - "criticality": level + "arguments": {"column": column, "limit": val_maybe_to_str(min_limit, include_sql_quotes=False)}, + }, + "filter": filter, + "name": f"{column}_not_less_than", + "criticality": level, } return None @staticmethod - def dq_generate_is_not_null(column: str, filter:str|None,level: str = "error", **params: dict): + def dq_generate_is_not_null(column: str, filter: str | None, level: str = "error", **params: dict): """ Generates a data quality rule to check if a column's value is not null. @@ -136,16 +135,16 @@ def dq_generate_is_not_null(column: str, filter:str|None,level: str = "error", * A dictionary representing the data quality rule. """ params = params or {} - + return { "check": {"function": "is_not_null", "arguments": {"column": column}}, - "filter":filter, + "filter": filter, "name": f"{column}_is_null", - "criticality": level + "criticality": level, } @staticmethod - def dq_generate_is_not_null_or_empty(column: str, filter:str|None,level: str = "error", **params: dict): + def dq_generate_is_not_null_or_empty(column: str, filter: str | None, level: str = "error", **params: dict): """ Generates a data quality rule to check if a column's value is not null or empty. @@ -157,14 +156,15 @@ def dq_generate_is_not_null_or_empty(column: str, filter:str|None,level: str = " Returns: A dictionary representing the data quality rule. """ - + return { "check": { "function": "is_not_null_and_not_empty", - "arguments": {"column": column, "trim_strings": params.get("trim_strings", True)}}, - "filter": filter, + "arguments": {"column": column, "trim_strings": params.get("trim_strings", True)}, + }, + "filter": filter, "name": f"{column}_is_null_or_empty", - "criticality": level + "criticality": level, } _checks_mapping = { diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index b89ab1bb9..7b744d818 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -32,8 +32,7 @@ class DQProfile: description: str | None = None parameters: dict[str, Any] | None = None filter: str | None = None - - + class DQProfiler(DQEngineBase): """Data Quality Profiler class to profile input data.""" @@ -114,11 +113,10 @@ def profile( return summary_stats, dq_rules self._profile(df, df_columns, dq_rules, options, summary_stats, total_count) - filter=options.get("filter",None) - if filter: + filter = options.get("filter", None) + if filter: for rule in dq_rules: - rule.filter=filter - + rule.filter = filter return summary_stats, dq_rules @@ -378,7 +376,7 @@ def _calculate_metrics( """ max_nulls = opts.get("max_null_ratio", 0) trim_strings = opts.get("trim_strings", True) - + dst = df.select(field_name).dropna() if typ == T.StringType() and trim_strings: col_name = dst.columns[0] @@ -396,13 +394,11 @@ def _calculate_metrics( name="is_not_null", column=field_name, description=f"Column {field_name} has {null_percentage * 100:.1f}% of null values " - f"(allowed {max_nulls * 100:.1f}%)", + f"(allowed {max_nulls * 100:.1f}%)", ) ) else: - dq_rules.append( - DQProfile(name="is_not_null", column=field_name) - ) + dq_rules.append(DQProfile(name="is_not_null", column=field_name)) if self._type_supports_distinct(typ): dst2 = dst.dropDuplicates() cnt = dst2.count() @@ -546,7 +542,7 @@ def _extract_min_max( if opts is None: opts = {} - + outlier_cols = opts.get("outlier_columns", []) column = dst.columns[0] if opts.get("remove_outliers", True) and ( diff --git a/tests/integration/test_load_checks_from_uc_volume.py b/tests/integration/test_load_checks_from_uc_volume.py index 2eb7046cd..6659d78f4 100644 --- a/tests/integration/test_load_checks_from_uc_volume.py +++ b/tests/integration/test_load_checks_from_uc_volume.py @@ -9,35 +9,33 @@ TEST_CHECKS = [ { "criticality": "error", - "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}, "filter":None}, + "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}, "filter": None}, }, { "check": { "function": "is_not_null", "arguments": {"column": "next_scheduled_date"}, - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, - + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + }, "name": "next_scheduled_date_is_null", "criticality": "Error", - - } + }, ] EXPECTED_CHECKS = [ { "criticality": "error", - "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {},"filter":None}, + "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}, "filter": None}, }, { "check": { "function": "is_not_null", "arguments": {"column": "next_scheduled_date"}, - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, - + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + }, "name": "next_scheduled_date_is_null", "criticality": "Error", - - } + }, ] diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index 848b342c4..2c5f1d35f 100644 --- a/tests/integration/test_profiler.py +++ b/tests/integration/test_profiler.py @@ -69,7 +69,11 @@ def test_profiler(spark, ws): expected_rules = [ DQProfile(name="is_not_null", column="t1", description=None, parameters=None, filter=None), DQProfile( - name="min_max", column="t1", description="Real min/max values were used", parameters={"min": 1, "max": 3}, filter=None + name="min_max", + column="t1", + description="Real min/max values were used", + parameters={"min": 1, "max": 3}, + filter=None, ), DQProfile(name='is_not_null', column='d1', description=None, parameters=None, filter=None), DQProfile( @@ -77,9 +81,11 @@ def test_profiler(spark, ws): column='d1', description='Real min/max values were used', parameters={'max': Decimal('333323.00'), 'min': Decimal('1.23')}, - filter=None + filter=None, + ), + DQProfile( + name='is_not_null_or_empty', column='t2', description=None, parameters={'trim_strings': True}, filter=None ), - DQProfile(name='is_not_null_or_empty', column='t2', description=None, parameters={'trim_strings': True},filter=None), DQProfile(name="is_not_null", column="s1.ns1", description=None, parameters=None), DQProfile( name="min_max", @@ -89,18 +95,18 @@ def test_profiler(spark, ws): "min": datetime(2023, 1, 6, 0, 0, tzinfo=timezone.utc), "max": datetime(2023, 1, 9, 0, 0, tzinfo=timezone.utc), }, - filter=None + filter=None, ), - DQProfile(name="is_not_null", column="s1.s2.ns2", description=None, parameters=None,filter=None), - DQProfile(name="is_not_null", column="s1.s2.ns3", description=None, parameters=None,filter=None), + DQProfile(name="is_not_null", column="s1.s2.ns2", description=None, parameters=None, filter=None), + DQProfile(name="is_not_null", column="s1.s2.ns3", description=None, parameters=None, filter=None), DQProfile( name="min_max", column="s1.s2.ns3", description="Real min/max values were used", parameters={"min": date(2023, 1, 6), "max": date(2023, 1, 8)}, - filter=None + filter=None, ), - DQProfile(name="is_not_null", column="b1", description=None, parameters=None,filter=None), + DQProfile(name="is_not_null", column="b1", description=None, parameters=None, filter=None), ] print(stats) assert len(stats.keys()) > 0 @@ -1341,7 +1347,7 @@ def test_profile_with_dataset_filter(spark, ws): custom_options = { "sample_fraction": None, - "round":False, + "round": False, "limit": None, "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", } @@ -1355,36 +1361,33 @@ def test_profile_with_dataset_filter(spark, ws): DQProfile( name="is_not_null", column="machine_id", - description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - + description=None, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", ), DQProfile( name="is_not_null", column="maintenance_type", description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", ), DQProfile( name="is_not_null", column="maintenance_date", description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", ), DQProfile( name="min_max", column="maintenance_date", description="Real min/max values were used", - parameters={ - "min": date(2025, 4, 29), - "max": date(2025, 7, 30)}, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - ), + parameters={"min": date(2025, 4, 29), "max": date(2025, 7, 30)}, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + ), DQProfile( name="is_not_null", column="cost", description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", ), DQProfile( name="min_max", @@ -1400,16 +1403,14 @@ def test_profile_with_dataset_filter(spark, ws): name="is_not_null", column="next_scheduled_date", description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", ), DQProfile( name="min_max", column="next_scheduled_date", description="Real min/max values were used", - parameters={ - "min": date(2025, 7, 15), - "max": date(2025, 12, 1)}, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + parameters={"min": date(2025, 7, 15), "max": date(2025, 12, 1)}, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", ), DQProfile( name="is_not_null", diff --git a/tests/integration/test_rules_generator.py b/tests/integration/test_rules_generator.py index 93db395d2..9636098d2 100644 --- a/tests/integration/test_rules_generator.py +++ b/tests/integration/test_rules_generator.py @@ -1,5 +1,6 @@ import datetime from decimal import Decimal + # from datetime import date, datetime, timezone from databricks.labs.dqx.profiler.generator import DQGenerator @@ -9,7 +10,7 @@ DQProfile( name="is_not_null", column="vendor_id", description="Column vendor_id has 0.3% of null values (allowed 1.0%)" ), - DQProfile(name="is_in", column="vendor_id",parameters={"in": ["1", "4", "2"]}), + DQProfile(name="is_in", column="vendor_id", parameters={"in": ["1", "4", "2"]}), DQProfile(name="is_not_null_or_empty", column="vendor_id", parameters={"trim_strings": True}), DQProfile( name="min_max", @@ -45,32 +46,31 @@ def test_generate_dq_rules(ws): expected = [ { "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"}}, - "filter":None, + "filter": None, "name": "vendor_id_is_null", "criticality": "error", }, { - "check": { - "function": "is_in_list", - "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}}, - "filter":None, + "check": {"function": "is_in_list", "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}}, + "filter": None, "name": "vendor_id_other_value", "criticality": "error", }, { "check": { "function": "is_not_null_and_not_empty", - "arguments": {"column": "vendor_id", "trim_strings": True} + "arguments": {"column": "vendor_id", "trim_strings": True}, }, - "filter":None, + "filter": None, "name": "vendor_id_is_null_or_empty", "criticality": "error", }, { "check": { "function": "is_in_range", - "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}}, - "filter":None, + "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}, + }, + "filter": None, "name": "rate_code_id_isnt_in_range", "criticality": "error", }, @@ -84,31 +84,31 @@ def test_generate_dq_rules_warn(ws): expected = [ { "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"}}, - "filter":None, + "filter": None, "name": "vendor_id_is_null", "criticality": "warn", }, { - "check": { - "function": "is_in_list", - "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}}, - "filter":None, + "check": {"function": "is_in_list", "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}}, + "filter": None, "name": "vendor_id_other_value", "criticality": "warn", }, { "check": { "function": "is_not_null_and_not_empty", - "arguments": {"column": "vendor_id", "trim_strings": True}}, - "filter":None, + "arguments": {"column": "vendor_id", "trim_strings": True}, + }, + "filter": None, "name": "vendor_id_is_null_or_empty", "criticality": "warn", }, { "check": { "function": "is_in_range", - "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}}, - "filter":None, + "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}, + }, + "filter": None, "name": "rate_code_id_isnt_in_range", "criticality": "warn", }, @@ -127,98 +127,80 @@ def test_generate_dq_no_rules(ws): expectations = generator.generate_dq_rules(None, level="warn") assert not expectations + def test_generate_dq_rules_dataframe_filter(ws): generator = DQGenerator(ws) - test_rules=[ - DQProfile(name="is_not_null", - column="machine_id", - description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"), - - DQProfile(name="is_in", - column="vendor_id", - parameters={"in": ["1", "4", "2"]}, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"), - - DQProfile(name="is_not_null", - column="cost", - description=None, - ), - - DQProfile(name="is_not_null", - column="next_scheduled_date", - description=None, - ), - - DQProfile(name="is_not_null", - column="safety_check_passed", - description=None, - ), - - DQProfile(name="is_not_null_or_empty", - column="vendor_id", - parameters={"trim_strings": True}), - - ] + test_rules = [ + DQProfile( + name="is_not_null", + column="machine_id", + description=None, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + ), + DQProfile( + name="is_in", + column="vendor_id", + parameters={"in": ["1", "4", "2"]}, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + ), + DQProfile( + name="is_not_null", + column="cost", + description=None, + ), + DQProfile( + name="is_not_null", + column="next_scheduled_date", + description=None, + ), + DQProfile( + name="is_not_null", + column="safety_check_passed", + description=None, + ), + DQProfile(name="is_not_null_or_empty", column="vendor_id", parameters={"trim_strings": True}), + ] expectations = generator.generate_dq_rules(test_rules) - + expected = [ { - "check": {"function": "is_not_null", - "arguments": {"column": "machine_id"}}, - "filter":"machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + "check": {"function": "is_not_null", "arguments": {"column": "machine_id"}}, + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", "name": "machine_id_is_null", - "criticality": "error" - + "criticality": "error", }, { - 'check': {"function": "is_in_list", - "arguments": {"allowed": ['1','4','2'],"column": "vendor_id"}}, - - "filter":"machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", - "criticality": "error", - "name": "vendor_id_other_value" - }, + 'check': {"function": "is_in_list", "arguments": {"allowed": ['1', '4', '2'], "column": "vendor_id"}}, + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + "criticality": "error", + "name": "vendor_id_other_value", + }, { - "check": { - "function": "is_not_null", - "arguments": {"column": "cost"}}, - + "check": {"function": "is_not_null", "arguments": {"column": "cost"}}, "filter": None, "name": "cost_is_null", "criticality": "error", - }, - { - "check": { - "function": "is_not_null", - "arguments": {"column": "next_scheduled_date"}}, - + "check": {"function": "is_not_null", "arguments": {"column": "next_scheduled_date"}}, "filter": None, "name": "next_scheduled_date_is_null", "criticality": "error", - }, - { - "check": { - "function": "is_not_null", - "arguments": {"column": "safety_check_passed"}}, + "check": {"function": "is_not_null", "arguments": {"column": "safety_check_passed"}}, "filter": None, - "name": "safety_check_passed_is_null", "criticality": "error", - }, { "check": { "function": "is_not_null_and_not_empty", - "arguments": {"column": "vendor_id", "trim_strings": True}}, + "arguments": {"column": "vendor_id", "trim_strings": True}, + }, "filter": None, "name": "vendor_id_is_null_or_empty", - "criticality": "error" - } - + "criticality": "error", + }, ] assert expectations == expected diff --git a/tests/integration/test_save_and_load_checks_from_table.py b/tests/integration/test_save_and_load_checks_from_table.py index ac514a604..383509e06 100644 --- a/tests/integration/test_save_and_load_checks_from_table.py +++ b/tests/integration/test_save_and_load_checks_from_table.py @@ -51,9 +51,7 @@ { "name": "column_not_less_than", "criticality": "warn", - "check": { - "function": "is_not_less_than", - "arguments": {"column": "col_2", "limit": 1}}, + "check": {"function": "is_not_less_than", "arguments": {"column": "col_2", "limit": 1}}, "filter": "Col_3 >1", "user_metadata": {"check_type": "standardization", "check_owner": "someone_else@email.com"}, }, @@ -65,7 +63,6 @@ ] - def test_load_checks_when_checks_table_does_not_exist(ws, make_schema, make_random, spark): catalog_name = "main" schema_name = make_schema(catalog_name=catalog_name).name diff --git a/tests/integration/test_save_checks_to_workspace_file.py b/tests/integration/test_save_checks_to_workspace_file.py index b27f47213..35b5a5f51 100644 --- a/tests/integration/test_save_checks_to_workspace_file.py +++ b/tests/integration/test_save_checks_to_workspace_file.py @@ -18,30 +18,24 @@ { "criticality": "error", "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}}, - }, + }, { "check": { "function": "is_not_null", "arguments": {"column": "next_scheduled_date"}, - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, - + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + }, "name": "next_scheduled_date_is_null", "criticality": "Error", - }, { - "check": { - "function": "is_not_null", - "arguments": {"column": "cost"}, - "filter": None}, - - "name": "cost_is_null", - "criticality": "error", - - } - + "check": {"function": "is_not_null", "arguments": {"column": "cost"}, "filter": None}, + "name": "cost_is_null", + "criticality": "error", + }, ] + def test_save_checks_in_workspace_file_as_yaml(ws, spark, installation_ctx): installation_ctx.installation.save(installation_ctx.config) install_dir = installation_ctx.installation.install_folder() From 889a5574739695c34aa92670e422a14ad7de3fff Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Tue, 9 Sep 2025 19:45:07 -0300 Subject: [PATCH 12/34] First commit: Installing DQX from feature branch --- src/databricks/labs/dqx/profiler/profiler.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 7b744d818..92e3a7566 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -54,7 +54,12 @@ def __init__(self, workspace_client: WorkspaceClient, spark: SparkSession | None "sample_fraction": 0.3, # fraction of data to sample (30%) "sample_seed": None, # seed for sampling "limit": 1000, # limit the number of samples +<<<<<<< HEAD "filter": None, # filter to apply to the dataset +======= + "filter":{}, # filter to apply before profiling + +>>>>>>> ce34e87 (First commit: Installing DQX from feature branch) } @staticmethod From b4f29c97c030acfbbe98d002bace47f3da8c7f20 Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Wed, 10 Sep 2025 23:45:47 -0300 Subject: [PATCH 13/34] Implement methods in the DQProfiler class to filter dataframes before submited to profile. --- src/databricks/labs/dqx/profiler/profiler.py | 70 ++++++++++++++++++++ 1 file changed, 70 insertions(+) diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 92e3a7566..15d0fd252 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -54,12 +54,16 @@ def __init__(self, workspace_client: WorkspaceClient, spark: SparkSession | None "sample_fraction": 0.3, # fraction of data to sample (30%) "sample_seed": None, # seed for sampling "limit": 1000, # limit the number of samples +<<<<<<< HEAD <<<<<<< HEAD "filter": None, # filter to apply to the dataset ======= "filter":{}, # filter to apply before profiling >>>>>>> ce34e87 (First commit: Installing DQX from feature branch) +======= + "dataset_filter": None, # filter to apply to the dataset +>>>>>>> 33f55cc (Implement methods in the DQProfiler class to filter dataframes before submited to profile.) } @staticmethod @@ -327,10 +331,17 @@ def _sample(df: DataFrame, opts: dict[str, Any]) -> DataFrame: sample_fraction = opts.get("sample_fraction", None) sample_seed = opts.get("sample_seed", None) limit = opts.get("limit", None) +<<<<<<< HEAD filter = opts.get("filter", None) if filter: df = df.filter(filter) +======= + filter = opts.get("dataset_filter", None) + + if filter: + df = DQProfiler._filter_dataframe(df, filter) +>>>>>>> 33f55cc (Implement methods in the DQProfiler class to filter dataframes before submited to profile.) if sample_fraction: df = df.sample(withReplacement=False, fraction=sample_fraction, seed=sample_seed) if limit: @@ -821,3 +832,62 @@ def _round_decimal(value: decimal.Decimal, direction: str) -> decimal.Decimal: if direction == "up": return value.to_integral_value(rounding=decimal.ROUND_CEILING) return value + + @staticmethod + def _filter_dataframe(df: DataFrame, filter: dict[str, Any]) -> DataFrame: + """ + Filters a DataFrame based on a dynamic filter. + + Args: + df: The input DataFrame to filter. + filter: A dictionary where keys are column names and values are the conditions to filter on. + + Returns: + A filtered DataFrame. + """ + + # Get the list of columns in the DataFrame + df_columns = df.columns + + for column, condition in filter.items(): + if column not in df_columns: + raise ValueError(f"Column '{column}' does not exist in the DataFrame.") + if isinstance(condition, (list, tuple)): + # If the condition is a list or tuple, use the `isin` filter + df = df.filter(F.col(column).isin(condition)) + elif isinstance(condition, dict): + # If the condition is a dictionary, handle operators like >, <, etc. + for operator, value in condition.items(): + df = DQProfiler._apply_filter_operator(df, column, operator, value) + else: + # Default equality filter + df = df.filter(F.col(column) == condition) + return df + + @staticmethod + def _apply_filter_operator(df: DataFrame, column: str, operator: str, value: Any) -> DataFrame: + """ + Applies a specific operator-based condition to a DataFrame column. + + Args: + df: The DataFrame to filter. + column: The column name to apply the operator on. + operator: The operator to use (e.g., "gt", "lt"). + value: The value to compare against. + + Returns: + A filtered DataFrame. + """ + operators = { + "gt": F.col(column) > value, + "lt": F.col(column) < value, + "gte": F.col(column) >= value, + "lte": F.col(column) <= value, + "eq": F.col(column) == value, + "neq": F.col(column) != value, + } + + if operator not in operators: + raise ValueError(f"Unsupported operator '{operator}' for column '{column}'.") + + return df.filter(operators[operator]) From 6b92cdcaa9384ebee4057e778192ee36b70d6c4b Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Thu, 11 Sep 2025 22:36:48 -0300 Subject: [PATCH 14/34] Refactor _sample method to utilize filtered dataframe --- src/databricks/labs/dqx/profiler/profiler.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 15d0fd252..9080453c7 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -339,13 +339,18 @@ def _sample(df: DataFrame, opts: dict[str, Any]) -> DataFrame: ======= filter = opts.get("dataset_filter", None) +<<<<<<< HEAD if filter: df = DQProfiler._filter_dataframe(df, filter) >>>>>>> 33f55cc (Implement methods in the DQProfiler class to filter dataframes before submited to profile.) +======= + + df_filter = DQProfiler._filter_dataframe(df, filter) +>>>>>>> 48cd64d (Refactor _sample method to utilize filtered dataframe) if sample_fraction: - df = df.sample(withReplacement=False, fraction=sample_fraction, seed=sample_seed) + df = df_filter.sample(withReplacement=False, fraction=sample_fraction, seed=sample_seed) if limit: - df = df.limit(limit) + df = df_filter.limit(limit) return df @@ -845,6 +850,8 @@ def _filter_dataframe(df: DataFrame, filter: dict[str, Any]) -> DataFrame: Returns: A filtered DataFrame. """ + if not filter: + return df # Get the list of columns in the DataFrame df_columns = df.columns @@ -860,8 +867,8 @@ def _filter_dataframe(df: DataFrame, filter: dict[str, Any]) -> DataFrame: for operator, value in condition.items(): df = DQProfiler._apply_filter_operator(df, column, operator, value) else: - # Default equality filter - df = df.filter(F.col(column) == condition) + # If the condition is not a list, tuple, or dictionary, raise an error + raise ValueError(f"Unsupported filter condition for column '{column}': {condition}") return df @staticmethod From 505d51d438e1cd59bdf314cdab4aca571ab41d0f Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Fri, 12 Sep 2025 23:40:13 -0300 Subject: [PATCH 15/34] Refactor _sample method to utilize return the filtered dataframe --- src/databricks/labs/dqx/profiler/profiler.py | 10 +++++++--- tests/unit/test_profiler.py | 1 + 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 9080453c7..6bae20fc9 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -339,6 +339,7 @@ def _sample(df: DataFrame, opts: dict[str, Any]) -> DataFrame: ======= filter = opts.get("dataset_filter", None) +<<<<<<< HEAD <<<<<<< HEAD if filter: df = DQProfiler._filter_dataframe(df, filter) @@ -347,10 +348,13 @@ def _sample(df: DataFrame, opts: dict[str, Any]) -> DataFrame: df_filter = DQProfiler._filter_dataframe(df, filter) >>>>>>> 48cd64d (Refactor _sample method to utilize filtered dataframe) +======= + df = DQProfiler._filter_dataframe(df, filter) +>>>>>>> c77a6eb (Refactor _sample method to utilize return the filtered dataframe) if sample_fraction: - df = df_filter.sample(withReplacement=False, fraction=sample_fraction, seed=sample_seed) + df = df.sample(withReplacement=False, fraction=sample_fraction, seed=sample_seed) if limit: - df = df_filter.limit(limit) + df = df.limit(limit) return df @@ -868,7 +872,7 @@ def _filter_dataframe(df: DataFrame, filter: dict[str, Any]) -> DataFrame: df = DQProfiler._apply_filter_operator(df, column, operator, value) else: # If the condition is not a list, tuple, or dictionary, raise an error - raise ValueError(f"Unsupported filter condition for column '{column}': {condition}") + raise ValueError(f"Unsupported filter condition for column '{column}': {condition}") return df @staticmethod diff --git a/tests/unit/test_profiler.py b/tests/unit/test_profiler.py index 62b0b14d6..bfd033502 100644 --- a/tests/unit/test_profiler.py +++ b/tests/unit/test_profiler.py @@ -2,6 +2,7 @@ from databricks.labs.dqx.profiler.profiler import DQProfiler + def test_get_columns_or_fields(): inp = T.StructType( [ From 0ed896ea4d9c9e4c627fe9db9c262a171081ac47 Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Sat, 13 Sep 2025 16:47:32 -0300 Subject: [PATCH 16/34] Implement print statement to test dataset_filter feature --- src/databricks/labs/dqx/profiler/profiler.py | 12 +- tests/integration/test_profiler.py | 117 +++++++++++++++++++ tests/unit/test_profiler.py | 1 + 3 files changed, 128 insertions(+), 2 deletions(-) diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 6bae20fc9..666a4ea38 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -55,6 +55,7 @@ def __init__(self, workspace_client: WorkspaceClient, spark: SparkSession | None "sample_seed": None, # seed for sampling "limit": 1000, # limit the number of samples <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD "filter": None, # filter to apply to the dataset ======= @@ -64,6 +65,9 @@ def __init__(self, workspace_client: WorkspaceClient, spark: SparkSession | None ======= "dataset_filter": None, # filter to apply to the dataset >>>>>>> 33f55cc (Implement methods in the DQProfiler class to filter dataframes before submited to profile.) +======= + # "dataset_filter": None, # filter to apply to the dataset +>>>>>>> a3a71dc (Implement print statement to test dataset_filter feature) } @staticmethod @@ -112,15 +116,19 @@ def profile( options = {} options = {**self.default_profile_options, **options} # merge default options with user-provided options +<<<<<<< HEAD +======= + print(options) +>>>>>>> a3a71dc (Implement print statement to test dataset_filter feature) df = self._sample(df, options) - + print(df) dq_rules: list[DQProfile] = [] total_count = df.count() summary_stats = self._get_df_summary_as_dict(df) if total_count == 0: return summary_stats, dq_rules - + print(summary_stats) self._profile(df, df_columns, dq_rules, options, summary_stats, total_count) filter = options.get("filter", None) if filter: diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index 2c5f1d35f..b7b443b74 100644 --- a/tests/integration/test_profiler.py +++ b/tests/integration/test_profiler.py @@ -1264,6 +1264,7 @@ def test_profile_tables_no_tables_or_patterns(ws): with pytest.raises(ValueError, match="Either 'tables' or 'patterns' must be provided"): profiler.profile_tables() +<<<<<<< HEAD def test_profile_with_dataset_filter(spark, ws): schema = T.StructType( @@ -1423,3 +1424,119 @@ def test_profile_with_dataset_filter(spark, ws): assert filtered_df.count() == 3 assert len(stats.keys()) > 0 assert rules == expected_rules +======= +def test_profile_sample_dataset_filter(spark, ws): + schema = T.StructType( + [ + T.StructField("maintenance_id", T.StringType(), False), + T.StructField("machine_id", T.StringType(), False), + T.StructField("maintenance_type", T.StringType(), True), + T.StructField("maintenance_date", T.DateType(), True), + T.StructField("duration_minutes", T.IntegerType(), True), + T.StructField("cost", T.DecimalType(10, 2), True), + T.StructField("next_scheduled_date", T.DateType(), True), + T.StructField("work_order_id", T.StringType(), True), + T.StructField("safety_check_passed", T.BooleanType(), True), + T.StructField("parts_list", T.ArrayType(T.StringType()), True), + T.StructField("ingest_date", T.DateType(), True), + ] + ) + maintenance_data = [ + # Ingest date 2025-04-28 + ( + "MTN-001", + "MCH-001", + "preventive", + datetime.strptime("2025-04-01", "%Y-%m-%d").date(), + 120, + Decimal("450.00"), + datetime.strptime("2025-07-01", "%Y-%m-%d").date(), + "WO-001", + True, + ["filter", "gasket"], + datetime.strptime("2025-04-28", "%Y-%m-%d").date(), + ), + ( + "MTN-002", + "MCH-002", + "corrective", + datetime.strptime("2025-04-15", "%Y-%m-%d").date(), + 240, + Decimal("1200.50"), + datetime.strptime("2026-04-01", "%Y-%m-%d").date(), + "WO-002", + False, + ["motor"], + datetime.strptime("2025-04-28", "%Y-%m-%d").date(), + ), # Future date issue + # Ingest date 2025-04-29 + ( + "MTN-003", + "MCH-003", + None, + datetime.strptime("2025-04-20", "%Y-%m-%d").date(), + -30, + Decimal("-500.00"), + datetime.strptime("2024-04-20", "%Y-%m-%d").date(), + "INVALID", + None, + [], + datetime.strptime("2025-04-29", "%Y-%m-%d").date(), + ), # Multiple issues + ( + "MTN-004", + "MCH-001", + "predictive", + datetime.strptime("2025-04-25", "%Y-%m-%d").date(), + 180, + Decimal("800.00"), + datetime.strptime("2025-10-01", "%Y-%m-%d").date(), + "WO-003", + True, + ["sensor"], + datetime.strptime("2025-04-29", "%Y-%m-%d").date(), + ), + # Ingest date 2025-04-30 + ( + "MTN-005", + "MCH-002", + "preventive", + datetime.strptime("2025-04-29", "%Y-%m-%d").date(), + 90, + Decimal("300.00"), + datetime.strptime("2025-07-15", "%Y-%m-%d").date(), + "WO-004", + True, + ["valve"], + datetime.strptime("2025-04-30", "%Y-%m-%d").date(), + ), + ( + "MTN-006", + "MCH-003", + "corrective", + datetime.strptime("2025-04-30", "%Y-%m-%d").date(), + 60, + Decimal("150.00"), + datetime.strptime("2025-08-01", "%Y-%m-%d").date(), + "WO-005", + False, + ["pump"], + datetime.strptime("2025-04-30", "%Y-%m-%d").date(), + ), + ] + + + input_df = spark.createDataFrame(maintenance_data, schema=schema) + + custom_options = { + "sample_fraction": None, + "limit": None, + "dataset_filter": {"machine_id":['MCH-002', 'MCH-003'], "maintenance_type":{'eq': 'preventive'}} + } + + + profiler = DQProfiler(ws) + filtered_df = profiler._sample(input_df, opts=custom_options) + assert filtered_df.count() == 1 + +>>>>>>> a3a71dc (Implement print statement to test dataset_filter feature) diff --git a/tests/unit/test_profiler.py b/tests/unit/test_profiler.py index bfd033502..76a355b2f 100644 --- a/tests/unit/test_profiler.py +++ b/tests/unit/test_profiler.py @@ -29,3 +29,4 @@ def test_get_columns_or_fields(): T.StructField("ss1.s2.ns3", T.DateType()), ] assert fields == expected + From 6a1745800e8aa924587602fcd20c6ce1e00378d3 Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Sun, 14 Sep 2025 22:27:32 -0300 Subject: [PATCH 17/34] Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter --- src/databricks/labs/dqx/profiler/profiler.py | 100 +++++------ tests/integration/test_profiler.py | 165 ++++++++++++------- 2 files changed, 143 insertions(+), 122 deletions(-) diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 666a4ea38..79edafd37 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -56,6 +56,7 @@ def __init__(self, workspace_client: WorkspaceClient, spark: SparkSession | None "limit": 1000, # limit the number of samples <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD "filter": None, # filter to apply to the dataset ======= @@ -68,6 +69,9 @@ def __init__(self, workspace_client: WorkspaceClient, spark: SparkSession | None ======= # "dataset_filter": None, # filter to apply to the dataset >>>>>>> a3a71dc (Implement print statement to test dataset_filter feature) +======= + "dataset_filter_expression": None, # filter to apply to the dataset +>>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) } @staticmethod @@ -117,18 +121,22 @@ def profile( options = {**self.default_profile_options, **options} # merge default options with user-provided options <<<<<<< HEAD +<<<<<<< HEAD ======= print(options) >>>>>>> a3a71dc (Implement print statement to test dataset_filter feature) +======= + +>>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) df = self._sample(df, options) - print(df) + dq_rules: list[DQProfile] = [] total_count = df.count() summary_stats = self._get_df_summary_as_dict(df) if total_count == 0: return summary_stats, dq_rules - print(summary_stats) + self._profile(df, df_columns, dq_rules, options, summary_stats, total_count) filter = options.get("filter", None) if filter: @@ -339,6 +347,7 @@ def _sample(df: DataFrame, opts: dict[str, Any]) -> DataFrame: sample_fraction = opts.get("sample_fraction", None) sample_seed = opts.get("sample_seed", None) limit = opts.get("limit", None) +<<<<<<< HEAD <<<<<<< HEAD filter = opts.get("filter", None) @@ -359,6 +368,12 @@ def _sample(df: DataFrame, opts: dict[str, Any]) -> DataFrame: ======= df = DQProfiler._filter_dataframe(df, filter) >>>>>>> c77a6eb (Refactor _sample method to utilize return the filtered dataframe) +======= + filter = opts.get("dataset_filter_expression", None) + + if filter: + df = df.filter(filter) +>>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) if sample_fraction: df = df.sample(withReplacement=False, fraction=sample_fraction, seed=sample_seed) if limit: @@ -409,6 +424,7 @@ def _calculate_metrics( """ max_nulls = opts.get("max_null_ratio", 0) trim_strings = opts.get("trim_strings", True) + filter = opts.get("dataset_filter_expression", None) dst = df.select(field_name).dropna() if typ == T.StringType() and trim_strings: @@ -428,20 +444,27 @@ def _calculate_metrics( column=field_name, description=f"Column {field_name} has {null_percentage * 100:.1f}% of null values " f"(allowed {max_nulls * 100:.1f}%)", + parameters={"dataset_filter_expression": filter}, ) ) else: - dq_rules.append(DQProfile(name="is_not_null", column=field_name)) + dq_rules.append(DQProfile(name="is_not_null", column=field_name, parameters={"dataset_filter_expression": filter})) if self._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( +<<<<<<< HEAD DQProfile( name="is_in", column=field_name, parameters={"in": [row[0] for row in dst2.collect()]}, ) +======= + DQProfile(name="is_in", + column=field_name, + parameters={"in": [row[0] for row in dst2.collect()], "dataset_filter_expression": filter}) +>>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) ) if ( typ == T.StringType() @@ -453,11 +476,17 @@ def _calculate_metrics( cnt = dst2.count() if cnt <= (metrics["count"] * opts.get("max_empty_ratio", 0)): dq_rules.append( +<<<<<<< HEAD DQProfile( name="is_not_null_or_empty", column=field_name, parameters={"trim_strings": trim_strings}, ) +======= + DQProfile(name="is_not_null_or_empty", + column=field_name, + parameters={"trim_strings": trim_strings, "dataset_filter_expression": filter}) +>>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) ) if metrics["count_non_null"] > 0 and self._type_supports_min_max(typ): rule = self._extract_min_max(dst, field_name, typ, metrics, opts) @@ -575,7 +604,7 @@ def _extract_min_max( if opts is None: opts = {} - + filter= opts.get("dataset_filter_expression", None) outlier_cols = opts.get("outlier_columns", []) column = dst.columns[0] if opts.get("remove_outliers", True) and ( @@ -611,10 +640,14 @@ def _extract_min_max( logger.info(f"Can't get min/max for field {col_name}") if descr and min_limit and max_limit: return DQProfile( +<<<<<<< HEAD name="min_max", column=col_name, parameters={"min": min_limit, "max": max_limit}, description=descr, +======= + name="min_max", column=col_name, parameters={"min": min_limit, "max": max_limit,"dataset_filter_expression": filter}, description=descr +>>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) ) return None @@ -850,63 +883,4 @@ def _round_decimal(value: decimal.Decimal, direction: str) -> decimal.Decimal: return value.to_integral_value(rounding=decimal.ROUND_CEILING) return value - @staticmethod - def _filter_dataframe(df: DataFrame, filter: dict[str, Any]) -> DataFrame: - """ - Filters a DataFrame based on a dynamic filter. - - Args: - df: The input DataFrame to filter. - filter: A dictionary where keys are column names and values are the conditions to filter on. - - Returns: - A filtered DataFrame. - """ - if not filter: - return df - - # Get the list of columns in the DataFrame - df_columns = df.columns - - for column, condition in filter.items(): - if column not in df_columns: - raise ValueError(f"Column '{column}' does not exist in the DataFrame.") - if isinstance(condition, (list, tuple)): - # If the condition is a list or tuple, use the `isin` filter - df = df.filter(F.col(column).isin(condition)) - elif isinstance(condition, dict): - # If the condition is a dictionary, handle operators like >, <, etc. - for operator, value in condition.items(): - df = DQProfiler._apply_filter_operator(df, column, operator, value) - else: - # If the condition is not a list, tuple, or dictionary, raise an error - raise ValueError(f"Unsupported filter condition for column '{column}': {condition}") - return df - - @staticmethod - def _apply_filter_operator(df: DataFrame, column: str, operator: str, value: Any) -> DataFrame: - """ - Applies a specific operator-based condition to a DataFrame column. - - Args: - df: The DataFrame to filter. - column: The column name to apply the operator on. - operator: The operator to use (e.g., "gt", "lt"). - value: The value to compare against. - - Returns: - A filtered DataFrame. - """ - operators = { - "gt": F.col(column) > value, - "lt": F.col(column) < value, - "gte": F.col(column) >= value, - "lte": F.col(column) <= value, - "eq": F.col(column) == value, - "neq": F.col(column) != value, - } - - if operator not in operators: - raise ValueError(f"Unsupported operator '{operator}' for column '{column}'.") - return df.filter(operators[operator]) diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index b7b443b74..c3a742008 100644 --- a/tests/integration/test_profiler.py +++ b/tests/integration/test_profiler.py @@ -1265,6 +1265,7 @@ def test_profile_tables_no_tables_or_patterns(ws): profiler.profile_tables() <<<<<<< HEAD +<<<<<<< HEAD def test_profile_with_dataset_filter(spark, ws): schema = T.StructType( @@ -1429,100 +1430,87 @@ def test_profile_sample_dataset_filter(spark, ws): schema = T.StructType( [ T.StructField("maintenance_id", T.StringType(), False), +======= +def test_profile_with_dataset_filter(spark, ws): + schema = T.StructType( [ + +>>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) T.StructField("machine_id", T.StringType(), False), T.StructField("maintenance_type", T.StringType(), True), - T.StructField("maintenance_date", T.DateType(), True), - T.StructField("duration_minutes", T.IntegerType(), True), + T.StructField("maintenance_date", T.DateType(), True), T.StructField("cost", T.DecimalType(10, 2), True), T.StructField("next_scheduled_date", T.DateType(), True), - T.StructField("work_order_id", T.StringType(), True), T.StructField("safety_check_passed", T.BooleanType(), True), - T.StructField("parts_list", T.ArrayType(T.StringType()), True), - T.StructField("ingest_date", T.DateType(), True), + ] ) maintenance_data = [ - # Ingest date 2025-04-28 + ( - "MTN-001", "MCH-001", "preventive", - datetime.strptime("2025-04-01", "%Y-%m-%d").date(), - 120, + date(2025, 4,1), Decimal("450.00"), - datetime.strptime("2025-07-01", "%Y-%m-%d").date(), - "WO-001", - True, - ["filter", "gasket"], - datetime.strptime("2025-04-28", "%Y-%m-%d").date(), + date(2025,7,1), + True, ), ( - "MTN-002", "MCH-002", "corrective", - datetime.strptime("2025-04-15", "%Y-%m-%d").date(), - 240, + date(2025, 4,15), Decimal("1200.50"), - datetime.strptime("2026-04-01", "%Y-%m-%d").date(), - "WO-002", - False, - ["motor"], - datetime.strptime("2025-04-28", "%Y-%m-%d").date(), - ), # Future date issue - # Ingest date 2025-04-29 + date(2026,4,1), + False, + ), ( - "MTN-003", "MCH-003", None, - datetime.strptime("2025-04-20", "%Y-%m-%d").date(), - -30, + date(2025,4,20), Decimal("-500.00"), - datetime.strptime("2024-04-20", "%Y-%m-%d").date(), - "INVALID", - None, - [], - datetime.strptime("2025-04-29", "%Y-%m-%d").date(), - ), # Multiple issues + date(2024,4,20), + None, + ), ( - "MTN-004", "MCH-001", "predictive", - datetime.strptime("2025-04-25", "%Y-%m-%d").date(), - 180, + date(2025,4,25), Decimal("800.00"), - datetime.strptime("2025-10-01", "%Y-%m-%d").date(), - "WO-003", - True, - ["sensor"], - datetime.strptime("2025-04-29", "%Y-%m-%d").date(), + date(2025,10,1), + True, ), - # Ingest date 2025-04-30 + ( - "MTN-005", "MCH-002", "preventive", - datetime.strptime("2025-04-29", "%Y-%m-%d").date(), - 90, + date(2025,4,29), Decimal("300.00"), - datetime.strptime("2025-07-15", "%Y-%m-%d").date(), - "WO-004", - True, - ["valve"], - datetime.strptime("2025-04-30", "%Y-%m-%d").date(), + date(2025,7,15), + True, ), ( - "MTN-006", "MCH-003", "corrective", - datetime.strptime("2025-04-30", "%Y-%m-%d").date(), - 60, + date(2025,4,30), Decimal("150.00"), - datetime.strptime("2025-08-01", "%Y-%m-%d").date(), - "WO-005", - False, - ["pump"], - datetime.strptime("2025-04-30", "%Y-%m-%d").date(), + date(2025,8,1), + False, + ), + ( + "MCH-002", + "preventive", + date(2025,5,30), + Decimal("150.00"), + date(2025,9,1), + True, ), + ( + "MCH-002", + "preventive", + date(2025,7,30), + Decimal("100.00"), + date(2025,12,1), + True, + ) ] @@ -1531,12 +1519,71 @@ def test_profile_sample_dataset_filter(spark, ws): custom_options = { "sample_fraction": None, "limit": None, - "dataset_filter": {"machine_id":['MCH-002', 'MCH-003'], "maintenance_type":{'eq': 'preventive'}} + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" } profiler = DQProfiler(ws) filtered_df = profiler._sample(input_df, opts=custom_options) +<<<<<<< HEAD assert filtered_df.count() == 1 >>>>>>> a3a71dc (Implement print statement to test dataset_filter feature) +======= + + stats, rules = profiler.profile(input_df,options=custom_options) + + expected_rules = [ + DQProfile(name="is_not_null", + column="machine_id", + description=None, + parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + DQProfile(name="is_not_null", + column="maintenance_type", + description=None, + parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + DQProfile(name="is_not_null", + column="maintenance_date", + description=None, + parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + DQProfile( + name="min_max", + column="maintenance_date", + description="Real min/max values were used", + parameters={"min": date(2025, 4, 29), + "max": date(2025, 7, 30), + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + ), + DQProfile(name="is_not_null", + column="cost", + description=None, + parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + DQProfile( + name="min_max", + column="cost", + parameters={"min": Decimal('100.00') , "max": Decimal('300.00'),"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + description="Real min/max values were used" + ), + DQProfile(name="is_not_null", + column="next_scheduled_date", + description=None, + parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + DQProfile( + name="min_max", + column="next_scheduled_date", + description="Real min/max values were used", + parameters={"min": date(2025, 7, 15), "max": date(2025, 12, 1),"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + ), + DQProfile(name="is_not_null", + column="safety_check_passed", + description=None, + parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + + ] + + + assert filtered_df.count() == 3 + assert len(stats.keys()) > 0 + assert rules == expected_rules + +>>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) From 0732a0bf002dc972f4043c93cdcc10f8b785cd99 Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Fri, 19 Sep 2025 09:53:00 -0300 Subject: [PATCH 18/34] Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file. --- src/databricks/labs/dqx/profiler/generator.py | 15 ++ src/databricks/labs/dqx/profiler/profiler.py | 39 +++- tests/integration/test_profiler.py | 191 +++++++++++------- tests/integration/test_rules_generator.py | 110 ++++++++++ .../test_save_checks_to_workspace_file.py | 15 +- tests/unit/test_profiler.py | 2 - 6 files changed, 286 insertions(+), 86 deletions(-) diff --git a/src/databricks/labs/dqx/profiler/generator.py b/src/databricks/labs/dqx/profiler/generator.py index 7a3566b5f..c866f8ada 100644 --- a/src/databricks/labs/dqx/profiler/generator.py +++ b/src/databricks/labs/dqx/profiler/generator.py @@ -61,6 +61,7 @@ def dq_generate_is_in(column: str, filter: str | None, level: str = "error", **p "filter": filter, "name": f"{column}_other_value", "criticality": level, + "filter": params.get("dataset_filter_expression", None), } @staticmethod @@ -78,6 +79,7 @@ def dq_generate_min_max(column: str, filter: str | None, level: str = "error", * """ min_limit = params.get("min") max_limit = params.get("max") + filter = params.get("dataset_filter_expression", None) if not isinstance(min_limit, int) or not isinstance(max_limit, int): return None # TODO handle timestamp and dates: https://github.com/databrickslabs/dqx/issues/71 @@ -95,6 +97,7 @@ def dq_generate_min_max(column: str, filter: str | None, level: str = "error", * "filter": filter, "name": f"{column}_isnt_in_range", "criticality": level, + "filter": filter, } if max_limit is not None: @@ -106,6 +109,7 @@ def dq_generate_min_max(column: str, filter: str | None, level: str = "error", * "filter": filter, "name": f"{column}_not_greater_than", "criticality": level, + "filter": filter, } if min_limit is not None: @@ -117,6 +121,7 @@ def dq_generate_min_max(column: str, filter: str | None, level: str = "error", * "filter": filter, "name": f"{column}_not_less_than", "criticality": level, + "filter": filter, } return None @@ -135,12 +140,17 @@ def dq_generate_is_not_null(column: str, filter: str | None, level: str = "error A dictionary representing the data quality rule. """ params = params or {} +<<<<<<< HEAD +======= + filter = params.get("dataset_filter_expression", None) +>>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) return { "check": {"function": "is_not_null", "arguments": {"column": column}}, "filter": filter, "name": f"{column}_is_null", "criticality": level, + "filter": filter, } @staticmethod @@ -156,7 +166,11 @@ def dq_generate_is_not_null_or_empty(column: str, filter: str | None, level: str Returns: A dictionary representing the data quality rule. """ +<<<<<<< HEAD +======= + filter = params.get("dataset_filter_expression", None) +>>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) return { "check": { "function": "is_not_null_and_not_empty", @@ -165,6 +179,7 @@ def dq_generate_is_not_null_or_empty(column: str, filter: str | None, level: str "filter": filter, "name": f"{column}_is_null_or_empty", "criticality": level, + "filter": filter, } _checks_mapping = { diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 79edafd37..51517f60c 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -122,6 +122,7 @@ def profile( options = {**self.default_profile_options, **options} # merge default options with user-provided options <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD ======= print(options) @@ -129,14 +130,17 @@ def profile( ======= >>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) +======= + +>>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) df = self._sample(df, options) - + dq_rules: list[DQProfile] = [] total_count = df.count() summary_stats = self._get_df_summary_as_dict(df) if total_count == 0: return summary_stats, dq_rules - + self._profile(df, df_columns, dq_rules, options, summary_stats, total_count) filter = options.get("filter", None) if filter: @@ -448,12 +452,15 @@ def _calculate_metrics( ) ) else: - dq_rules.append(DQProfile(name="is_not_null", column=field_name, parameters={"dataset_filter_expression": filter})) + dq_rules.append( + DQProfile(name="is_not_null", column=field_name, parameters={"dataset_filter_expression": filter}) + ) if self._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( +<<<<<<< HEAD <<<<<<< HEAD DQProfile( name="is_in", @@ -465,6 +472,13 @@ def _calculate_metrics( column=field_name, parameters={"in": [row[0] for row in dst2.collect()], "dataset_filter_expression": filter}) >>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) +======= + DQProfile( + name="is_in", + column=field_name, + parameters={"in": [row[0] for row in dst2.collect()], "dataset_filter_expression": filter}, + ) +>>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) ) if ( typ == T.StringType() @@ -476,6 +490,7 @@ def _calculate_metrics( cnt = dst2.count() if cnt <= (metrics["count"] * opts.get("max_empty_ratio", 0)): dq_rules.append( +<<<<<<< HEAD <<<<<<< HEAD DQProfile( name="is_not_null_or_empty", @@ -487,6 +502,13 @@ def _calculate_metrics( column=field_name, parameters={"trim_strings": trim_strings, "dataset_filter_expression": filter}) >>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) +======= + DQProfile( + name="is_not_null_or_empty", + column=field_name, + parameters={"trim_strings": trim_strings, "dataset_filter_expression": filter}, + ) +>>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) ) if metrics["count_non_null"] > 0 and self._type_supports_min_max(typ): rule = self._extract_min_max(dst, field_name, typ, metrics, opts) @@ -604,7 +626,7 @@ def _extract_min_max( if opts is None: opts = {} - filter= opts.get("dataset_filter_expression", None) + filter = opts.get("dataset_filter_expression", None) outlier_cols = opts.get("outlier_columns", []) column = dst.columns[0] if opts.get("remove_outliers", True) and ( @@ -640,6 +662,7 @@ def _extract_min_max( logger.info(f"Can't get min/max for field {col_name}") if descr and min_limit and max_limit: return DQProfile( +<<<<<<< HEAD <<<<<<< HEAD name="min_max", column=col_name, @@ -648,6 +671,12 @@ def _extract_min_max( ======= name="min_max", column=col_name, parameters={"min": min_limit, "max": max_limit,"dataset_filter_expression": filter}, description=descr >>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) +======= + name="min_max", + column=col_name, + parameters={"min": min_limit, "max": max_limit, "dataset_filter_expression": filter}, + description=descr, +>>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) ) return None @@ -882,5 +911,3 @@ def _round_decimal(value: decimal.Decimal, direction: str) -> decimal.Decimal: if direction == "up": return value.to_integral_value(rounding=decimal.ROUND_CEILING) return value - - diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index c3a742008..442a67733 100644 --- a/tests/integration/test_profiler.py +++ b/tests/integration/test_profiler.py @@ -1266,6 +1266,7 @@ def test_profile_tables_no_tables_or_patterns(ws): <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD def test_profile_with_dataset_filter(spark, ws): schema = T.StructType( @@ -1435,93 +1436,94 @@ def test_profile_with_dataset_filter(spark, ws): schema = T.StructType( [ >>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) +======= + +def test_profile_with_dataset_filter(spark, ws): + schema = T.StructType( + [ +>>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) T.StructField("machine_id", T.StringType(), False), T.StructField("maintenance_type", T.StringType(), True), - T.StructField("maintenance_date", T.DateType(), True), + T.StructField("maintenance_date", T.DateType(), True), T.StructField("cost", T.DecimalType(10, 2), True), T.StructField("next_scheduled_date", T.DateType(), True), T.StructField("safety_check_passed", T.BooleanType(), True), - ] ) maintenance_data = [ - ( "MCH-001", "preventive", - date(2025, 4,1), + date(2025, 4, 1), Decimal("450.00"), - date(2025,7,1), - True, + date(2025, 7, 1), + True, ), ( "MCH-002", "corrective", - date(2025, 4,15), + date(2025, 4, 15), Decimal("1200.50"), - date(2026,4,1), - False, - ), + date(2026, 4, 1), + False, + ), ( "MCH-003", None, - date(2025,4,20), + date(2025, 4, 20), Decimal("-500.00"), - date(2024,4,20), - None, - ), + date(2024, 4, 20), + None, + ), ( "MCH-001", "predictive", - date(2025,4,25), + date(2025, 4, 25), Decimal("800.00"), - date(2025,10,1), - True, + date(2025, 10, 1), + True, ), - ( "MCH-002", "preventive", - date(2025,4,29), + date(2025, 4, 29), Decimal("300.00"), - date(2025,7,15), - True, + date(2025, 7, 15), + True, ), ( "MCH-003", "corrective", - date(2025,4,30), + date(2025, 4, 30), Decimal("150.00"), - date(2025,8,1), - False, + date(2025, 8, 1), + False, ), ( "MCH-002", "preventive", - date(2025,5,30), + date(2025, 5, 30), Decimal("150.00"), - date(2025,9,1), - True, + date(2025, 9, 1), + True, ), ( "MCH-002", "preventive", - date(2025,7,30), + date(2025, 7, 30), Decimal("100.00"), - date(2025,12,1), - True, - ) - ] - - + date(2025, 12, 1), + True, + ), + ] + input_df = spark.createDataFrame(maintenance_data, schema=schema) - + custom_options = { - "sample_fraction": None, - "limit": None, - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - } - + "sample_fraction": None, + "limit": None, + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + } profiler = DQProfiler(ws) filtered_df = profiler._sample(input_df, opts=custom_options) @@ -1531,59 +1533,94 @@ def test_profile_with_dataset_filter(spark, ws): >>>>>>> a3a71dc (Implement print statement to test dataset_filter feature) ======= - stats, rules = profiler.profile(input_df,options=custom_options) + stats, rules = profiler.profile(input_df, options=custom_options) expected_rules = [ - DQProfile(name="is_not_null", - column="machine_id", - description=None, - parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), - DQProfile(name="is_not_null", - column="maintenance_type", - description=None, - parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), - DQProfile(name="is_not_null", - column="maintenance_date", - description=None, - parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + DQProfile( + name="is_not_null", + column="machine_id", + description=None, + parameters={ + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + }, + ), + DQProfile( + name="is_not_null", + column="maintenance_type", + description=None, + parameters={ + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + }, + ), + DQProfile( + name="is_not_null", + column="maintenance_date", + description=None, + parameters={ + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + }, + ), DQProfile( name="min_max", column="maintenance_date", description="Real min/max values were used", - parameters={"min": date(2025, 4, 29), - "max": date(2025, 7, 30), - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + parameters={ + "min": date(2025, 4, 29), + "max": date(2025, 7, 30), + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + }, ), - DQProfile(name="is_not_null", - column="cost", - description=None, - parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), DQProfile( - name="min_max", - column="cost", - parameters={"min": Decimal('100.00') , "max": Decimal('300.00'),"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, - description="Real min/max values were used" - ), - DQProfile(name="is_not_null", - column="next_scheduled_date", - description=None, - parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + name="is_not_null", + column="cost", + description=None, + parameters={ + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + }, + ), + DQProfile( + name="min_max", + column="cost", + parameters={ + "min": Decimal('100.00'), + "max": Decimal('300.00'), + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + }, + description="Real min/max values were used", + ), + DQProfile( + name="is_not_null", + column="next_scheduled_date", + description=None, + parameters={ + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + }, + ), DQProfile( name="min_max", column="next_scheduled_date", description="Real min/max values were used", - parameters={"min": date(2025, 7, 15), "max": date(2025, 12, 1),"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + parameters={ + "min": date(2025, 7, 15), + "max": date(2025, 12, 1), + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + }, ), - DQProfile(name="is_not_null", - column="safety_check_passed", - description=None, - parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + DQProfile( + name="is_not_null", + column="safety_check_passed", + description=None, + parameters={ + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + }, + ), + ] - ] - - assert filtered_df.count() == 3 assert len(stats.keys()) > 0 assert rules == expected_rules +<<<<<<< HEAD >>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) +======= +>>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) diff --git a/tests/integration/test_rules_generator.py b/tests/integration/test_rules_generator.py index 9636098d2..423740273 100644 --- a/tests/integration/test_rules_generator.py +++ b/tests/integration/test_rules_generator.py @@ -1,5 +1,6 @@ import datetime from decimal import Decimal +# from datetime import date, datetime, timezone # from datetime import date, datetime, timezone @@ -127,6 +128,7 @@ def test_generate_dq_no_rules(ws): expectations = generator.generate_dq_rules(None, level="warn") assert not expectations +<<<<<<< HEAD def test_generate_dq_rules_dataframe_filter(ws): generator = DQGenerator(ws) @@ -202,5 +204,113 @@ def test_generate_dq_rules_dataframe_filter(ws): "name": "vendor_id_is_null_or_empty", "criticality": "error", }, +======= +def test_generate_dq_rules_dataframe_filter(ws): + generator = DQGenerator(ws) + test_rules=[ + DQProfile(name="is_not_null", + column="machine_id", + description=None, + parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + + DQProfile( + name="min_max", + column="maintenance_date", + description="Real min/max values were used", + parameters={"min": datetime.date(2025, 4, 29), + "max": datetime.date(2025, 7, 30), + "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + ), + DQProfile(name="is_not_null", + column="cost", + description=None, + parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + DQProfile( + name="min_max", + column="cost", + description="Real min/max values were used", + parameters={"min": Decimal('100.00') , "max": Decimal('300.00'),"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + ), + DQProfile(name="is_not_null", + column="next_scheduled_date", + description=None, + parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + DQProfile( + name="min_max", + column="next_scheduled_date", + description="Real min/max values were used", + parameters={"min": datetime.date(2025, 7, 15), "max": datetime.date(2025, 12, 1),"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + ), + DQProfile(name="is_not_null", + column="safety_check_passed", + description=None, + parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + + ] + expectations = generator.generate_dq_rules(test_rules) + + expected = [ + { + "check": {"function": "is_not_null", "arguments": {"column": "machine_id"}}, + "name": "machine_id_is_null", + "criticality": "Error", + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + }, + { + "check": { + "function": "min_max", + "arguments": {"column": "maintenance_date","min": datetime.date(2025, 4, 29), "max": datetime.date(2025, 7, 30)} + }, + "name": "maintenance_date_isnt_in_range", + "criticality": "Error", + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + }, + { + "check": { + "function": "is_not_null", + "arguments": {"column": "cost"}}, + + "name": "cost_is_null", + "criticality": "Error", + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + }, + { + "check": { + "function": "min_max", + "arguments": {"column": "cost","min": Decimal('100.00') , "max": Decimal('300.00')} + }, + "name": "cost_isnt_in_range", + "criticality": "Error", + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + }, + { + "check": { + "function": "is_not_null", + "arguments": {"column": "next_scheduled_date"}}, + + "name": "next_scheduled_date_is_null", + "criticality": "Error", + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + }, + { + "check": { + "function": "min_max", + "arguments": {"column": "next_scheduled_date","min": datetime.date(2025, 7, 15), "max": datetime.date(2025, 12, 1)} + }, + "name": "cost_isnt_in_range", + "criticality": "Error", + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + }, + { + "check": { + "function": "is_not_null", + "arguments": {"column": "safety_check_passed"}}, + + "name": "safety_check_passed_is_null", + "criticality": "Error", + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + } + +>>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) ] assert expectations == expected diff --git a/tests/integration/test_save_checks_to_workspace_file.py b/tests/integration/test_save_checks_to_workspace_file.py index 35b5a5f51..337f59ac0 100644 --- a/tests/integration/test_save_checks_to_workspace_file.py +++ b/tests/integration/test_save_checks_to_workspace_file.py @@ -18,6 +18,7 @@ { "criticality": "error", "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}}, +<<<<<<< HEAD }, { "check": { @@ -33,9 +34,21 @@ "name": "cost_is_null", "criticality": "error", }, +======= + }, + { + "check": { + "function": "is_not_null", + "arguments": {"column": "next_scheduled_date"}}, + + "name": "next_scheduled_date_is_null", + "criticality": "Error", + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + } + +>>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) ] - def test_save_checks_in_workspace_file_as_yaml(ws, spark, installation_ctx): installation_ctx.installation.save(installation_ctx.config) install_dir = installation_ctx.installation.install_folder() diff --git a/tests/unit/test_profiler.py b/tests/unit/test_profiler.py index 76a355b2f..62b0b14d6 100644 --- a/tests/unit/test_profiler.py +++ b/tests/unit/test_profiler.py @@ -2,7 +2,6 @@ from databricks.labs.dqx.profiler.profiler import DQProfiler - def test_get_columns_or_fields(): inp = T.StructType( [ @@ -29,4 +28,3 @@ def test_get_columns_or_fields(): T.StructField("ss1.s2.ns3", T.DateType()), ] assert fields == expected - From 5dc8243172b9ac7eaa77e767da4ff3ca68b643d6 Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Sat, 20 Sep 2025 20:14:41 -0300 Subject: [PATCH 19/34] Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter. --- src/databricks/labs/dqx/profiler/generator.py | 53 ++++++++++++----- src/databricks/labs/dqx/profiler/profiler.py | 37 +++++++++--- tests/integration/test_profiler.py | 43 +++++--------- tests/integration/test_rules_generator.py | 58 +++++++++---------- 4 files changed, 111 insertions(+), 80 deletions(-) diff --git a/src/databricks/labs/dqx/profiler/generator.py b/src/databricks/labs/dqx/profiler/generator.py index c866f8ada..82c51dba0 100644 --- a/src/databricks/labs/dqx/profiler/generator.py +++ b/src/databricks/labs/dqx/profiler/generator.py @@ -57,11 +57,14 @@ def dq_generate_is_in(column: str, filter: str | None, level: str = "error", **p A dictionary representing the data quality rule. """ return { +<<<<<<< HEAD "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"]}}, "filter": filter, +======= + "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"],"filter": params.get("filter", None)}}, +>>>>>>> 07d81a3 (Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.) "name": f"{column}_other_value", - "criticality": level, - "filter": params.get("dataset_filter_expression", None), + "criticality": level } @staticmethod @@ -79,7 +82,7 @@ def dq_generate_min_max(column: str, filter: str | None, level: str = "error", * """ min_limit = params.get("min") max_limit = params.get("max") - filter = params.get("dataset_filter_expression", None) + filter = params.get("filter", None) if not isinstance(min_limit, int) or not isinstance(max_limit, int): return None # TODO handle timestamp and dates: https://github.com/databrickslabs/dqx/issues/71 @@ -92,36 +95,50 @@ def dq_generate_min_max(column: str, filter: str | None, level: str = "error", * "column": column, "min_limit": val_maybe_to_str(min_limit, include_sql_quotes=False), "max_limit": val_maybe_to_str(max_limit, include_sql_quotes=False), + "filter": params.get("filter", None) }, }, "filter": filter, "name": f"{column}_isnt_in_range", - "criticality": level, - "filter": filter, + "criticality": level } if max_limit is not None: return { "check": { "function": "is_not_greater_than", +<<<<<<< HEAD "arguments": {"column": column, "limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}, +======= + "arguments": { + "column": column, + "limit": val_maybe_to_str(max_limit, include_sql_quotes=False), + "filter": params.get("filter", None) + }, +>>>>>>> 07d81a3 (Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.) }, "filter": filter, "name": f"{column}_not_greater_than", - "criticality": level, - "filter": filter, + "criticality": level } if min_limit is not None: return { "check": { "function": "is_not_less_than", +<<<<<<< HEAD "arguments": {"column": column, "limit": val_maybe_to_str(min_limit, include_sql_quotes=False)}, +======= + "arguments": { + "column": column, + "limit": val_maybe_to_str(min_limit, include_sql_quotes=False), + "filter": params.get("filter", None) + }, +>>>>>>> 07d81a3 (Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.) }, "filter": filter, "name": f"{column}_not_less_than", - "criticality": level, - "filter": filter, + "criticality": level } return None @@ -141,6 +158,7 @@ def dq_generate_is_not_null(column: str, filter: str | None, level: str = "error """ params = params or {} <<<<<<< HEAD +<<<<<<< HEAD ======= filter = params.get("dataset_filter_expression", None) @@ -148,9 +166,13 @@ def dq_generate_is_not_null(column: str, filter: str | None, level: str = "error return { "check": {"function": "is_not_null", "arguments": {"column": column}}, "filter": filter, +======= + filter = params.get("filter", None) + return { + "check": {"function": "is_not_null", "arguments": {"column": column, "filter": params.get("filter", None)}}, +>>>>>>> 07d81a3 (Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.) "name": f"{column}_is_null", - "criticality": level, - "filter": filter, + "criticality": level } @staticmethod @@ -167,19 +189,22 @@ def dq_generate_is_not_null_or_empty(column: str, filter: str | None, level: str A dictionary representing the data quality rule. """ <<<<<<< HEAD +<<<<<<< HEAD ======= filter = params.get("dataset_filter_expression", None) >>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) +======= + filter = params.get("filter", None) +>>>>>>> 07d81a3 (Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.) return { "check": { "function": "is_not_null_and_not_empty", - "arguments": {"column": column, "trim_strings": params.get("trim_strings", True)}, + "arguments": {"column": column, "trim_strings": params.get("trim_strings", True), "filter": params.get("filter", None)}, }, "filter": filter, "name": f"{column}_is_null_or_empty", - "criticality": level, - "filter": filter, + "criticality": level } _checks_mapping = { diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 51517f60c..035234f94 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -32,7 +32,12 @@ class DQProfile: description: str | None = None parameters: dict[str, Any] | None = None filter: str | None = None +<<<<<<< HEAD +======= + + +>>>>>>> 07d81a3 (Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.) class DQProfiler(DQEngineBase): """Data Quality Profiler class to profile input data.""" @@ -57,6 +62,7 @@ def __init__(self, workspace_client: WorkspaceClient, spark: SparkSession | None <<<<<<< HEAD <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD "filter": None, # filter to apply to the dataset ======= @@ -72,6 +78,9 @@ def __init__(self, workspace_client: WorkspaceClient, spark: SparkSession | None ======= "dataset_filter_expression": None, # filter to apply to the dataset >>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) +======= + "filter": None, # filter to apply to the dataset +>>>>>>> 07d81a3 (Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.) } @staticmethod @@ -142,10 +151,18 @@ def profile( return summary_stats, dq_rules self._profile(df, df_columns, dq_rules, options, summary_stats, total_count) +<<<<<<< HEAD filter = options.get("filter", None) if filter: for rule in dq_rules: rule.filter = filter +======= + filter=options.get("filter",None) + if filter: + for rule in dq_rules: + rule.filter=filter + +>>>>>>> 07d81a3 (Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.) return summary_stats, dq_rules @@ -352,6 +369,7 @@ def _sample(df: DataFrame, opts: dict[str, Any]) -> DataFrame: sample_seed = opts.get("sample_seed", None) limit = opts.get("limit", None) <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD filter = opts.get("filter", None) @@ -374,6 +392,9 @@ def _sample(df: DataFrame, opts: dict[str, Any]) -> DataFrame: >>>>>>> c77a6eb (Refactor _sample method to utilize return the filtered dataframe) ======= filter = opts.get("dataset_filter_expression", None) +======= + filter = opts.get("filter", None) +>>>>>>> 07d81a3 (Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.) if filter: df = df.filter(filter) @@ -428,8 +449,7 @@ def _calculate_metrics( """ max_nulls = opts.get("max_null_ratio", 0) trim_strings = opts.get("trim_strings", True) - filter = opts.get("dataset_filter_expression", None) - + dst = df.select(field_name).dropna() if typ == T.StringType() and trim_strings: col_name = dst.columns[0] @@ -447,13 +467,12 @@ def _calculate_metrics( name="is_not_null", column=field_name, description=f"Column {field_name} has {null_percentage * 100:.1f}% of null values " - f"(allowed {max_nulls * 100:.1f}%)", - parameters={"dataset_filter_expression": filter}, + f"(allowed {max_nulls * 100:.1f}%)", ) ) else: dq_rules.append( - DQProfile(name="is_not_null", column=field_name, parameters={"dataset_filter_expression": filter}) + DQProfile(name="is_not_null", column=field_name) ) if self._type_supports_distinct(typ): dst2 = dst.dropDuplicates() @@ -476,7 +495,7 @@ def _calculate_metrics( DQProfile( name="is_in", column=field_name, - parameters={"in": [row[0] for row in dst2.collect()], "dataset_filter_expression": filter}, + parameters={"in": [row[0] for row in dst2.collect()]}, ) >>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) ) @@ -506,7 +525,7 @@ def _calculate_metrics( DQProfile( name="is_not_null_or_empty", column=field_name, - parameters={"trim_strings": trim_strings, "dataset_filter_expression": filter}, + parameters={"trim_strings": trim_strings}, ) >>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) ) @@ -626,7 +645,7 @@ def _extract_min_max( if opts is None: opts = {} - filter = opts.get("dataset_filter_expression", None) + outlier_cols = opts.get("outlier_columns", []) column = dst.columns[0] if opts.get("remove_outliers", True) and ( @@ -674,7 +693,7 @@ def _extract_min_max( ======= name="min_max", column=col_name, - parameters={"min": min_limit, "max": max_limit, "dataset_filter_expression": filter}, + parameters={"min": min_limit, "max": max_limit}, description=descr, >>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) ) diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index 442a67733..ffc0c03b7 100644 --- a/tests/integration/test_profiler.py +++ b/tests/integration/test_profiler.py @@ -1522,7 +1522,7 @@ def test_profile_with_dataset_filter(spark, ws): custom_options = { "sample_fraction": None, "limit": None, - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", } profiler = DQProfiler(ws) @@ -1539,26 +1539,21 @@ def test_profile_with_dataset_filter(spark, ws): DQProfile( name="is_not_null", column="machine_id", - description=None, - parameters={ - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - }, + description=None, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + ), DQProfile( name="is_not_null", column="maintenance_type", description=None, - parameters={ - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - }, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), DQProfile( name="is_not_null", column="maintenance_date", description=None, - parameters={ - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - }, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), DQProfile( name="min_max", @@ -1566,17 +1561,14 @@ def test_profile_with_dataset_filter(spark, ws): description="Real min/max values were used", parameters={ "min": date(2025, 4, 29), - "max": date(2025, 7, 30), - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", - }, - ), + "max": date(2025, 7, 30)}, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + ), DQProfile( name="is_not_null", column="cost", description=None, - parameters={ - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - }, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), DQProfile( name="min_max", @@ -1584,17 +1576,15 @@ def test_profile_with_dataset_filter(spark, ws): parameters={ "min": Decimal('100.00'), "max": Decimal('300.00'), - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", }, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", description="Real min/max values were used", ), DQProfile( name="is_not_null", column="next_scheduled_date", description=None, - parameters={ - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - }, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), DQProfile( name="min_max", @@ -1602,17 +1592,14 @@ def test_profile_with_dataset_filter(spark, ws): description="Real min/max values were used", parameters={ "min": date(2025, 7, 15), - "max": date(2025, 12, 1), - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", - }, + "max": date(2025, 12, 1)}, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), DQProfile( name="is_not_null", column="safety_check_passed", description=None, - parameters={ - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - }, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", ), ] diff --git a/tests/integration/test_rules_generator.py b/tests/integration/test_rules_generator.py index 423740273..c5faa808f 100644 --- a/tests/integration/test_rules_generator.py +++ b/tests/integration/test_rules_generator.py @@ -11,7 +11,7 @@ DQProfile( name="is_not_null", column="vendor_id", description="Column vendor_id has 0.3% of null values (allowed 1.0%)" ), - DQProfile(name="is_in", column="vendor_id", parameters={"in": ["1", "4", "2"]}), + DQProfile(name="is_in", column="vendor_id",parameters={"in": ["1", "4", "2"]}), DQProfile(name="is_not_null_or_empty", column="vendor_id", parameters={"trim_strings": True}), DQProfile( name="min_max", @@ -211,7 +211,7 @@ def test_generate_dq_rules_dataframe_filter(ws): DQProfile(name="is_not_null", column="machine_id", description=None, - parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + parameters={"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), DQProfile( name="min_max", @@ -219,96 +219,96 @@ def test_generate_dq_rules_dataframe_filter(ws): description="Real min/max values were used", parameters={"min": datetime.date(2025, 4, 29), "max": datetime.date(2025, 7, 30), - "dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, ), DQProfile(name="is_not_null", column="cost", description=None, - parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + parameters={"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), DQProfile( name="min_max", column="cost", description="Real min/max values were used", - parameters={"min": Decimal('100.00') , "max": Decimal('300.00'),"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + parameters={"min": Decimal('100.00') , "max": Decimal('300.00'),"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, ), DQProfile(name="is_not_null", column="next_scheduled_date", description=None, - parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + parameters={"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), DQProfile( name="min_max", column="next_scheduled_date", description="Real min/max values were used", - parameters={"min": datetime.date(2025, 7, 15), "max": datetime.date(2025, 12, 1),"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + parameters={"min": datetime.date(2025, 7, 15), "max": datetime.date(2025, 12, 1),"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, ), DQProfile(name="is_not_null", column="safety_check_passed", description=None, - parameters={"dataset_filter_expression": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + parameters={"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), ] expectations = generator.generate_dq_rules(test_rules) expected = [ { - "check": {"function": "is_not_null", "arguments": {"column": "machine_id"}}, + "check": {"function": "is_not_null", "arguments": {"column": "machine_id","filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}}, "name": "machine_id_is_null", - "criticality": "Error", - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + "criticality": "error" + }, { "check": { "function": "min_max", - "arguments": {"column": "maintenance_date","min": datetime.date(2025, 4, 29), "max": datetime.date(2025, 7, 30)} + "arguments": {"column": "maintenance_date","min": datetime.date(2025, 4, 29), "max": datetime.date(2025, 7, 30), "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} }, "name": "maintenance_date_isnt_in_range", - "criticality": "Error", - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + "criticality": "error", + }, { "check": { "function": "is_not_null", - "arguments": {"column": "cost"}}, + "arguments": {"column": "cost","filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}}, "name": "cost_is_null", - "criticality": "Error", - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + "criticality": "error", + }, { "check": { "function": "min_max", - "arguments": {"column": "cost","min": Decimal('100.00') , "max": Decimal('300.00')} + "arguments": {"column": "cost","min": Decimal('100.00') , "max": Decimal('300.00'),"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'",} }, "name": "cost_isnt_in_range", - "criticality": "Error", - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + "criticality": "error", + }, { "check": { "function": "is_not_null", - "arguments": {"column": "next_scheduled_date"}}, + "arguments": {"column": "next_scheduled_date","filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}}, "name": "next_scheduled_date_is_null", - "criticality": "Error", - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + "criticality": "error", + }, { "check": { "function": "min_max", - "arguments": {"column": "next_scheduled_date","min": datetime.date(2025, 7, 15), "max": datetime.date(2025, 12, 1)} + "arguments": {"column": "next_scheduled_date","min": datetime.date(2025, 7, 15), "max": datetime.date(2025, 12, 1), "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} }, "name": "cost_isnt_in_range", - "criticality": "Error", - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + "criticality": "error", + }, { "check": { "function": "is_not_null", - "arguments": {"column": "safety_check_passed"}}, + "arguments": {"column": "safety_check_passed", "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}}, "name": "safety_check_passed_is_null", - "criticality": "Error", - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + "criticality": "error", + } >>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) From 4494393ba9fa87bb8a5bde158b7fdc38d3385d8e Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Sat, 20 Sep 2025 23:34:49 -0300 Subject: [PATCH 20/34] Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict. --- src/databricks/labs/dqx/profiler/generator.py | 66 ++++++++- tests/integration/test_rules_generator.py | 140 ++++++++++++------ 2 files changed, 161 insertions(+), 45 deletions(-) diff --git a/src/databricks/labs/dqx/profiler/generator.py b/src/databricks/labs/dqx/profiler/generator.py index 82c51dba0..e689fba90 100644 --- a/src/databricks/labs/dqx/profiler/generator.py +++ b/src/databricks/labs/dqx/profiler/generator.py @@ -25,7 +25,11 @@ def generate_dq_rules(self, profiles: list[DQProfile] | None = None, level: str if profiles is None: profiles = [] dq_rules = [] +<<<<<<< HEAD +======= + +>>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) for profile in profiles: rule_name = profile.name column = profile.column @@ -34,7 +38,11 @@ def generate_dq_rules(self, profiles: list[DQProfile] | None = None, level: str if rule_name not in self._checks_mapping: logger.info(f"No rule '{rule_name}' for column '{column}'. skipping...") continue +<<<<<<< HEAD expr = self._checks_mapping[rule_name](column, filter, level, **params) +======= + expr = self._checks_mapping[rule_name](column, filter,level,**params) +>>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) if expr: dq_rules.append(expr) @@ -44,7 +52,11 @@ def generate_dq_rules(self, profiles: list[DQProfile] | None = None, level: str return dq_rules @staticmethod +<<<<<<< HEAD def dq_generate_is_in(column: str, filter: str | None, level: str = "error", **params: dict): +======= + def dq_generate_is_in(column: str, filter:str|None,level: str = "error",**params: dict ): +>>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) """ Generates a data quality rule to check if a column's value is in a specified list. @@ -57,18 +69,26 @@ def dq_generate_is_in(column: str, filter: str | None, level: str = "error", **p A dictionary representing the data quality rule. """ return { +<<<<<<< HEAD <<<<<<< HEAD "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"]}}, "filter": filter, ======= "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"],"filter": params.get("filter", None)}}, >>>>>>> 07d81a3 (Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.) +======= + "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"]},"filter":filter}, +>>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) "name": f"{column}_other_value", "criticality": level } @staticmethod +<<<<<<< HEAD def dq_generate_min_max(column: str, filter: str | None, level: str = "error", **params: dict): +======= + def dq_generate_min_max(column: str, filter:str|None,level: str = "error", **params: dict): +>>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) """ Generates a data quality rule to check if a column's value is within a specified range. @@ -82,7 +102,7 @@ def dq_generate_min_max(column: str, filter: str | None, level: str = "error", * """ min_limit = params.get("min") max_limit = params.get("max") - filter = params.get("filter", None) + if not isinstance(min_limit, int) or not isinstance(max_limit, int): return None # TODO handle timestamp and dates: https://github.com/databrickslabs/dqx/issues/71 @@ -94,6 +114,7 @@ def dq_generate_min_max(column: str, filter: str | None, level: str = "error", * "arguments": { "column": column, "min_limit": val_maybe_to_str(min_limit, include_sql_quotes=False), +<<<<<<< HEAD "max_limit": val_maybe_to_str(max_limit, include_sql_quotes=False), "filter": params.get("filter", None) }, @@ -101,6 +122,12 @@ def dq_generate_min_max(column: str, filter: str | None, level: str = "error", * "filter": filter, "name": f"{column}_isnt_in_range", "criticality": level +======= + "max_limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}, + "filter": filter}, + "name": f"{column}_isnt_in_range", + "criticality": level +>>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) } if max_limit is not None: @@ -112,6 +139,7 @@ def dq_generate_min_max(column: str, filter: str | None, level: str = "error", * ======= "arguments": { "column": column, +<<<<<<< HEAD "limit": val_maybe_to_str(max_limit, include_sql_quotes=False), "filter": params.get("filter", None) }, @@ -120,6 +148,12 @@ def dq_generate_min_max(column: str, filter: str | None, level: str = "error", * "filter": filter, "name": f"{column}_not_greater_than", "criticality": level +======= + "limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}, + "filter": filter}, + "name": f"{column}_not_greater_than", + "criticality": level +>>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) } if min_limit is not None: @@ -131,6 +165,7 @@ def dq_generate_min_max(column: str, filter: str | None, level: str = "error", * ======= "arguments": { "column": column, +<<<<<<< HEAD "limit": val_maybe_to_str(min_limit, include_sql_quotes=False), "filter": params.get("filter", None) }, @@ -139,12 +174,22 @@ def dq_generate_min_max(column: str, filter: str | None, level: str = "error", * "filter": filter, "name": f"{column}_not_less_than", "criticality": level +======= + "limit": val_maybe_to_str(min_limit, include_sql_quotes=False)}, + "filter": filter}, + "name": f"{column}_not_less_than", + "criticality": level +>>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) } return None @staticmethod +<<<<<<< HEAD def dq_generate_is_not_null(column: str, filter: str | None, level: str = "error", **params: dict): +======= + def dq_generate_is_not_null(column: str, filter:str|None,level: str = "error", **params: dict): +>>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) """ Generates a data quality rule to check if a column's value is not null. @@ -159,6 +204,7 @@ def dq_generate_is_not_null(column: str, filter: str | None, level: str = "error params = params or {} <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD ======= filter = params.get("dataset_filter_expression", None) @@ -171,12 +217,21 @@ def dq_generate_is_not_null(column: str, filter: str | None, level: str = "error return { "check": {"function": "is_not_null", "arguments": {"column": column, "filter": params.get("filter", None)}}, >>>>>>> 07d81a3 (Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.) +======= + + return { + "check": {"function": "is_not_null", "arguments": {"column": column},"filter":filter}, +>>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) "name": f"{column}_is_null", "criticality": level } @staticmethod +<<<<<<< HEAD def dq_generate_is_not_null_or_empty(column: str, filter: str | None, level: str = "error", **params: dict): +======= + def dq_generate_is_not_null_or_empty(column: str, filter:str|None,level: str = "error", **params: dict): +>>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) """ Generates a data quality rule to check if a column's value is not null or empty. @@ -190,6 +245,7 @@ def dq_generate_is_not_null_or_empty(column: str, filter: str | None, level: str """ <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD ======= filter = params.get("dataset_filter_expression", None) @@ -203,6 +259,14 @@ def dq_generate_is_not_null_or_empty(column: str, filter: str | None, level: str "arguments": {"column": column, "trim_strings": params.get("trim_strings", True), "filter": params.get("filter", None)}, }, "filter": filter, +======= + + return { + "check": { + "function": "is_not_null_and_not_empty", + "arguments": {"column": column, "trim_strings": params.get("trim_strings", True)}, + "filter": filter}, +>>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) "name": f"{column}_is_null_or_empty", "criticality": level } diff --git a/tests/integration/test_rules_generator.py b/tests/integration/test_rules_generator.py index c5faa808f..682ba6787 100644 --- a/tests/integration/test_rules_generator.py +++ b/tests/integration/test_rules_generator.py @@ -46,21 +46,32 @@ def test_generate_dq_rules(ws): expectations = generator.generate_dq_rules(test_rules) expected = [ { +<<<<<<< HEAD "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"}}, "filter": None, +======= + "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"},"filter":None}, +>>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) "name": "vendor_id_is_null", "criticality": "error", }, { +<<<<<<< HEAD "check": {"function": "is_in_list", "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}}, "filter": None, +======= + "check": { + "function": "is_in_list", + "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]},"filter":None, + }, +>>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) "name": "vendor_id_other_value", "criticality": "error", }, { "check": { "function": "is_not_null_and_not_empty", - "arguments": {"column": "vendor_id", "trim_strings": True}, + "arguments": {"column": "vendor_id", "trim_strings": True},"filter":None }, "filter": None, "name": "vendor_id_is_null_or_empty", @@ -70,6 +81,7 @@ def test_generate_dq_rules(ws): "check": { "function": "is_in_range", "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}, + "filter":None }, "filter": None, "name": "rate_code_id_isnt_in_range", @@ -84,21 +96,32 @@ def test_generate_dq_rules_warn(ws): expectations = generator.generate_dq_rules(test_rules, level="warn") expected = [ { +<<<<<<< HEAD "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"}}, "filter": None, +======= + "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"},"filter":None}, +>>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) "name": "vendor_id_is_null", "criticality": "warn", }, { +<<<<<<< HEAD "check": {"function": "is_in_list", "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}}, "filter": None, +======= + "check": { + "function": "is_in_list", + "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]},"filter":None, + }, +>>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) "name": "vendor_id_other_value", "criticality": "warn", }, { "check": { "function": "is_not_null_and_not_empty", - "arguments": {"column": "vendor_id", "trim_strings": True}, + "arguments": {"column": "vendor_id", "trim_strings": True},"filter":None, }, "filter": None, "name": "vendor_id_is_null_or_empty", @@ -107,7 +130,7 @@ def test_generate_dq_rules_warn(ws): { "check": { "function": "is_in_range", - "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}, + "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265},"filter":None, }, "filter": None, "name": "rate_code_id_isnt_in_range", @@ -211,105 +234,134 @@ def test_generate_dq_rules_dataframe_filter(ws): DQProfile(name="is_not_null", column="machine_id", description=None, - parameters={"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"), + DQProfile(name="is_in", + column="vendor_id", + parameters={"in": ["1", "4", "2"]}, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"), DQProfile( name="min_max", column="maintenance_date", description="Real min/max values were used", parameters={"min": datetime.date(2025, 4, 29), - "max": datetime.date(2025, 7, 30), - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + "max": datetime.date(2025, 7, 30)}, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", ), DQProfile(name="is_not_null", column="cost", description=None, - parameters={"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), DQProfile( name="min_max", column="cost", description="Real min/max values were used", - parameters={"min": Decimal('100.00') , "max": Decimal('300.00'),"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + parameters={"min": Decimal('100.00') , "max": Decimal('300.00')}, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), DQProfile(name="is_not_null", column="next_scheduled_date", description=None, - parameters={"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), - DQProfile( - name="min_max", - column="next_scheduled_date", - description="Real min/max values were used", - parameters={"min": datetime.date(2025, 7, 15), "max": datetime.date(2025, 12, 1),"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, - ), + filter= "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"), + DQProfile(name="is_not_null", column="safety_check_passed", description=None, - parameters={"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} ), + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), + + DQProfile(name="is_not_null_or_empty", + column="vendor_id", + parameters={"trim_strings": True}, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), ] expectations = generator.generate_dq_rules(test_rules) - + expected = [ { - "check": {"function": "is_not_null", "arguments": {"column": "machine_id","filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}}, + "check": {"function": "is_not_null", + "arguments": {"column": "machine_id"}, + "filter":"machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, "name": "machine_id_is_null", "criticality": "error" }, { - "check": { - "function": "min_max", - "arguments": {"column": "maintenance_date","min": datetime.date(2025, 4, 29), "max": datetime.date(2025, 7, 30), "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} - }, - "name": "maintenance_date_isnt_in_range", + 'check': {"function": "is_in_list", + "arguments": {"allowed": ['1','4','2'],"column": "vendor_id"}, + "filter":"machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, "criticality": "error", - + "name": "vendor_id_other_value" }, + + # { + # "check": { + # "function": "min_max", + # "arguments": {"column": "maintenance_date","min": datetime.date(2025, 4, 29), "max": datetime.date(2025, 7, 30)}, + # "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + # }, + # "name": "maintenance_date_isnt_in_range", + # "criticality": "error", + + # }, { "check": { "function": "is_not_null", - "arguments": {"column": "cost","filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}}, + "arguments": {"column": "cost"}, + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, "name": "cost_is_null", "criticality": "error", }, - { - "check": { - "function": "min_max", - "arguments": {"column": "cost","min": Decimal('100.00') , "max": Decimal('300.00'),"filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'",} - }, - "name": "cost_isnt_in_range", - "criticality": "error", + # { + # "check": { + # "function": "min_max", + # "arguments": {"column": "cost","min": Decimal('100.00') , "max": Decimal('300.00')}, + # "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + # }, + # "name": "cost_isnt_in_range", + # "criticality": "error", - }, + # }, { "check": { "function": "is_not_null", - "arguments": {"column": "next_scheduled_date","filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}}, + "arguments": {"column": "next_scheduled_date"}, + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, "name": "next_scheduled_date_is_null", "criticality": "error", }, - { - "check": { - "function": "min_max", - "arguments": {"column": "next_scheduled_date","min": datetime.date(2025, 7, 15), "max": datetime.date(2025, 12, 1), "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"} - }, - "name": "cost_isnt_in_range", - "criticality": "error", + # { + # "check": { + # "function": "min_max", + # "arguments": {"column": "next_scheduled_date","min": datetime.date(2025, 7, 15), "max": datetime.date(2025, 12, 1)}, + # "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + # }, + # "name": "cost_isnt_in_range", + # "criticality": "error", - }, + # }, { "check": { "function": "is_not_null", - "arguments": {"column": "safety_check_passed", "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}}, + "arguments": {"column": "safety_check_passed"}, + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, "name": "safety_check_passed_is_null", "criticality": "error", - } + }, + { + "check": { + "function": "is_not_null_and_not_empty", + "arguments": {"column": "vendor_id", "trim_strings": True}, + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + "name": "vendor_id_is_null_or_empty", + "criticality": "error" + } >>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) ] From bb5977d3ca10e633bd6959fcce8be9a1124c091f Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Mon, 22 Sep 2025 10:49:50 -0300 Subject: [PATCH 21/34] Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig. --- src/databricks/labs/dqx/profiler/generator.py | 30 ++++ .../test_load_checks_from_uc_volume.py | 26 ++++ tests/integration/test_profiler.py | 29 +++- tests/integration/test_rules_generator.py | 133 +++++++++--------- .../test_save_and_load_checks_from_table.py | 7 + .../test_save_checks_to_workspace_file.py | 15 +- 6 files changed, 169 insertions(+), 71 deletions(-) diff --git a/src/databricks/labs/dqx/profiler/generator.py b/src/databricks/labs/dqx/profiler/generator.py index e689fba90..cdd6346ef 100644 --- a/src/databricks/labs/dqx/profiler/generator.py +++ b/src/databricks/labs/dqx/profiler/generator.py @@ -70,6 +70,7 @@ def dq_generate_is_in(column: str, filter:str|None,level: str = "error",**params """ return { <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"]}}, "filter": filter, @@ -79,6 +80,10 @@ def dq_generate_is_in(column: str, filter:str|None,level: str = "error",**params ======= "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"]},"filter":filter}, >>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) +======= + "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"]}}, + "filter":filter, +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": f"{column}_other_value", "criticality": level } @@ -114,6 +119,7 @@ def dq_generate_min_max(column: str, filter:str|None,level: str = "error", **par "arguments": { "column": column, "min_limit": val_maybe_to_str(min_limit, include_sql_quotes=False), +<<<<<<< HEAD <<<<<<< HEAD "max_limit": val_maybe_to_str(max_limit, include_sql_quotes=False), "filter": params.get("filter", None) @@ -125,6 +131,10 @@ def dq_generate_min_max(column: str, filter:str|None,level: str = "error", **par ======= "max_limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}, "filter": filter}, +======= + "max_limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}}, + "filter": filter, +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": f"{column}_isnt_in_range", "criticality": level >>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) @@ -139,6 +149,7 @@ def dq_generate_min_max(column: str, filter:str|None,level: str = "error", **par ======= "arguments": { "column": column, +<<<<<<< HEAD <<<<<<< HEAD "limit": val_maybe_to_str(max_limit, include_sql_quotes=False), "filter": params.get("filter", None) @@ -151,6 +162,10 @@ def dq_generate_min_max(column: str, filter:str|None,level: str = "error", **par ======= "limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}, "filter": filter}, +======= + "limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}}, + "filter": filter, +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": f"{column}_not_greater_than", "criticality": level >>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) @@ -165,6 +180,7 @@ def dq_generate_min_max(column: str, filter:str|None,level: str = "error", **par ======= "arguments": { "column": column, +<<<<<<< HEAD <<<<<<< HEAD "limit": val_maybe_to_str(min_limit, include_sql_quotes=False), "filter": params.get("filter", None) @@ -177,6 +193,10 @@ def dq_generate_min_max(column: str, filter:str|None,level: str = "error", **par ======= "limit": val_maybe_to_str(min_limit, include_sql_quotes=False)}, "filter": filter}, +======= + "limit": val_maybe_to_str(min_limit, include_sql_quotes=False)}}, + "filter": filter, +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": f"{column}_not_less_than", "criticality": level >>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) @@ -220,8 +240,13 @@ def dq_generate_is_not_null(column: str, filter:str|None,level: str = "error", * ======= return { +<<<<<<< HEAD "check": {"function": "is_not_null", "arguments": {"column": column},"filter":filter}, >>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) +======= + "check": {"function": "is_not_null", "arguments": {"column": column}}, + "filter":filter, +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": f"{column}_is_null", "criticality": level } @@ -264,9 +289,14 @@ def dq_generate_is_not_null_or_empty(column: str, filter:str|None,level: str = " return { "check": { "function": "is_not_null_and_not_empty", +<<<<<<< HEAD "arguments": {"column": column, "trim_strings": params.get("trim_strings", True)}, "filter": filter}, >>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) +======= + "arguments": {"column": column, "trim_strings": params.get("trim_strings", True)}}, + "filter": filter, +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": f"{column}_is_null_or_empty", "criticality": level } diff --git a/tests/integration/test_load_checks_from_uc_volume.py b/tests/integration/test_load_checks_from_uc_volume.py index 6659d78f4..acd164db2 100644 --- a/tests/integration/test_load_checks_from_uc_volume.py +++ b/tests/integration/test_load_checks_from_uc_volume.py @@ -9,33 +9,59 @@ TEST_CHECKS = [ { "criticality": "error", +<<<<<<< HEAD "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}, "filter": None}, +======= + "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}, "filter":None}, +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) }, { "check": { "function": "is_not_null", "arguments": {"column": "next_scheduled_date"}, +<<<<<<< HEAD "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", }, "name": "next_scheduled_date_is_null", "criticality": "Error", }, +======= + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + + "name": "next_scheduled_date_is_null", + "criticality": "Error", + + } +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) ] EXPECTED_CHECKS = [ { "criticality": "error", +<<<<<<< HEAD "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}, "filter": None}, +======= + "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {},"filter":None}, +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) }, { "check": { "function": "is_not_null", "arguments": {"column": "next_scheduled_date"}, +<<<<<<< HEAD "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", }, "name": "next_scheduled_date_is_null", "criticality": "Error", }, +======= + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + + "name": "next_scheduled_date_is_null", + "criticality": "Error", + + } +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) ] diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index ffc0c03b7..87c26311e 100644 --- a/tests/integration/test_profiler.py +++ b/tests/integration/test_profiler.py @@ -69,11 +69,15 @@ def test_profiler(spark, ws): expected_rules = [ DQProfile(name="is_not_null", column="t1", description=None, parameters=None, filter=None), DQProfile( +<<<<<<< HEAD name="min_max", column="t1", description="Real min/max values were used", parameters={"min": 1, "max": 3}, filter=None, +======= + name="min_max", column="t1", description="Real min/max values were used", parameters={"min": 1, "max": 3}, filter=None +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) ), DQProfile(name='is_not_null', column='d1', description=None, parameters=None, filter=None), DQProfile( @@ -81,11 +85,17 @@ def test_profiler(spark, ws): column='d1', description='Real min/max values were used', parameters={'max': Decimal('333323.00'), 'min': Decimal('1.23')}, +<<<<<<< HEAD filter=None, ), DQProfile( name='is_not_null_or_empty', column='t2', description=None, parameters={'trim_strings': True}, filter=None ), +======= + filter=None + ), + DQProfile(name='is_not_null_or_empty', column='t2', description=None, parameters={'trim_strings': True},filter=None), +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) DQProfile(name="is_not_null", column="s1.ns1", description=None, parameters=None), DQProfile( name="min_max", @@ -95,18 +105,31 @@ def test_profiler(spark, ws): "min": datetime(2023, 1, 6, 0, 0, tzinfo=timezone.utc), "max": datetime(2023, 1, 9, 0, 0, tzinfo=timezone.utc), }, +<<<<<<< HEAD filter=None, ), DQProfile(name="is_not_null", column="s1.s2.ns2", description=None, parameters=None, filter=None), DQProfile(name="is_not_null", column="s1.s2.ns3", description=None, parameters=None, filter=None), +======= + filter=None + ), + DQProfile(name="is_not_null", column="s1.s2.ns2", description=None, parameters=None,filter=None), + DQProfile(name="is_not_null", column="s1.s2.ns3", description=None, parameters=None,filter=None), +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) DQProfile( name="min_max", column="s1.s2.ns3", description="Real min/max values were used", parameters={"min": date(2023, 1, 6), "max": date(2023, 1, 8)}, +<<<<<<< HEAD filter=None, ), DQProfile(name="is_not_null", column="b1", description=None, parameters=None, filter=None), +======= + filter=None + ), + DQProfile(name="is_not_null", column="b1", description=None, parameters=None,filter=None), +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) ] print(stats) assert len(stats.keys()) > 0 @@ -1317,6 +1340,7 @@ def test_profile_with_dataset_filter(spark, ws): "preventive", date(2025, 4, 29), Decimal("300.50"), +<<<<<<< HEAD date(2025, 7, 15), True, ), @@ -1488,6 +1512,8 @@ def test_profile_with_dataset_filter(spark, ws): "preventive", date(2025, 4, 29), Decimal("300.00"), +======= +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) date(2025, 7, 15), True, ), @@ -1521,6 +1547,7 @@ def test_profile_with_dataset_filter(spark, ws): custom_options = { "sample_fraction": None, + "round":False, "limit": None, "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", } @@ -1575,7 +1602,7 @@ def test_profile_with_dataset_filter(spark, ws): column="cost", parameters={ "min": Decimal('100.00'), - "max": Decimal('300.00'), + "max": Decimal('300.50'), }, filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", description="Real min/max values were used", diff --git a/tests/integration/test_rules_generator.py b/tests/integration/test_rules_generator.py index 682ba6787..3012d9268 100644 --- a/tests/integration/test_rules_generator.py +++ b/tests/integration/test_rules_generator.py @@ -46,12 +46,17 @@ def test_generate_dq_rules(ws): expectations = generator.generate_dq_rules(test_rules) expected = [ { +<<<<<<< HEAD <<<<<<< HEAD "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"}}, "filter": None, ======= "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"},"filter":None}, >>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) +======= + "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"}}, + "filter":None, +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": "vendor_id_is_null", "criticality": "error", }, @@ -62,28 +67,42 @@ def test_generate_dq_rules(ws): ======= "check": { "function": "is_in_list", +<<<<<<< HEAD "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]},"filter":None, }, >>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) +======= + "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}}, + "filter":None, +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": "vendor_id_other_value", "criticality": "error", }, { "check": { "function": "is_not_null_and_not_empty", - "arguments": {"column": "vendor_id", "trim_strings": True},"filter":None + "arguments": {"column": "vendor_id", "trim_strings": True} }, +<<<<<<< HEAD "filter": None, +======= + "filter":None, +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": "vendor_id_is_null_or_empty", "criticality": "error", }, { "check": { "function": "is_in_range", +<<<<<<< HEAD "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}, "filter":None }, "filter": None, +======= + "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}}, + "filter":None, +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": "rate_code_id_isnt_in_range", "criticality": "error", }, @@ -96,12 +115,17 @@ def test_generate_dq_rules_warn(ws): expectations = generator.generate_dq_rules(test_rules, level="warn") expected = [ { +<<<<<<< HEAD <<<<<<< HEAD "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"}}, "filter": None, ======= "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"},"filter":None}, >>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) +======= + "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"}}, + "filter":None, +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": "vendor_id_is_null", "criticality": "warn", }, @@ -112,27 +136,42 @@ def test_generate_dq_rules_warn(ws): ======= "check": { "function": "is_in_list", +<<<<<<< HEAD "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]},"filter":None, }, >>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) +======= + "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}}, + "filter":None, +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": "vendor_id_other_value", "criticality": "warn", }, { "check": { "function": "is_not_null_and_not_empty", +<<<<<<< HEAD "arguments": {"column": "vendor_id", "trim_strings": True},"filter":None, }, "filter": None, +======= + "arguments": {"column": "vendor_id", "trim_strings": True}}, + "filter":None, +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": "vendor_id_is_null_or_empty", "criticality": "warn", }, { "check": { "function": "is_in_range", +<<<<<<< HEAD "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265},"filter":None, }, "filter": None, +======= + "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}}, + "filter":None, +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": "rate_code_id_isnt_in_range", "criticality": "warn", }, @@ -240,39 +279,25 @@ def test_generate_dq_rules_dataframe_filter(ws): column="vendor_id", parameters={"in": ["1", "4", "2"]}, filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"), - DQProfile( - name="min_max", - column="maintenance_date", - description="Real min/max values were used", - parameters={"min": datetime.date(2025, 4, 29), - "max": datetime.date(2025, 7, 30)}, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", - ), + DQProfile(name="is_not_null", column="cost", description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), - DQProfile( - name="min_max", - column="cost", - description="Real min/max values were used", - parameters={"min": Decimal('100.00') , "max": Decimal('300.00')}, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - ), + ), + DQProfile(name="is_not_null", column="next_scheduled_date", description=None, - filter= "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"), + ), DQProfile(name="is_not_null", column="safety_check_passed", description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), + ), DQProfile(name="is_not_null_or_empty", column="vendor_id", - parameters={"trim_strings": True}, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" ), + parameters={"trim_strings": True}), ] expectations = generator.generate_dq_rules(test_rules) @@ -280,75 +305,47 @@ def test_generate_dq_rules_dataframe_filter(ws): expected = [ { "check": {"function": "is_not_null", - "arguments": {"column": "machine_id"}, - "filter":"machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + "arguments": {"column": "machine_id"}}, + "filter":"machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", "name": "machine_id_is_null", "criticality": "error" }, { 'check': {"function": "is_in_list", - "arguments": {"allowed": ['1','4','2'],"column": "vendor_id"}, - "filter":"machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, - "criticality": "error", + "arguments": {"allowed": ['1','4','2'],"column": "vendor_id"}}, + + "filter":"machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + "criticality": "error", "name": "vendor_id_other_value" - }, - - # { - # "check": { - # "function": "min_max", - # "arguments": {"column": "maintenance_date","min": datetime.date(2025, 4, 29), "max": datetime.date(2025, 7, 30)}, - # "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - # }, - # "name": "maintenance_date_isnt_in_range", - # "criticality": "error", - - # }, + }, { "check": { "function": "is_not_null", - "arguments": {"column": "cost"}, - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, - + "arguments": {"column": "cost"}}, + + "filter": None, "name": "cost_is_null", "criticality": "error", }, - # { - # "check": { - # "function": "min_max", - # "arguments": {"column": "cost","min": Decimal('100.00') , "max": Decimal('300.00')}, - # "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - # }, - # "name": "cost_isnt_in_range", - # "criticality": "error", - - # }, + { "check": { "function": "is_not_null", - "arguments": {"column": "next_scheduled_date"}, - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, - + "arguments": {"column": "next_scheduled_date"}}, + + "filter": None, "name": "next_scheduled_date_is_null", "criticality": "error", }, - # { - # "check": { - # "function": "min_max", - # "arguments": {"column": "next_scheduled_date","min": datetime.date(2025, 7, 15), "max": datetime.date(2025, 12, 1)}, - # "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - # }, - # "name": "cost_isnt_in_range", - # "criticality": "error", - - # }, + { "check": { "function": "is_not_null", - "arguments": {"column": "safety_check_passed"}, - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + "arguments": {"column": "safety_check_passed"}}, + "filter": None, "name": "safety_check_passed_is_null", "criticality": "error", @@ -357,8 +354,8 @@ def test_generate_dq_rules_dataframe_filter(ws): { "check": { "function": "is_not_null_and_not_empty", - "arguments": {"column": "vendor_id", "trim_strings": True}, - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, + "arguments": {"column": "vendor_id", "trim_strings": True}}, + "filter": None, "name": "vendor_id_is_null_or_empty", "criticality": "error" } diff --git a/tests/integration/test_save_and_load_checks_from_table.py b/tests/integration/test_save_and_load_checks_from_table.py index 383509e06..1b5571388 100644 --- a/tests/integration/test_save_and_load_checks_from_table.py +++ b/tests/integration/test_save_and_load_checks_from_table.py @@ -51,7 +51,13 @@ { "name": "column_not_less_than", "criticality": "warn", +<<<<<<< HEAD "check": {"function": "is_not_less_than", "arguments": {"column": "col_2", "limit": 1}}, +======= + "check": { + "function": "is_not_less_than", + "arguments": {"column": "col_2", "limit": 1}}, +>>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "filter": "Col_3 >1", "user_metadata": {"check_type": "standardization", "check_owner": "someone_else@email.com"}, }, @@ -63,6 +69,7 @@ ] + def test_load_checks_when_checks_table_does_not_exist(ws, make_schema, make_random, spark): catalog_name = "main" schema_name = make_schema(catalog_name=catalog_name).name diff --git a/tests/integration/test_save_checks_to_workspace_file.py b/tests/integration/test_save_checks_to_workspace_file.py index 337f59ac0..68bd5fa51 100644 --- a/tests/integration/test_save_checks_to_workspace_file.py +++ b/tests/integration/test_save_checks_to_workspace_file.py @@ -39,11 +39,22 @@ { "check": { "function": "is_not_null", - "arguments": {"column": "next_scheduled_date"}}, + "arguments": {"column": "next_scheduled_date"}, + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, "name": "next_scheduled_date_is_null", "criticality": "Error", - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + + }, + { + "check": { + "function": "is_not_null", + "arguments": {"column": "cost"}, + "filter": None}, + + "name": "cost_is_null", + "criticality": "error", + } >>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) From 5f1b18567ba28756ccb16ab3a82e80a6681df920 Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Thu, 25 Sep 2025 12:39:29 -0300 Subject: [PATCH 22/34] resolving merge conflict markersin the code and deleting local files --- .gitignore | 1 + databricks.yml | 19 -- src/databricks/labs/dqx/profiler/generator.py | 149 ++------------- src/databricks/labs/dqx/profiler/profiler.py | 125 +------------ .../test_load_checks_from_uc_volume.py | 26 --- tests/integration/test_profiler.py | 56 +----- tests/integration/test_rules_generator.py | 171 +----------------- .../test_save_and_load_checks_from_table.py | 6 - .../test_save_checks_to_workspace_file.py | 25 --- 9 files changed, 37 insertions(+), 541 deletions(-) delete mode 100644 databricks.yml diff --git a/.gitignore b/.gitignore index 770857d8a..1f6e39ffe 100644 --- a/.gitignore +++ b/.gitignore @@ -171,4 +171,5 @@ docs/dqx/docs/reference/api !docs/dqx/docs/reference/api/index.mdx .databricks +databricks.yml diff --git a/databricks.yml b/databricks.yml deleted file mode 100644 index 6bdb7731f..000000000 --- a/databricks.yml +++ /dev/null @@ -1,19 +0,0 @@ -# This is a Databricks asset bundle definition for dqx. -# The Databricks extension requires databricks.yml configuration file. -# See https://docs.databricks.com/dev-tools/bundles/index.html for documentation. - -bundle: - name: dqx - -targets: - dev: - mode: development - default: true - workspace: - host: https://dbc-08fa3e7e-3bcc.cloud.databricks.com - - ## Optionally, there could be 'staging' or 'prod' targets here. - # - # prod: - # workspace: - # host: https://dbc-08fa3e7e-3bcc.cloud.databricks.com diff --git a/src/databricks/labs/dqx/profiler/generator.py b/src/databricks/labs/dqx/profiler/generator.py index cdd6346ef..56d59839d 100644 --- a/src/databricks/labs/dqx/profiler/generator.py +++ b/src/databricks/labs/dqx/profiler/generator.py @@ -25,11 +25,6 @@ def generate_dq_rules(self, profiles: list[DQProfile] | None = None, level: str if profiles is None: profiles = [] dq_rules = [] -<<<<<<< HEAD - -======= - ->>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) for profile in profiles: rule_name = profile.name column = profile.column @@ -38,11 +33,7 @@ def generate_dq_rules(self, profiles: list[DQProfile] | None = None, level: str if rule_name not in self._checks_mapping: logger.info(f"No rule '{rule_name}' for column '{column}'. skipping...") continue -<<<<<<< HEAD expr = self._checks_mapping[rule_name](column, filter, level, **params) -======= - expr = self._checks_mapping[rule_name](column, filter,level,**params) ->>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) if expr: dq_rules.append(expr) @@ -52,11 +43,7 @@ def generate_dq_rules(self, profiles: list[DQProfile] | None = None, level: str return dq_rules @staticmethod -<<<<<<< HEAD def dq_generate_is_in(column: str, filter: str | None, level: str = "error", **params: dict): -======= - def dq_generate_is_in(column: str, filter:str|None,level: str = "error",**params: dict ): ->>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) """ Generates a data quality rule to check if a column's value is in a specified list. @@ -69,31 +56,16 @@ def dq_generate_is_in(column: str, filter:str|None,level: str = "error",**params A dictionary representing the data quality rule. """ return { -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD + "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"]}}, "filter": filter, -======= - "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"],"filter": params.get("filter", None)}}, ->>>>>>> 07d81a3 (Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.) -======= - "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"]},"filter":filter}, ->>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) -======= - "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"]}}, - "filter":filter, ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": f"{column}_other_value", - "criticality": level + "criticality": level, + } @staticmethod -<<<<<<< HEAD def dq_generate_min_max(column: str, filter: str | None, level: str = "error", **params: dict): -======= - def dq_generate_min_max(column: str, filter:str|None,level: str = "error", **params: dict): ->>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) """ Generates a data quality rule to check if a column's value is within a specified range. @@ -119,97 +91,46 @@ def dq_generate_min_max(column: str, filter:str|None,level: str = "error", **par "arguments": { "column": column, "min_limit": val_maybe_to_str(min_limit, include_sql_quotes=False), -<<<<<<< HEAD -<<<<<<< HEAD "max_limit": val_maybe_to_str(max_limit, include_sql_quotes=False), - "filter": params.get("filter", None) + }, }, "filter": filter, "name": f"{column}_isnt_in_range", "criticality": level -======= - "max_limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}, - "filter": filter}, -======= - "max_limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}}, - "filter": filter, ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) - "name": f"{column}_isnt_in_range", - "criticality": level ->>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) + } if max_limit is not None: return { "check": { "function": "is_not_greater_than", -<<<<<<< HEAD "arguments": {"column": column, "limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}, -======= - "arguments": { - "column": column, -<<<<<<< HEAD -<<<<<<< HEAD - "limit": val_maybe_to_str(max_limit, include_sql_quotes=False), - "filter": params.get("filter", None) - }, ->>>>>>> 07d81a3 (Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.) + }, "filter": filter, "name": f"{column}_not_greater_than", "criticality": level -======= - "limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}, - "filter": filter}, -======= - "limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}}, - "filter": filter, ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) - "name": f"{column}_not_greater_than", - "criticality": level ->>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) + } if min_limit is not None: return { "check": { "function": "is_not_less_than", -<<<<<<< HEAD "arguments": {"column": column, "limit": val_maybe_to_str(min_limit, include_sql_quotes=False)}, -======= - "arguments": { - "column": column, -<<<<<<< HEAD -<<<<<<< HEAD - "limit": val_maybe_to_str(min_limit, include_sql_quotes=False), - "filter": params.get("filter", None) - }, ->>>>>>> 07d81a3 (Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.) }, - "filter": filter, - "name": f"{column}_not_less_than", - "criticality": level -======= - "limit": val_maybe_to_str(min_limit, include_sql_quotes=False)}, - "filter": filter}, -======= - "limit": val_maybe_to_str(min_limit, include_sql_quotes=False)}}, - "filter": filter, ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) + "filter": filter, + "name": f"{column}_not_less_than", "criticality": level ->>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) + } return None @staticmethod -<<<<<<< HEAD def dq_generate_is_not_null(column: str, filter: str | None, level: str = "error", **params: dict): -======= - def dq_generate_is_not_null(column: str, filter:str|None,level: str = "error", **params: dict): ->>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) """ Generates a data quality rule to check if a column's value is not null. @@ -222,41 +143,18 @@ def dq_generate_is_not_null(column: str, filter:str|None,level: str = "error", * A dictionary representing the data quality rule. """ params = params or {} -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD -======= - filter = params.get("dataset_filter_expression", None) ->>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) - return { - "check": {"function": "is_not_null", "arguments": {"column": column}}, - "filter": filter, -======= - filter = params.get("filter", None) - return { - "check": {"function": "is_not_null", "arguments": {"column": column, "filter": params.get("filter", None)}}, ->>>>>>> 07d81a3 (Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.) -======= return { -<<<<<<< HEAD - "check": {"function": "is_not_null", "arguments": {"column": column},"filter":filter}, ->>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) -======= + "check": {"function": "is_not_null", "arguments": {"column": column}}, "filter":filter, ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": f"{column}_is_null", "criticality": level } @staticmethod -<<<<<<< HEAD def dq_generate_is_not_null_or_empty(column: str, filter: str | None, level: str = "error", **params: dict): -======= - def dq_generate_is_not_null_or_empty(column: str, filter:str|None,level: str = "error", **params: dict): ->>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) """ Generates a data quality rule to check if a column's value is not null or empty. @@ -268,35 +166,14 @@ def dq_generate_is_not_null_or_empty(column: str, filter:str|None,level: str = " Returns: A dictionary representing the data quality rule. """ -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD - -======= - filter = params.get("dataset_filter_expression", None) ->>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) -======= - filter = params.get("filter", None) ->>>>>>> 07d81a3 (Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.) - return { - "check": { - "function": "is_not_null_and_not_empty", - "arguments": {"column": column, "trim_strings": params.get("trim_strings", True), "filter": params.get("filter", None)}, - }, - "filter": filter, -======= return { "check": { "function": "is_not_null_and_not_empty", -<<<<<<< HEAD - "arguments": {"column": column, "trim_strings": params.get("trim_strings", True)}, - "filter": filter}, ->>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) -======= + "arguments": {"column": column, "trim_strings": params.get("trim_strings", True)}}, "filter": filter, ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) + "name": f"{column}_is_null_or_empty", "criticality": level } diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 035234f94..db729a56a 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -32,12 +32,6 @@ class DQProfile: description: str | None = None parameters: dict[str, Any] | None = None filter: str | None = None -<<<<<<< HEAD - -======= - - ->>>>>>> 07d81a3 (Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.) class DQProfiler(DQEngineBase): """Data Quality Profiler class to profile input data.""" @@ -59,30 +53,10 @@ def __init__(self, workspace_client: WorkspaceClient, spark: SparkSession | None "sample_fraction": 0.3, # fraction of data to sample (30%) "sample_seed": None, # seed for sampling "limit": 1000, # limit the number of samples -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD - "filter": None, # filter to apply to the dataset -======= - "filter":{}, # filter to apply before profiling - ->>>>>>> ce34e87 (First commit: Installing DQX from feature branch) -======= - "dataset_filter": None, # filter to apply to the dataset ->>>>>>> 33f55cc (Implement methods in the DQProfiler class to filter dataframes before submited to profile.) -======= - # "dataset_filter": None, # filter to apply to the dataset ->>>>>>> a3a71dc (Implement print statement to test dataset_filter feature) -======= - "dataset_filter_expression": None, # filter to apply to the dataset ->>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) -======= "filter": None, # filter to apply to the dataset ->>>>>>> 07d81a3 (Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.) } + @staticmethod def get_columns_or_fields(columns: list[T.StructField]) -> list[T.StructField]: """ @@ -129,19 +103,6 @@ def profile( options = {} options = {**self.default_profile_options, **options} # merge default options with user-provided options -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD - -======= - print(options) ->>>>>>> a3a71dc (Implement print statement to test dataset_filter feature) -======= - ->>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) -======= - ->>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) df = self._sample(df, options) dq_rules: list[DQProfile] = [] @@ -151,18 +112,10 @@ def profile( return summary_stats, dq_rules self._profile(df, df_columns, dq_rules, options, summary_stats, total_count) -<<<<<<< HEAD filter = options.get("filter", None) if filter: for rule in dq_rules: rule.filter = filter -======= - filter=options.get("filter",None) - if filter: - for rule in dq_rules: - rule.filter=filter - ->>>>>>> 07d81a3 (Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.) return summary_stats, dq_rules @@ -368,37 +321,13 @@ def _sample(df: DataFrame, opts: dict[str, Any]) -> DataFrame: sample_fraction = opts.get("sample_fraction", None) sample_seed = opts.get("sample_seed", None) limit = opts.get("limit", None) -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD + filter = opts.get("filter", None) - if filter: - df = df.filter(filter) -======= - filter = opts.get("dataset_filter", None) -<<<<<<< HEAD -<<<<<<< HEAD - if filter: - df = DQProfiler._filter_dataframe(df, filter) ->>>>>>> 33f55cc (Implement methods in the DQProfiler class to filter dataframes before submited to profile.) -======= - - df_filter = DQProfiler._filter_dataframe(df, filter) ->>>>>>> 48cd64d (Refactor _sample method to utilize filtered dataframe) -======= - df = DQProfiler._filter_dataframe(df, filter) ->>>>>>> c77a6eb (Refactor _sample method to utilize return the filtered dataframe) -======= - filter = opts.get("dataset_filter_expression", None) -======= - filter = opts.get("filter", None) ->>>>>>> 07d81a3 (Refactor filter implementation in the profiler class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.) if filter: df = df.filter(filter) ->>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) if sample_fraction: df = df.sample(withReplacement=False, fraction=sample_fraction, seed=sample_seed) if limit: @@ -479,26 +408,13 @@ def _calculate_metrics( cnt = dst2.count() if 0 < cnt < total_count * opts["distinct_ratio"] and cnt < opts["max_in_count"]: dq_rules.append( -<<<<<<< HEAD -<<<<<<< HEAD - DQProfile( - name="is_in", - column=field_name, - parameters={"in": [row[0] for row in dst2.collect()]}, - ) -======= - DQProfile(name="is_in", - column=field_name, - parameters={"in": [row[0] for row in dst2.collect()], "dataset_filter_expression": filter}) ->>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) -======= DQProfile( name="is_in", column=field_name, - parameters={"in": [row[0] for row in dst2.collect()]}, - ) ->>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) + parameters={"in": [row[0] for row in dst2.collect()]}) ) + + if ( typ == T.StringType() and not any( # does not make sense to add is_not_null_or_empty if is_not_null already exists @@ -509,25 +425,11 @@ def _calculate_metrics( cnt = dst2.count() if cnt <= (metrics["count"] * opts.get("max_empty_ratio", 0)): dq_rules.append( -<<<<<<< HEAD -<<<<<<< HEAD - DQProfile( - name="is_not_null_or_empty", - column=field_name, - parameters={"trim_strings": trim_strings}, - ) -======= - DQProfile(name="is_not_null_or_empty", - column=field_name, - parameters={"trim_strings": trim_strings, "dataset_filter_expression": filter}) ->>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) -======= DQProfile( name="is_not_null_or_empty", column=field_name, - parameters={"trim_strings": trim_strings}, - ) ->>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) + parameters={"trim_strings": trim_strings},) + ) if metrics["count_non_null"] > 0 and self._type_supports_min_max(typ): rule = self._extract_min_max(dst, field_name, typ, metrics, opts) @@ -681,21 +583,12 @@ def _extract_min_max( logger.info(f"Can't get min/max for field {col_name}") if descr and min_limit and max_limit: return DQProfile( -<<<<<<< HEAD -<<<<<<< HEAD - name="min_max", - column=col_name, - parameters={"min": min_limit, "max": max_limit}, - description=descr, -======= - name="min_max", column=col_name, parameters={"min": min_limit, "max": max_limit,"dataset_filter_expression": filter}, description=descr ->>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) -======= + name="min_max", column=col_name, parameters={"min": min_limit, "max": max_limit}, description=descr, ->>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) + ) return None diff --git a/tests/integration/test_load_checks_from_uc_volume.py b/tests/integration/test_load_checks_from_uc_volume.py index acd164db2..6659d78f4 100644 --- a/tests/integration/test_load_checks_from_uc_volume.py +++ b/tests/integration/test_load_checks_from_uc_volume.py @@ -9,59 +9,33 @@ TEST_CHECKS = [ { "criticality": "error", -<<<<<<< HEAD "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}, "filter": None}, -======= - "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}, "filter":None}, ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) }, { "check": { "function": "is_not_null", "arguments": {"column": "next_scheduled_date"}, -<<<<<<< HEAD "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", }, "name": "next_scheduled_date_is_null", "criticality": "Error", }, -======= - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, - - "name": "next_scheduled_date_is_null", - "criticality": "Error", - - } ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) ] EXPECTED_CHECKS = [ { "criticality": "error", -<<<<<<< HEAD "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}, "filter": None}, -======= - "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {},"filter":None}, ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) }, { "check": { "function": "is_not_null", "arguments": {"column": "next_scheduled_date"}, -<<<<<<< HEAD "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", }, "name": "next_scheduled_date_is_null", "criticality": "Error", }, -======= - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, - - "name": "next_scheduled_date_is_null", - "criticality": "Error", - - } ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) ] diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index 87c26311e..7bf989bbe 100644 --- a/tests/integration/test_profiler.py +++ b/tests/integration/test_profiler.py @@ -69,15 +69,11 @@ def test_profiler(spark, ws): expected_rules = [ DQProfile(name="is_not_null", column="t1", description=None, parameters=None, filter=None), DQProfile( -<<<<<<< HEAD name="min_max", column="t1", description="Real min/max values were used", parameters={"min": 1, "max": 3}, filter=None, -======= - name="min_max", column="t1", description="Real min/max values were used", parameters={"min": 1, "max": 3}, filter=None ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) ), DQProfile(name='is_not_null', column='d1', description=None, parameters=None, filter=None), DQProfile( @@ -85,17 +81,11 @@ def test_profiler(spark, ws): column='d1', description='Real min/max values were used', parameters={'max': Decimal('333323.00'), 'min': Decimal('1.23')}, -<<<<<<< HEAD filter=None, ), DQProfile( name='is_not_null_or_empty', column='t2', description=None, parameters={'trim_strings': True}, filter=None ), -======= - filter=None - ), - DQProfile(name='is_not_null_or_empty', column='t2', description=None, parameters={'trim_strings': True},filter=None), ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) DQProfile(name="is_not_null", column="s1.ns1", description=None, parameters=None), DQProfile( name="min_max", @@ -105,31 +95,18 @@ def test_profiler(spark, ws): "min": datetime(2023, 1, 6, 0, 0, tzinfo=timezone.utc), "max": datetime(2023, 1, 9, 0, 0, tzinfo=timezone.utc), }, -<<<<<<< HEAD filter=None, ), DQProfile(name="is_not_null", column="s1.s2.ns2", description=None, parameters=None, filter=None), DQProfile(name="is_not_null", column="s1.s2.ns3", description=None, parameters=None, filter=None), -======= - filter=None - ), - DQProfile(name="is_not_null", column="s1.s2.ns2", description=None, parameters=None,filter=None), - DQProfile(name="is_not_null", column="s1.s2.ns3", description=None, parameters=None,filter=None), ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) DQProfile( name="min_max", column="s1.s2.ns3", description="Real min/max values were used", parameters={"min": date(2023, 1, 6), "max": date(2023, 1, 8)}, -<<<<<<< HEAD filter=None, ), DQProfile(name="is_not_null", column="b1", description=None, parameters=None, filter=None), -======= - filter=None - ), - DQProfile(name="is_not_null", column="b1", description=None, parameters=None,filter=None), ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) ] print(stats) assert len(stats.keys()) > 0 @@ -1287,9 +1264,6 @@ def test_profile_tables_no_tables_or_patterns(ws): with pytest.raises(ValueError, match="Either 'tables' or 'patterns' must be provided"): profiler.profile_tables() -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD def test_profile_with_dataset_filter(spark, ws): schema = T.StructType( @@ -1340,7 +1314,6 @@ def test_profile_with_dataset_filter(spark, ws): "preventive", date(2025, 4, 29), Decimal("300.50"), -<<<<<<< HEAD date(2025, 7, 15), True, ), @@ -1450,22 +1423,12 @@ def test_profile_with_dataset_filter(spark, ws): assert filtered_df.count() == 3 assert len(stats.keys()) > 0 assert rules == expected_rules -======= -def test_profile_sample_dataset_filter(spark, ws): - schema = T.StructType( - [ - T.StructField("maintenance_id", T.StringType(), False), -======= -def test_profile_with_dataset_filter(spark, ws): - schema = T.StructType( [ - ->>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) -======= + def test_profile_with_dataset_filter(spark, ws): schema = T.StructType( [ ->>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) + T.StructField("machine_id", T.StringType(), False), T.StructField("maintenance_type", T.StringType(), True), T.StructField("maintenance_date", T.DateType(), True), @@ -1512,8 +1475,6 @@ def test_profile_with_dataset_filter(spark, ws): "preventive", date(2025, 4, 29), Decimal("300.00"), -======= ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) date(2025, 7, 15), True, ), @@ -1554,12 +1515,8 @@ def test_profile_with_dataset_filter(spark, ws): profiler = DQProfiler(ws) filtered_df = profiler._sample(input_df, opts=custom_options) -<<<<<<< HEAD - assert filtered_df.count() == 1 - ->>>>>>> a3a71dc (Implement print statement to test dataset_filter feature) -======= + stats, rules = profiler.profile(input_df, options=custom_options) expected_rules = [ @@ -1602,7 +1559,7 @@ def test_profile_with_dataset_filter(spark, ws): column="cost", parameters={ "min": Decimal('100.00'), - "max": Decimal('300.50'), + "max": Decimal('300.00'), }, filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", description="Real min/max values were used", @@ -1633,8 +1590,3 @@ def test_profile_with_dataset_filter(spark, ws): assert filtered_df.count() == 3 assert len(stats.keys()) > 0 assert rules == expected_rules -<<<<<<< HEAD - ->>>>>>> 8c91719 (Modifying DQProfiler class to generate DQProfile class instances with dataset_filter_expression as one of his parameter) -======= ->>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) diff --git a/tests/integration/test_rules_generator.py b/tests/integration/test_rules_generator.py index 3012d9268..28183d2b5 100644 --- a/tests/integration/test_rules_generator.py +++ b/tests/integration/test_rules_generator.py @@ -2,7 +2,6 @@ from decimal import Decimal # from datetime import date, datetime, timezone -# from datetime import date, datetime, timezone from databricks.labs.dqx.profiler.generator import DQGenerator from databricks.labs.dqx.profiler.profiler import DQProfile @@ -46,35 +45,17 @@ def test_generate_dq_rules(ws): expectations = generator.generate_dq_rules(test_rules) expected = [ { -<<<<<<< HEAD -<<<<<<< HEAD + "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"}}, "filter": None, -======= - "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"},"filter":None}, ->>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) -======= - "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"}}, - "filter":None, ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": "vendor_id_is_null", "criticality": "error", }, { -<<<<<<< HEAD + "check": {"function": "is_in_list", "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}}, "filter": None, -======= - "check": { - "function": "is_in_list", -<<<<<<< HEAD - "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]},"filter":None, - }, ->>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) -======= - "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}}, - "filter":None, ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) + "name": "vendor_id_other_value", "criticality": "error", }, @@ -83,26 +64,17 @@ def test_generate_dq_rules(ws): "function": "is_not_null_and_not_empty", "arguments": {"column": "vendor_id", "trim_strings": True} }, -<<<<<<< HEAD "filter": None, -======= - "filter":None, ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": "vendor_id_is_null_or_empty", "criticality": "error", }, { "check": { "function": "is_in_range", -<<<<<<< HEAD "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}, - "filter":None + }, "filter": None, -======= - "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}}, - "filter":None, ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": "rate_code_id_isnt_in_range", "criticality": "error", }, @@ -115,63 +87,36 @@ def test_generate_dq_rules_warn(ws): expectations = generator.generate_dq_rules(test_rules, level="warn") expected = [ { -<<<<<<< HEAD -<<<<<<< HEAD + + "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"}}, "filter": None, -======= - "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"},"filter":None}, ->>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) -======= - "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"}}, - "filter":None, ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) + "name": "vendor_id_is_null", "criticality": "warn", }, { -<<<<<<< HEAD + "check": {"function": "is_in_list", "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}}, "filter": None, -======= - "check": { - "function": "is_in_list", -<<<<<<< HEAD - "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]},"filter":None, - }, ->>>>>>> ef060b8 (Refactor filter implementation in the DQGenerator class. Changing from inserting key, value pair in parameter attribute to adding new class attribute called filter.Filter is now another key for the check dict.) -======= - "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}}, - "filter":None, ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": "vendor_id_other_value", "criticality": "warn", }, { "check": { "function": "is_not_null_and_not_empty", -<<<<<<< HEAD - "arguments": {"column": "vendor_id", "trim_strings": True},"filter":None, + "arguments": {"column": "vendor_id", "trim_strings": True} }, "filter": None, -======= - "arguments": {"column": "vendor_id", "trim_strings": True}}, - "filter":None, ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": "vendor_id_is_null_or_empty", "criticality": "warn", }, { "check": { "function": "is_in_range", -<<<<<<< HEAD - "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265},"filter":None, + "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265} }, "filter": None, -======= - "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}}, - "filter":None, ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "name": "rate_code_id_isnt_in_range", "criticality": "warn", }, @@ -190,7 +135,6 @@ def test_generate_dq_no_rules(ws): expectations = generator.generate_dq_rules(None, level="warn") assert not expectations -<<<<<<< HEAD def test_generate_dq_rules_dataframe_filter(ws): generator = DQGenerator(ws) @@ -266,100 +210,5 @@ def test_generate_dq_rules_dataframe_filter(ws): "name": "vendor_id_is_null_or_empty", "criticality": "error", }, -======= -def test_generate_dq_rules_dataframe_filter(ws): - generator = DQGenerator(ws) - test_rules=[ - DQProfile(name="is_not_null", - column="machine_id", - description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"), - - DQProfile(name="is_in", - column="vendor_id", - parameters={"in": ["1", "4", "2"]}, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"), - - DQProfile(name="is_not_null", - column="cost", - description=None, - ), - - DQProfile(name="is_not_null", - column="next_scheduled_date", - description=None, - ), - - DQProfile(name="is_not_null", - column="safety_check_passed", - description=None, - ), - - DQProfile(name="is_not_null_or_empty", - column="vendor_id", - parameters={"trim_strings": True}), - - ] - expectations = generator.generate_dq_rules(test_rules) - - expected = [ - { - "check": {"function": "is_not_null", - "arguments": {"column": "machine_id"}}, - "filter":"machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", - "name": "machine_id_is_null", - "criticality": "error" - - }, - { - 'check': {"function": "is_in_list", - "arguments": {"allowed": ['1','4','2'],"column": "vendor_id"}}, - - "filter":"machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", - "criticality": "error", - "name": "vendor_id_other_value" - }, - { - "check": { - "function": "is_not_null", - "arguments": {"column": "cost"}}, - - "filter": None, - "name": "cost_is_null", - "criticality": "error", - - }, - - { - "check": { - "function": "is_not_null", - "arguments": {"column": "next_scheduled_date"}}, - - "filter": None, - "name": "next_scheduled_date_is_null", - "criticality": "error", - - }, - - { - "check": { - "function": "is_not_null", - "arguments": {"column": "safety_check_passed"}}, - "filter": None, - - "name": "safety_check_passed_is_null", - "criticality": "error", - - }, - { - "check": { - "function": "is_not_null_and_not_empty", - "arguments": {"column": "vendor_id", "trim_strings": True}}, - "filter": None, - "name": "vendor_id_is_null_or_empty", - "criticality": "error" - } - ->>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) ] assert expectations == expected diff --git a/tests/integration/test_save_and_load_checks_from_table.py b/tests/integration/test_save_and_load_checks_from_table.py index 1b5571388..01f4475b1 100644 --- a/tests/integration/test_save_and_load_checks_from_table.py +++ b/tests/integration/test_save_and_load_checks_from_table.py @@ -51,13 +51,7 @@ { "name": "column_not_less_than", "criticality": "warn", -<<<<<<< HEAD "check": {"function": "is_not_less_than", "arguments": {"column": "col_2", "limit": 1}}, -======= - "check": { - "function": "is_not_less_than", - "arguments": {"column": "col_2", "limit": 1}}, ->>>>>>> 6326af5 (Creating and running tests for profiler and generator class. Creating and running test for save and load checks FileChecksStorageConfig, WorkspaceFileChecksStorageConfig, InstallationChecksStorageConfig, TableChecksStorageConfig, VolumeFileChecksStorageConfig.) "filter": "Col_3 >1", "user_metadata": {"check_type": "standardization", "check_owner": "someone_else@email.com"}, }, diff --git a/tests/integration/test_save_checks_to_workspace_file.py b/tests/integration/test_save_checks_to_workspace_file.py index 68bd5fa51..eda7e3fb5 100644 --- a/tests/integration/test_save_checks_to_workspace_file.py +++ b/tests/integration/test_save_checks_to_workspace_file.py @@ -18,7 +18,6 @@ { "criticality": "error", "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}}, -<<<<<<< HEAD }, { "check": { @@ -34,30 +33,6 @@ "name": "cost_is_null", "criticality": "error", }, -======= - }, - { - "check": { - "function": "is_not_null", - "arguments": {"column": "next_scheduled_date"}, - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'"}, - - "name": "next_scheduled_date_is_null", - "criticality": "Error", - - }, - { - "check": { - "function": "is_not_null", - "arguments": {"column": "cost"}, - "filter": None}, - - "name": "cost_is_null", - "criticality": "error", - - } - ->>>>>>> 0377f1e (Modifying DQGenerator class to show filter in the generate_dq_rules method. Added test in the test_rules_generator.py file.) ] def test_save_checks_in_workspace_file_as_yaml(ws, spark, installation_ctx): From ebe714b4e810be0a656ed40dabbeab0efcadad1a Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Fri, 26 Sep 2025 19:50:45 -0300 Subject: [PATCH 23/34] Refactor DQGenerator class to only return filter key in the rule dict if filter is not None --- src/databricks/labs/dqx/profiler/generator.py | 35 ++++++------ tests/integration/test_rules_generator.py | 55 +++++++++++++++++++ 2 files changed, 71 insertions(+), 19 deletions(-) diff --git a/src/databricks/labs/dqx/profiler/generator.py b/src/databricks/labs/dqx/profiler/generator.py index 56d59839d..691b7762c 100644 --- a/src/databricks/labs/dqx/profiler/generator.py +++ b/src/databricks/labs/dqx/profiler/generator.py @@ -29,21 +29,24 @@ def generate_dq_rules(self, profiles: list[DQProfile] | None = None, level: str rule_name = profile.name column = profile.column params = profile.parameters or {} - filter = profile.filter + filter = profile.filter if rule_name not in self._checks_mapping: logger.info(f"No rule '{rule_name}' for column '{column}'. skipping...") continue - expr = self._checks_mapping[rule_name](column, filter, level, **params) + expr = self._checks_mapping[rule_name](column, level, **params) + if expr: - dq_rules.append(expr) - + if filter: + expr["filter"] = filter + dq_rules.append(expr) + status = DQEngine.validate_checks(dq_rules) assert not status.has_errors return dq_rules @staticmethod - def dq_generate_is_in(column: str, filter: str | None, level: str = "error", **params: dict): + def dq_generate_is_in(column: str, level: str = "error", **params: dict): """ Generates a data quality rule to check if a column's value is in a specified list. @@ -57,15 +60,14 @@ def dq_generate_is_in(column: str, filter: str | None, level: str = "error", **p """ return { - "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"]}}, - "filter": filter, + "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"]}}, "name": f"{column}_other_value", "criticality": level, } @staticmethod - def dq_generate_min_max(column: str, filter: str | None, level: str = "error", **params: dict): + def dq_generate_min_max(column: str, level: str = "error", **params: dict): """ Generates a data quality rule to check if a column's value is within a specified range. @@ -94,8 +96,7 @@ def dq_generate_min_max(column: str, filter: str | None, level: str = "error", * "max_limit": val_maybe_to_str(max_limit, include_sql_quotes=False), }, - }, - "filter": filter, + }, "name": f"{column}_isnt_in_range", "criticality": level @@ -107,8 +108,7 @@ def dq_generate_min_max(column: str, filter: str | None, level: str = "error", * "function": "is_not_greater_than", "arguments": {"column": column, "limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}, - }, - "filter": filter, + }, "name": f"{column}_not_greater_than", "criticality": level @@ -119,8 +119,7 @@ def dq_generate_min_max(column: str, filter: str | None, level: str = "error", * "check": { "function": "is_not_less_than", "arguments": {"column": column, "limit": val_maybe_to_str(min_limit, include_sql_quotes=False)}, - }, - "filter": filter, + }, "name": f"{column}_not_less_than", "criticality": level @@ -130,7 +129,7 @@ def dq_generate_min_max(column: str, filter: str | None, level: str = "error", * return None @staticmethod - def dq_generate_is_not_null(column: str, filter: str | None, level: str = "error", **params: dict): + def dq_generate_is_not_null(column: str, level: str = "error", **params: dict): """ Generates a data quality rule to check if a column's value is not null. @@ -148,13 +147,12 @@ def dq_generate_is_not_null(column: str, filter: str | None, level: str = "error return { "check": {"function": "is_not_null", "arguments": {"column": column}}, - "filter":filter, "name": f"{column}_is_null", "criticality": level } @staticmethod - def dq_generate_is_not_null_or_empty(column: str, filter: str | None, level: str = "error", **params: dict): + def dq_generate_is_not_null_or_empty(column: str, level: str = "error", **params: dict): """ Generates a data quality rule to check if a column's value is not null or empty. @@ -172,8 +170,7 @@ def dq_generate_is_not_null_or_empty(column: str, filter: str | None, level: str "function": "is_not_null_and_not_empty", "arguments": {"column": column, "trim_strings": params.get("trim_strings", True)}}, - "filter": filter, - + "name": f"{column}_is_null_or_empty", "criticality": level } diff --git a/tests/integration/test_rules_generator.py b/tests/integration/test_rules_generator.py index 28183d2b5..a2573e75a 100644 --- a/tests/integration/test_rules_generator.py +++ b/tests/integration/test_rules_generator.py @@ -212,3 +212,58 @@ def test_generate_dq_rules_dataframe_filter(ws): }, ] assert expectations == expected + + +def test_generate_dq_rules_dataframe_filter_none(ws): + generator = DQGenerator(ws) + test_rules = [ + DQProfile( + name="is_not_null", + column="machine_id", + description=None, + filter=None, + ), + DQProfile( + name="is_in", + column="vendor_id", + parameters={"in": ["1", "4", "2"]}, + filter=None, + ), + + DQProfile( + name="is_not_null", + column="next_scheduled_date", + description=None, + filter=None, + ), + + DQProfile(name="is_not_null_or_empty", column="vendor_id", parameters={"trim_strings": True}, filter=None), + ] + expectations = generator.generate_dq_rules(test_rules) + + expected = [ + { + "check": {"function": "is_not_null", "arguments": {"column": "machine_id"}}, + "name": "machine_id_is_null", + "criticality": "error", + }, + { + "check": {"function": "is_in_list", "arguments": {"allowed": ['1', '4', '2'], "column": "vendor_id"}}, + "criticality": "error", + "name": "vendor_id_other_value", + }, + { + "check": {"function": "is_not_null", "arguments": {"column": "next_scheduled_date"}}, + "name": "next_scheduled_date_is_null", + "criticality": "error", + }, + { + "check": { + "function": "is_not_null_and_not_empty", + "arguments": {"column": "vendor_id", "trim_strings": True}, + }, + "name": "vendor_id_is_null_or_empty", + "criticality": "error", + } + ] + assert expectations == expected From 3a3a7a4bb5c2b0162f30204c03b88b204381ccf8 Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Fri, 26 Sep 2025 20:01:15 -0300 Subject: [PATCH 24/34] Formatting code with make fmt --- src/databricks/labs/dqx/profiler/generator.py | 46 +++++++------------ src/databricks/labs/dqx/profiler/profiler.py | 26 ++++------- tests/integration/test_profiler.py | 31 +++++-------- tests/integration/test_rules_generator.py | 23 +++------- .../test_save_and_load_checks_from_table.py | 1 - .../test_save_checks_to_workspace_file.py | 1 + 6 files changed, 45 insertions(+), 83 deletions(-) diff --git a/src/databricks/labs/dqx/profiler/generator.py b/src/databricks/labs/dqx/profiler/generator.py index 691b7762c..8df638d28 100644 --- a/src/databricks/labs/dqx/profiler/generator.py +++ b/src/databricks/labs/dqx/profiler/generator.py @@ -29,17 +29,17 @@ def generate_dq_rules(self, profiles: list[DQProfile] | None = None, level: str rule_name = profile.name column = profile.column params = profile.parameters or {} - filter = profile.filter + filter = profile.filter if rule_name not in self._checks_mapping: logger.info(f"No rule '{rule_name}' for column '{column}'. skipping...") continue expr = self._checks_mapping[rule_name](column, level, **params) - + if expr: if filter: expr["filter"] = filter - dq_rules.append(expr) - + dq_rules.append(expr) + status = DQEngine.validate_checks(dq_rules) assert not status.has_errors @@ -59,11 +59,9 @@ def dq_generate_is_in(column: str, level: str = "error", **params: dict): A dictionary representing the data quality rule. """ return { - - "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"]}}, + "check": {"function": "is_in_list", "arguments": {"column": column, "allowed": params["in"]}}, "name": f"{column}_other_value", "criticality": level, - } @staticmethod @@ -81,7 +79,6 @@ def dq_generate_min_max(column: str, level: str = "error", **params: dict): """ min_limit = params.get("min") max_limit = params.get("max") - if not isinstance(min_limit, int) or not isinstance(max_limit, int): return None # TODO handle timestamp and dates: https://github.com/databrickslabs/dqx/issues/71 @@ -94,12 +91,10 @@ def dq_generate_min_max(column: str, level: str = "error", **params: dict): "column": column, "min_limit": val_maybe_to_str(min_limit, include_sql_quotes=False), "max_limit": val_maybe_to_str(max_limit, include_sql_quotes=False), - }, - }, + }, "name": f"{column}_isnt_in_range", - "criticality": level - + "criticality": level, } if max_limit is not None: @@ -107,11 +102,9 @@ def dq_generate_min_max(column: str, level: str = "error", **params: dict): "check": { "function": "is_not_greater_than", "arguments": {"column": column, "limit": val_maybe_to_str(max_limit, include_sql_quotes=False)}, - - }, + }, "name": f"{column}_not_greater_than", - "criticality": level - + "criticality": level, } if min_limit is not None: @@ -119,11 +112,9 @@ def dq_generate_min_max(column: str, level: str = "error", **params: dict): "check": { "function": "is_not_less_than", "arguments": {"column": column, "limit": val_maybe_to_str(min_limit, include_sql_quotes=False)}, - }, - - "name": f"{column}_not_less_than", - "criticality": level - + }, + "name": f"{column}_not_less_than", + "criticality": level, } return None @@ -143,12 +134,10 @@ def dq_generate_is_not_null(column: str, level: str = "error", **params: dict): """ params = params or {} - return { - "check": {"function": "is_not_null", "arguments": {"column": column}}, "name": f"{column}_is_null", - "criticality": level + "criticality": level, } @staticmethod @@ -164,15 +153,14 @@ def dq_generate_is_not_null_or_empty(column: str, level: str = "error", **params Returns: A dictionary representing the data quality rule. """ - + return { "check": { "function": "is_not_null_and_not_empty", - - "arguments": {"column": column, "trim_strings": params.get("trim_strings", True)}}, - + "arguments": {"column": column, "trim_strings": params.get("trim_strings", True)}, + }, "name": f"{column}_is_null_or_empty", - "criticality": level + "criticality": level, } _checks_mapping = { diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index db729a56a..07ad9be2c 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -33,6 +33,7 @@ class DQProfile: parameters: dict[str, Any] | None = None filter: str | None = None + class DQProfiler(DQEngineBase): """Data Quality Profiler class to profile input data.""" @@ -56,7 +57,6 @@ def __init__(self, workspace_client: WorkspaceClient, spark: SparkSession | None "filter": None, # filter to apply to the dataset } - @staticmethod def get_columns_or_fields(columns: list[T.StructField]) -> list[T.StructField]: """ @@ -324,8 +324,6 @@ def _sample(df: DataFrame, opts: dict[str, Any]) -> DataFrame: filter = opts.get("filter", None) - - if filter: df = df.filter(filter) if sample_fraction: @@ -378,7 +376,7 @@ def _calculate_metrics( """ max_nulls = opts.get("max_null_ratio", 0) trim_strings = opts.get("trim_strings", True) - + dst = df.select(field_name).dropna() if typ == T.StringType() and trim_strings: col_name = dst.columns[0] @@ -396,25 +394,19 @@ def _calculate_metrics( name="is_not_null", column=field_name, description=f"Column {field_name} has {null_percentage * 100:.1f}% of null values " - f"(allowed {max_nulls * 100:.1f}%)", + f"(allowed {max_nulls * 100:.1f}%)", ) ) else: - dq_rules.append( - DQProfile(name="is_not_null", column=field_name) - ) + dq_rules.append(DQProfile(name="is_not_null", column=field_name)) if self._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( - DQProfile( - name="is_in", - column=field_name, - parameters={"in": [row[0] for row in dst2.collect()]}) + DQProfile(name="is_in", column=field_name, parameters={"in": [row[0] for row in dst2.collect()]}) ) - if ( typ == T.StringType() and not any( # does not make sense to add is_not_null_or_empty if is_not_null already exists @@ -428,8 +420,8 @@ def _calculate_metrics( DQProfile( name="is_not_null_or_empty", column=field_name, - parameters={"trim_strings": trim_strings},) - + parameters={"trim_strings": trim_strings}, + ) ) if metrics["count_non_null"] > 0 and self._type_supports_min_max(typ): rule = self._extract_min_max(dst, field_name, typ, metrics, opts) @@ -547,7 +539,7 @@ def _extract_min_max( if opts is None: opts = {} - + outlier_cols = opts.get("outlier_columns", []) column = dst.columns[0] if opts.get("remove_outliers", True) and ( @@ -583,12 +575,10 @@ def _extract_min_max( logger.info(f"Can't get min/max for field {col_name}") if descr and min_limit and max_limit: return DQProfile( - name="min_max", column=col_name, parameters={"min": min_limit, "max": max_limit}, description=descr, - ) return None diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index 7bf989bbe..0af7497e5 100644 --- a/tests/integration/test_profiler.py +++ b/tests/integration/test_profiler.py @@ -1428,7 +1428,6 @@ def test_profile_with_dataset_filter(spark, ws): def test_profile_with_dataset_filter(spark, ws): schema = T.StructType( [ - T.StructField("machine_id", T.StringType(), False), T.StructField("maintenance_type", T.StringType(), True), T.StructField("maintenance_date", T.DateType(), True), @@ -1508,7 +1507,7 @@ def test_profile_with_dataset_filter(spark, ws): custom_options = { "sample_fraction": None, - "round":False, + "round": False, "limit": None, "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", } @@ -1516,43 +1515,39 @@ def test_profile_with_dataset_filter(spark, ws): profiler = DQProfiler(ws) filtered_df = profiler._sample(input_df, opts=custom_options) - stats, rules = profiler.profile(input_df, options=custom_options) expected_rules = [ DQProfile( name="is_not_null", column="machine_id", - description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - + description=None, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", ), DQProfile( name="is_not_null", column="maintenance_type", description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", ), DQProfile( name="is_not_null", column="maintenance_date", description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", ), DQProfile( name="min_max", column="maintenance_date", description="Real min/max values were used", - parameters={ - "min": date(2025, 4, 29), - "max": date(2025, 7, 30)}, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" - ), + parameters={"min": date(2025, 4, 29), "max": date(2025, 7, 30)}, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + ), DQProfile( name="is_not_null", column="cost", description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", ), DQProfile( name="min_max", @@ -1568,16 +1563,14 @@ def test_profile_with_dataset_filter(spark, ws): name="is_not_null", column="next_scheduled_date", description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", ), DQProfile( name="min_max", column="next_scheduled_date", description="Real min/max values were used", - parameters={ - "min": date(2025, 7, 15), - "max": date(2025, 12, 1)}, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'" + parameters={"min": date(2025, 7, 15), "max": date(2025, 12, 1)}, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", ), DQProfile( name="is_not_null", diff --git a/tests/integration/test_rules_generator.py b/tests/integration/test_rules_generator.py index a2573e75a..3dc9cfa99 100644 --- a/tests/integration/test_rules_generator.py +++ b/tests/integration/test_rules_generator.py @@ -1,5 +1,6 @@ import datetime from decimal import Decimal + # from datetime import date, datetime, timezone @@ -10,7 +11,7 @@ DQProfile( name="is_not_null", column="vendor_id", description="Column vendor_id has 0.3% of null values (allowed 1.0%)" ), - DQProfile(name="is_in", column="vendor_id",parameters={"in": ["1", "4", "2"]}), + DQProfile(name="is_in", column="vendor_id", parameters={"in": ["1", "4", "2"]}), DQProfile(name="is_not_null_or_empty", column="vendor_id", parameters={"trim_strings": True}), DQProfile( name="min_max", @@ -45,24 +46,21 @@ def test_generate_dq_rules(ws): expectations = generator.generate_dq_rules(test_rules) expected = [ { - "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"}}, "filter": None, "name": "vendor_id_is_null", "criticality": "error", }, { - "check": {"function": "is_in_list", "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}}, "filter": None, - "name": "vendor_id_other_value", "criticality": "error", }, { "check": { "function": "is_not_null_and_not_empty", - "arguments": {"column": "vendor_id", "trim_strings": True} + "arguments": {"column": "vendor_id", "trim_strings": True}, }, "filter": None, "name": "vendor_id_is_null_or_empty", @@ -72,7 +70,6 @@ def test_generate_dq_rules(ws): "check": { "function": "is_in_range", "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}, - }, "filter": None, "name": "rate_code_id_isnt_in_range", @@ -87,16 +84,12 @@ def test_generate_dq_rules_warn(ws): expectations = generator.generate_dq_rules(test_rules, level="warn") expected = [ { - - "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"}}, "filter": None, - "name": "vendor_id_is_null", "criticality": "warn", }, { - "check": {"function": "is_in_list", "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}}, "filter": None, "name": "vendor_id_other_value", @@ -105,7 +98,7 @@ def test_generate_dq_rules_warn(ws): { "check": { "function": "is_not_null_and_not_empty", - "arguments": {"column": "vendor_id", "trim_strings": True} + "arguments": {"column": "vendor_id", "trim_strings": True}, }, "filter": None, "name": "vendor_id_is_null_or_empty", @@ -114,7 +107,7 @@ def test_generate_dq_rules_warn(ws): { "check": { "function": "is_in_range", - "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265} + "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}, }, "filter": None, "name": "rate_code_id_isnt_in_range", @@ -229,14 +222,12 @@ def test_generate_dq_rules_dataframe_filter_none(ws): parameters={"in": ["1", "4", "2"]}, filter=None, ), - DQProfile( name="is_not_null", column="next_scheduled_date", description=None, filter=None, ), - DQProfile(name="is_not_null_or_empty", column="vendor_id", parameters={"trim_strings": True}, filter=None), ] expectations = generator.generate_dq_rules(test_rules) @@ -256,7 +247,7 @@ def test_generate_dq_rules_dataframe_filter_none(ws): "check": {"function": "is_not_null", "arguments": {"column": "next_scheduled_date"}}, "name": "next_scheduled_date_is_null", "criticality": "error", - }, + }, { "check": { "function": "is_not_null_and_not_empty", @@ -264,6 +255,6 @@ def test_generate_dq_rules_dataframe_filter_none(ws): }, "name": "vendor_id_is_null_or_empty", "criticality": "error", - } + }, ] assert expectations == expected diff --git a/tests/integration/test_save_and_load_checks_from_table.py b/tests/integration/test_save_and_load_checks_from_table.py index 01f4475b1..383509e06 100644 --- a/tests/integration/test_save_and_load_checks_from_table.py +++ b/tests/integration/test_save_and_load_checks_from_table.py @@ -63,7 +63,6 @@ ] - def test_load_checks_when_checks_table_does_not_exist(ws, make_schema, make_random, spark): catalog_name = "main" schema_name = make_schema(catalog_name=catalog_name).name diff --git a/tests/integration/test_save_checks_to_workspace_file.py b/tests/integration/test_save_checks_to_workspace_file.py index eda7e3fb5..35b5a5f51 100644 --- a/tests/integration/test_save_checks_to_workspace_file.py +++ b/tests/integration/test_save_checks_to_workspace_file.py @@ -35,6 +35,7 @@ }, ] + def test_save_checks_in_workspace_file_as_yaml(ws, spark, installation_ctx): installation_ctx.installation.save(installation_ctx.config) install_dir = installation_ctx.installation.install_folder() From 1ea27587c7aa839d0d064ebe55bfdddd27497044 Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Tue, 30 Sep 2025 11:34:27 -0300 Subject: [PATCH 25/34] Fixing format issues --- src/databricks/labs/dqx/profiler/generator.py | 4 +- src/databricks/labs/dqx/profiler/profiler.py | 29 +++++--- .../test_load_checks_from_uc_volume.py | 73 +++++++++++++------ tests/integration/test_profiler.py | 70 +++++++----------- tests/integration/test_rules_generator.py | 33 +++------ 5 files changed, 112 insertions(+), 97 deletions(-) diff --git a/src/databricks/labs/dqx/profiler/generator.py b/src/databricks/labs/dqx/profiler/generator.py index 8df638d28..ae4e64df4 100644 --- a/src/databricks/labs/dqx/profiler/generator.py +++ b/src/databricks/labs/dqx/profiler/generator.py @@ -29,14 +29,14 @@ def generate_dq_rules(self, profiles: list[DQProfile] | None = None, level: str rule_name = profile.name column = profile.column params = profile.parameters or {} - filter = profile.filter + dataset_filter = profile.filter if rule_name not in self._checks_mapping: logger.info(f"No rule '{rule_name}' for column '{column}'. skipping...") continue expr = self._checks_mapping[rule_name](column, level, **params) if expr: - if filter: + if dataset_filter is not None: expr["filter"] = filter dq_rules.append(expr) diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 04755cc78..9be22b861 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -113,10 +113,6 @@ def profile( return summary_stats, dq_rules self._profile(df, df_columns, dq_rules, options, summary_stats, total_count) - filter = options.get("filter", None) - if filter: - for rule in dq_rules: - rule.filter = filter return summary_stats, dq_rules @@ -325,11 +321,10 @@ def _sample(df: DataFrame, opts: dict[str, Any]) -> DataFrame: sample_fraction = opts.get("sample_fraction", None) sample_seed = opts.get("sample_seed", None) limit = opts.get("limit", None) + filter_dataset = opts.get("filter", None) - filter = opts.get("filter", None) - - if filter: - df = df.filter(filter) + if filter_dataset: + df = df.filter(filter_dataset) if sample_fraction: df = df.sample(withReplacement=False, fraction=sample_fraction, seed=sample_seed) if limit: @@ -399,16 +394,28 @@ def _calculate_metrics( column=field_name, description=f"Column {field_name} has {null_percentage * 100:.1f}% of null values " f"(allowed {max_nulls * 100:.1f}%)", + filter=opts.get("filter", None), ) ) else: - dq_rules.append(DQProfile(name="is_not_null", column=field_name)) + dq_rules.append( + DQProfile( + name="is_not_null", + column=field_name, + filter=opts.get("filter", None), + ) + ) if self._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( - DQProfile(name="is_in", column=field_name, parameters={"in": [row[0] for row in dst2.collect()]}) + DQProfile( + name="is_in", + column=field_name, + parameters={"in": [row[0] for row in dst2.collect()]}, + filter=opts.get("filter", None), + ) ) if ( @@ -425,6 +432,7 @@ def _calculate_metrics( name="is_not_null_or_empty", column=field_name, parameters={"trim_strings": trim_strings}, + filter=opts.get("filter", None), ) ) if metrics["count_non_null"] > 0 and self._type_supports_min_max(typ): @@ -583,6 +591,7 @@ def _extract_min_max( column=col_name, parameters={"min": min_limit, "max": max_limit}, description=descr, + filter=opts.get("filter", None), ) return None diff --git a/tests/integration/test_load_checks_from_uc_volume.py b/tests/integration/test_load_checks_from_uc_volume.py index 00cb6e3bf..cbbcc20bd 100644 --- a/tests/integration/test_load_checks_from_uc_volume.py +++ b/tests/integration/test_load_checks_from_uc_volume.py @@ -9,33 +9,15 @@ TEST_CHECKS = [ { "criticality": "error", - "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}, "filter": None}, - }, - { - "check": { - "function": "is_not_null", - "arguments": {"column": "next_scheduled_date"}, - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", - }, - "name": "next_scheduled_date_is_null", - "criticality": "Error", - }, + "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}}, + } ] EXPECTED_CHECKS = [ { "criticality": "error", - "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}, "filter": None}, - }, - { - "check": { - "function": "is_not_null", - "arguments": {"column": "next_scheduled_date"}, - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", - }, - "name": "next_scheduled_date_is_null", - "criticality": "Error", - }, + "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}}, + } ] @@ -234,3 +216,50 @@ def test_save_checks_in_volume_file_as_json(ws, make_schema, make_volume, instal checks = dq_engine.load_checks(config=VolumeFileChecksStorageConfig(location=checks_path)) assert TEST_CHECKS == checks, "Checks were not saved correctly" + + +TEST_CHECKS_FILTER = [ + { + "criticality": "error", + "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}, "filter": None}, + }, + { + "check": { + "function": "is_not_null", + "arguments": {"column": "next_scheduled_date"}, + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + }, + "name": "next_scheduled_date_is_null", + "criticality": "Error", + }, +] + +EXPECTED_CHECKS_FILTER = [ + { + "criticality": "error", + "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}, "filter": None}, + }, + { + "check": { + "function": "is_not_null", + "arguments": {"column": "next_scheduled_date"}, + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + }, + "name": "next_scheduled_date_is_null", + "criticality": "Error", + }, +] + + +def test_save_and_load_checks_with_filters_from_volume(ws, make_schema, make_volume): + + catalog_name = "main" + schema_name = make_schema(catalog_name=catalog_name).name + volume_name = make_volume(catalog_name=catalog_name, schema_name=schema_name).name + volume = f"/Volumes/{catalog_name}/{schema_name}/{volume_name}/checks.yml" + + engine = DQEngine(ws) + config = VolumeFileChecksStorageConfig(location=volume) + engine.save_checks(checks=TEST_CHECKS_FILTER, config=config) + checks = engine.load_checks(config=config) + assert checks == EXPECTED_CHECKS_FILTER, "Checks were not loaded correctly." diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index 7004a20e6..991be9824 100644 --- a/tests/integration/test_profiler.py +++ b/tests/integration/test_profiler.py @@ -68,25 +68,18 @@ def test_profiler(spark, ws): stats, rules = profiler.profile(inp_df, options={"sample_fraction": None}) expected_rules = [ - DQProfile(name="is_not_null", column="t1", description=None, parameters=None, filter=None), + DQProfile(name="is_not_null", column="t1", description=None, parameters=None), DQProfile( - name="min_max", - column="t1", - description="Real min/max values were used", - parameters={"min": 1, "max": 3}, - filter=None, + name="min_max", column="t1", description="Real min/max values were used", parameters={"min": 1, "max": 3} ), - DQProfile(name='is_not_null', column='d1', description=None, parameters=None, filter=None), + DQProfile(name='is_not_null', column='d1', description=None, parameters=None), DQProfile( name='min_max', column='d1', description='Real min/max values were used', parameters={'max': Decimal('333323.00'), 'min': Decimal('1.23')}, - filter=None, - ), - DQProfile( - name='is_not_null_or_empty', column='t2', description=None, parameters={'trim_strings': True}, filter=None ), + DQProfile(name='is_not_null_or_empty', column='t2', description=None, parameters={'trim_strings': True}), DQProfile(name="is_not_null", column="s1.ns1", description=None, parameters=None), DQProfile( name="min_max", @@ -96,18 +89,16 @@ def test_profiler(spark, ws): "min": datetime(2023, 1, 6, 0, 0, tzinfo=timezone.utc), "max": datetime(2023, 1, 9, 0, 0, tzinfo=timezone.utc), }, - filter=None, ), - DQProfile(name="is_not_null", column="s1.s2.ns2", description=None, parameters=None, filter=None), - DQProfile(name="is_not_null", column="s1.s2.ns3", description=None, parameters=None, filter=None), + DQProfile(name="is_not_null", column="s1.s2.ns2", description=None, parameters=None), + DQProfile(name="is_not_null", column="s1.s2.ns3", description=None, parameters=None), DQProfile( name="min_max", column="s1.s2.ns3", description="Real min/max values were used", parameters={"min": date(2023, 1, 6), "max": date(2023, 1, 8)}, - filter=None, ), - DQProfile(name="is_not_null", column="b1", description=None, parameters=None, filter=None), + DQProfile(name="is_not_null", column="b1", description=None, parameters=None), ] print(stats) assert len(stats.keys()) > 0 @@ -1354,7 +1345,6 @@ def test_profile_with_dataset_filter(spark, ws): } profiler = DQProfiler(ws) - filtered_df = profiler._sample(input_df, opts=custom_options) stats, rules = profiler.profile(input_df, options=custom_options) @@ -1421,12 +1411,11 @@ def test_profile_with_dataset_filter(spark, ws): ), ] - assert filtered_df.count() == 3 assert len(stats.keys()) > 0 assert rules == expected_rules -def test_profile_with_dataset_filter(spark, ws): +def test_profile_with_no_filter(spark, ws): schema = T.StructType( [ T.StructField("machine_id", T.StringType(), False), @@ -1460,7 +1449,7 @@ def test_profile_with_dataset_filter(spark, ws): date(2025, 4, 20), Decimal("-500.00"), date(2024, 4, 20), - None, + False, ), ( "MCH-001", @@ -1474,7 +1463,7 @@ def test_profile_with_dataset_filter(spark, ws): "MCH-002", "preventive", date(2025, 4, 29), - Decimal("300.00"), + Decimal("300.50"), date(2025, 7, 15), True, ), @@ -1506,16 +1495,13 @@ def test_profile_with_dataset_filter(spark, ws): input_df = spark.createDataFrame(maintenance_data, schema=schema) + profiler = DQProfiler(ws) custom_options = { "sample_fraction": None, - "round": False, + "round": True, "limit": None, - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + "filter": None, } - - profiler = DQProfiler(ws) - filtered_df = profiler._sample(input_df, opts=custom_options) - stats, rules = profiler.profile(input_df, options=custom_options) expected_rules = [ @@ -1523,64 +1509,64 @@ def test_profile_with_dataset_filter(spark, ws): name="is_not_null", column="machine_id", description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + filter=None, ), DQProfile( - name="is_not_null", + name="is_not_null_or_empty", column="maintenance_type", description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + parameters={"trim_strings": True}, + filter=None, ), DQProfile( name="is_not_null", column="maintenance_date", description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + filter=None, ), DQProfile( name="min_max", column="maintenance_date", description="Real min/max values were used", - parameters={"min": date(2025, 4, 29), "max": date(2025, 7, 30)}, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + parameters={"min": date(2025, 4, 1), "max": date(2025, 7, 30)}, + filter=None, ), DQProfile( name="is_not_null", column="cost", description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + filter=None, ), DQProfile( name="min_max", column="cost", parameters={ - "min": Decimal('100.00'), - "max": Decimal('300.00'), + "min": Decimal('-500.00'), + "max": Decimal('1200.50'), }, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + filter=None, description="Real min/max values were used", ), DQProfile( name="is_not_null", column="next_scheduled_date", description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + filter=None, ), DQProfile( name="min_max", column="next_scheduled_date", description="Real min/max values were used", - parameters={"min": date(2025, 7, 15), "max": date(2025, 12, 1)}, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + parameters={"min": date(2024, 4, 20), "max": date(2026, 4, 1)}, + filter=None, ), DQProfile( name="is_not_null", column="safety_check_passed", description=None, - filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + filter=None, ), ] - assert filtered_df.count() == 3 assert len(stats.keys()) > 0 assert rules == expected_rules diff --git a/tests/integration/test_rules_generator.py b/tests/integration/test_rules_generator.py index 3dc9cfa99..c17db4b07 100644 --- a/tests/integration/test_rules_generator.py +++ b/tests/integration/test_rules_generator.py @@ -1,9 +1,6 @@ import datetime from decimal import Decimal -# from datetime import date, datetime, timezone - - from databricks.labs.dqx.profiler.generator import DQGenerator from databricks.labs.dqx.profiler.profiler import DQProfile @@ -47,13 +44,14 @@ def test_generate_dq_rules(ws): expected = [ { "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"}}, - "filter": None, "name": "vendor_id_is_null", "criticality": "error", }, { - "check": {"function": "is_in_list", "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}}, - "filter": None, + "check": { + "function": "is_in_list", + "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}, + }, "name": "vendor_id_other_value", "criticality": "error", }, @@ -62,7 +60,6 @@ def test_generate_dq_rules(ws): "function": "is_not_null_and_not_empty", "arguments": {"column": "vendor_id", "trim_strings": True}, }, - "filter": None, "name": "vendor_id_is_null_or_empty", "criticality": "error", }, @@ -71,7 +68,6 @@ def test_generate_dq_rules(ws): "function": "is_in_range", "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}, }, - "filter": None, "name": "rate_code_id_isnt_in_range", "criticality": "error", }, @@ -85,13 +81,14 @@ def test_generate_dq_rules_warn(ws): expected = [ { "check": {"function": "is_not_null", "arguments": {"column": "vendor_id"}}, - "filter": None, "name": "vendor_id_is_null", "criticality": "warn", }, { - "check": {"function": "is_in_list", "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}}, - "filter": None, + "check": { + "function": "is_in_list", + "arguments": {"column": "vendor_id", "allowed": ["1", "4", "2"]}, + }, "name": "vendor_id_other_value", "criticality": "warn", }, @@ -100,7 +97,6 @@ def test_generate_dq_rules_warn(ws): "function": "is_not_null_and_not_empty", "arguments": {"column": "vendor_id", "trim_strings": True}, }, - "filter": None, "name": "vendor_id_is_null_or_empty", "criticality": "warn", }, @@ -109,7 +105,6 @@ def test_generate_dq_rules_warn(ws): "function": "is_in_range", "arguments": {"column": "rate_code_id", "min_limit": 1, "max_limit": 265}, }, - "filter": None, "name": "rate_code_id_isnt_in_range", "criticality": "warn", }, @@ -131,7 +126,7 @@ def test_generate_dq_no_rules(ws): def test_generate_dq_rules_dataframe_filter(ws): generator = DQGenerator(ws) - test_rules = [ + test_rules_filter = [ DQProfile( name="is_not_null", column="machine_id", @@ -161,7 +156,7 @@ def test_generate_dq_rules_dataframe_filter(ws): ), DQProfile(name="is_not_null_or_empty", column="vendor_id", parameters={"trim_strings": True}), ] - expectations = generator.generate_dq_rules(test_rules) + expectations = generator.generate_dq_rules(test_rules_filter) expected = [ { @@ -178,19 +173,16 @@ def test_generate_dq_rules_dataframe_filter(ws): }, { "check": {"function": "is_not_null", "arguments": {"column": "cost"}}, - "filter": None, "name": "cost_is_null", "criticality": "error", }, { "check": {"function": "is_not_null", "arguments": {"column": "next_scheduled_date"}}, - "filter": None, "name": "next_scheduled_date_is_null", "criticality": "error", }, { "check": {"function": "is_not_null", "arguments": {"column": "safety_check_passed"}}, - "filter": None, "name": "safety_check_passed_is_null", "criticality": "error", }, @@ -199,7 +191,6 @@ def test_generate_dq_rules_dataframe_filter(ws): "function": "is_not_null_and_not_empty", "arguments": {"column": "vendor_id", "trim_strings": True}, }, - "filter": None, "name": "vendor_id_is_null_or_empty", "criticality": "error", }, @@ -209,7 +200,7 @@ def test_generate_dq_rules_dataframe_filter(ws): def test_generate_dq_rules_dataframe_filter_none(ws): generator = DQGenerator(ws) - test_rules = [ + test_rules_no_filter = [ DQProfile( name="is_not_null", column="machine_id", @@ -230,7 +221,7 @@ def test_generate_dq_rules_dataframe_filter_none(ws): ), DQProfile(name="is_not_null_or_empty", column="vendor_id", parameters={"trim_strings": True}, filter=None), ] - expectations = generator.generate_dq_rules(test_rules) + expectations = generator.generate_dq_rules(test_rules_no_filter) expected = [ { From a2aeb158f7e1a6c052855c2d87391e7a86578ed5 Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Tue, 30 Sep 2025 12:09:30 -0300 Subject: [PATCH 26/34] Refactor DQGenerator class after fmt sugestions modifications --- src/databricks/labs/dqx/profiler/generator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/labs/dqx/profiler/generator.py b/src/databricks/labs/dqx/profiler/generator.py index ae4e64df4..f0849b60b 100644 --- a/src/databricks/labs/dqx/profiler/generator.py +++ b/src/databricks/labs/dqx/profiler/generator.py @@ -37,7 +37,7 @@ def generate_dq_rules(self, profiles: list[DQProfile] | None = None, level: str if expr: if dataset_filter is not None: - expr["filter"] = filter + expr["filter"] = dataset_filter dq_rules.append(expr) status = DQEngine.validate_checks(dq_rules) From fe9b6e3887cfe2d0a6720ec6c13828b5b36ef096 Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Wed, 1 Oct 2025 09:22:20 -0300 Subject: [PATCH 27/34] Removing filter inside check key in the load checks from uc volume test file --- .../test_load_checks_from_uc_volume.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/tests/integration/test_load_checks_from_uc_volume.py b/tests/integration/test_load_checks_from_uc_volume.py index cbbcc20bd..96ea6cbb0 100644 --- a/tests/integration/test_load_checks_from_uc_volume.py +++ b/tests/integration/test_load_checks_from_uc_volume.py @@ -221,32 +221,31 @@ def test_save_checks_in_volume_file_as_json(ws, make_schema, make_volume, instal TEST_CHECKS_FILTER = [ { "criticality": "error", - "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}, "filter": None}, + "filter": None, + "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}}, }, { - "check": { - "function": "is_not_null", - "arguments": {"column": "next_scheduled_date"}, - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", - }, + "criticality": "error", + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + "check": {"function": "is_not_null", "arguments": {"column": "next_scheduled_date"}}, "name": "next_scheduled_date_is_null", - "criticality": "Error", }, ] EXPECTED_CHECKS_FILTER = [ { "criticality": "error", - "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}, "filter": None}, + "filter": None, + "check": {"function": "is_not_null", "for_each_column": ["col1", "col2"], "arguments": {}}, }, { + "criticality": "error", + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", "check": { "function": "is_not_null", "arguments": {"column": "next_scheduled_date"}, - "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", }, "name": "next_scheduled_date_is_null", - "criticality": "Error", }, ] From e01420b72f4a27578a5ce54d18e5b499434bc25f Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 1 Oct 2025 19:52:03 +0200 Subject: [PATCH 28/34] Update tests/integration/test_save_checks_to_workspace_file.py --- tests/integration/test_save_checks_to_workspace_file.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/test_save_checks_to_workspace_file.py b/tests/integration/test_save_checks_to_workspace_file.py index 550f79102..8746600a9 100644 --- a/tests/integration/test_save_checks_to_workspace_file.py +++ b/tests/integration/test_save_checks_to_workspace_file.py @@ -26,7 +26,7 @@ "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", }, "name": "next_scheduled_date_is_null", - "criticality": "Error", + "criticality": "error", }, { "check": {"function": "is_not_null", "arguments": {"column": "cost"}, "filter": None}, From 586293bb8d8a893a7698bbffdc2f1d44968f3835 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 1 Oct 2025 19:52:16 +0200 Subject: [PATCH 29/34] Update tests/integration/test_rules_generator.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- tests/integration/test_rules_generator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/test_rules_generator.py b/tests/integration/test_rules_generator.py index c17db4b07..3f5e6fc91 100644 --- a/tests/integration/test_rules_generator.py +++ b/tests/integration/test_rules_generator.py @@ -230,7 +230,7 @@ def test_generate_dq_rules_dataframe_filter_none(ws): "criticality": "error", }, { - "check": {"function": "is_in_list", "arguments": {"allowed": ['1', '4', '2'], "column": "vendor_id"}}, + "check": {"function": "is_in_list", "arguments": {"allowed": ["1", "4", "2"], "column": "vendor_id"}}, "criticality": "error", "name": "vendor_id_other_value", }, From f10b3c0fbb78d111703d5c8b4acdcab76977cd2c Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Wed, 1 Oct 2025 19:52:23 +0200 Subject: [PATCH 30/34] Update tests/integration/test_rules_generator.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- tests/integration/test_rules_generator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/test_rules_generator.py b/tests/integration/test_rules_generator.py index 3f5e6fc91..8fa49459c 100644 --- a/tests/integration/test_rules_generator.py +++ b/tests/integration/test_rules_generator.py @@ -166,7 +166,7 @@ def test_generate_dq_rules_dataframe_filter(ws): "criticality": "error", }, { - 'check': {"function": "is_in_list", "arguments": {"allowed": ['1', '4', '2'], "column": "vendor_id"}}, + "check": {"function": "is_in_list", "arguments": {"allowed": ["1", "4", "2"], "column": "vendor_id"}}, "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", "criticality": "error", "name": "vendor_id_other_value", From b7c8911a3f767ed443a899ce03172ee0ec1a07ee Mon Sep 17 00:00:00 2001 From: STEFANOVIVAS Date: Thu, 2 Oct 2025 11:30:08 -0300 Subject: [PATCH 31/34] Update doc to add filter feature reference. Add filter attribute in the ProfilerConfig dataclass so that workflows can use it. --- docs/dqx/docs/guide/data_profiling.mdx | 5 +++++ docs/dqx/docs/installation.mdx | 1 + docs/dqx/docs/reference/profiler.mdx | 2 ++ src/databricks/labs/dqx/config.py | 1 + 4 files changed, 9 insertions(+) diff --git a/docs/dqx/docs/guide/data_profiling.mdx b/docs/dqx/docs/guide/data_profiling.mdx index 0118c9371..ef996fd2d 100644 --- a/docs/dqx/docs/guide/data_profiling.mdx +++ b/docs/dqx/docs/guide/data_profiling.mdx @@ -240,6 +240,7 @@ The following fields from the [configuration file](/docs/installation/#configura - `sample_fraction`: fraction of data to sample for profiling. - `sample_seed`: seed for reproducible sampling. - `limit`: maximum number of records to analyze. + - `filter`: filter data source as a string SQL expression - `serverless_clusters`: whether to use serverless clusters for running the workflow (default: 'true'). Using serverless clusters is recommended as it allows for automated cluster management and scaling. - `profiler_spark_conf`: (optional) spark configuration to use for the profiler workflow, only applicable if `serverless_clusters` is set to 'false'. - `profiler_override_clusters`: (optional) cluster configuration to use for profiler workflow, only applicable if `serverless_clusters` is set to 'false'. @@ -258,6 +259,7 @@ Example of the configuration file (relevant fields only): limit: 1000 sample_fraction: 0.3 summary_stats_file: profile_summary_stats.yml + filter: "maintenance_type = 'preventive'" ``` ## Profiling options @@ -299,6 +301,9 @@ The profiler supports extensive configuration options to customize the profiling # Value rounding "round": True, # Round min/max values for cleaner rules + + # Filter Options + "filter": None, # Filter data source as a string SQL expression } ws = WorkspaceClient() diff --git a/docs/dqx/docs/installation.mdx b/docs/dqx/docs/installation.mdx index 0b1854f91..ffc59b167 100644 --- a/docs/dqx/docs/installation.mdx +++ b/docs/dqx/docs/installation.mdx @@ -223,6 +223,7 @@ run_configs: # <- list of run configurations, each run co sample_fraction: 0.3 # <- fraction of data to sample in the profiler (30%) sample_seed: 30 # <- optional seed for reproducible sampling limit: 1000 # <- limit the number of records to profile + filter: "maintenance_type = 'preventive'" # <- generate profile in a subset of your data source warehouse_id: your-warehouse-id # <- warehouse id for refreshing dashboard diff --git a/docs/dqx/docs/reference/profiler.mdx b/docs/dqx/docs/reference/profiler.mdx index 6ec62e2ec..3c9cae57b 100644 --- a/docs/dqx/docs/reference/profiler.mdx +++ b/docs/dqx/docs/reference/profiler.mdx @@ -68,6 +68,7 @@ The profiler supports extensive configuration options to customize behavior: | `sample_fraction` | `0.3` | Sample 30% of the data for profiling | | `sample_seed` | `None` | Seed for sampling (None = random) | | `limit` | `1000` | Maximum number of records to analyze | +| `filter` | `None` | Filter data source as a string SQL expression | ## DQProfile Structure @@ -80,6 +81,7 @@ class DQProfile: column: str # Column name the rule applies to description: str | None = None # Optional description of how the rule was generated parameters: dict[str, Any] | None = None # Optional parameters for the rule + filter: str | None = None # Optional filter to be applied to the data source ``` ## DQGenerator methods diff --git a/src/databricks/labs/dqx/config.py b/src/databricks/labs/dqx/config.py index 8ae65f777..0279d72ed 100644 --- a/src/databricks/labs/dqx/config.py +++ b/src/databricks/labs/dqx/config.py @@ -49,6 +49,7 @@ class ProfilerConfig: sample_fraction: float = 0.3 # fraction of data to sample (30%) sample_seed: int | None = None # seed for sampling limit: int = 1000 # limit the number of records to profile + filter: str | None = None # filter to apply to the data before profiling @dataclass From f64c20776246cc8cc278d10146a8d1a3fd77f756 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Thu, 2 Oct 2025 20:52:24 +0200 Subject: [PATCH 32/34] pass filter to profiler workflow refactor --- docs/dqx/docs/guide/data_profiling.mdx | 7 ++-- docs/dqx/docs/guide/quality_checks_apply.mdx | 1 + docs/dqx/docs/reference/profiler.mdx | 2 +- .../labs/dqx/profiler/profiler_runner.py | 1 + tests/integration/test_profiler.py | 33 +++++++++++++------ tests/integration/test_profiler_runner.py | 4 +-- 6 files changed, 32 insertions(+), 16 deletions(-) diff --git a/docs/dqx/docs/guide/data_profiling.mdx b/docs/dqx/docs/guide/data_profiling.mdx index ef996fd2d..0e3d71fb9 100644 --- a/docs/dqx/docs/guide/data_profiling.mdx +++ b/docs/dqx/docs/guide/data_profiling.mdx @@ -240,7 +240,7 @@ The following fields from the [configuration file](/docs/installation/#configura - `sample_fraction`: fraction of data to sample for profiling. - `sample_seed`: seed for reproducible sampling. - `limit`: maximum number of records to analyze. - - `filter`: filter data source as a string SQL expression + - `filter`: filter for the input data as a string SQL expression - `serverless_clusters`: whether to use serverless clusters for running the workflow (default: 'true'). Using serverless clusters is recommended as it allows for automated cluster management and scaling. - `profiler_spark_conf`: (optional) spark configuration to use for the profiler workflow, only applicable if `serverless_clusters` is set to 'false'. - `profiler_override_clusters`: (optional) cluster configuration to use for profiler workflow, only applicable if `serverless_clusters` is set to 'false'. @@ -303,7 +303,7 @@ The profiler supports extensive configuration options to customize the profiling "round": True, # Round min/max values for cleaner rules # Filter Options - "filter": None, # Filter data source as a string SQL expression + "filter": None, # No filter (SQL expression) } ws = WorkspaceClient() @@ -327,6 +327,7 @@ The profiler supports extensive configuration options to customize the profiling - `sample_fraction`: fraction of data to sample for profiling. - `sample_seed`: seed for reproducible sampling. - `limit`: maximum number of records to analyze. + - `filter`: string SQL expression to filter the input data You can open the configuration file to check available run configs and adjust the settings if needed: ```commandline @@ -404,7 +405,7 @@ Wildcard patterns are supported, allowing you to match table names and apply spe To ensure efficiency, summary statistics are computed on a **sampled subset** of your data, not the entire dataset. By default, the profiler proceeds as follows: -1. It **samples** the rows of the dataset you provide by using Spark’s `DataFrame.sample` (default `sample_fraction = 0.3`), which is probabilistic, hence the sampled row count is not guaranteed to be exact; the number may be slightly higher or lower. +1. It **samples** the rows of the dataset you provide by using Spark’s `DataFrame.sample` (default `sample_fraction = 0.3`), which is probabilistic, hence the sampled row count is not guaranteed to be exact; the number may be slightly higher or lower. No filters are applied by default. 2. A **limit** is applied to the sampled rows (default `limit = 1000`). 3. Statistics are computed on this final **sampled-and-limited** DataFrame. diff --git a/docs/dqx/docs/guide/quality_checks_apply.mdx b/docs/dqx/docs/guide/quality_checks_apply.mdx index 8297ff8bd..d45ab1081 100644 --- a/docs/dqx/docs/guide/quality_checks_apply.mdx +++ b/docs/dqx/docs/guide/quality_checks_apply.mdx @@ -493,6 +493,7 @@ The following fields from the [configuration file](/docs/installation/#configura - `sample_fraction`: fraction of data to sample for profiling. - `sample_seed`: seed for reproducible sampling. - `limit`: maximum number of records to analyze. + - `filter`: filter for the input data as a string SQL expression (default: None). - `extra_params`: (optional) extra parameters to pass to the jobs such as result column names and user_metadata - `custom_check_functions`: (optional) custom check functions defined in Python files that can be used in the quality checks. - `reference_tables`: (optional) reference tables that can be used in the quality checks. diff --git a/docs/dqx/docs/reference/profiler.mdx b/docs/dqx/docs/reference/profiler.mdx index 3c9cae57b..e6dd07120 100644 --- a/docs/dqx/docs/reference/profiler.mdx +++ b/docs/dqx/docs/reference/profiler.mdx @@ -68,7 +68,7 @@ The profiler supports extensive configuration options to customize behavior: | `sample_fraction` | `0.3` | Sample 30% of the data for profiling | | `sample_seed` | `None` | Seed for sampling (None = random) | | `limit` | `1000` | Maximum number of records to analyze | -| `filter` | `None` | Filter data source as a string SQL expression | +| `filter` | `None` | Filter for the input data as a string SQL expression | ## DQProfile Structure diff --git a/src/databricks/labs/dqx/profiler/profiler_runner.py b/src/databricks/labs/dqx/profiler/profiler_runner.py index 7626eee60..f6e2b38d9 100644 --- a/src/databricks/labs/dqx/profiler/profiler_runner.py +++ b/src/databricks/labs/dqx/profiler/profiler_runner.py @@ -57,6 +57,7 @@ def run( "sample_fraction": profiler_config.sample_fraction, "sample_seed": profiler_config.sample_seed, "limit": profiler_config.limit, + "filter": profiler_config.filter, }, ) checks = self.generator.generate_dq_rules(profiles) # use default criticality level "error" diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index 991be9824..26922159e 100644 --- a/tests/integration/test_profiler.py +++ b/tests/integration/test_profiler.py @@ -203,7 +203,7 @@ def test_profiler_rounding_midnight_behavior(spark, ws): def test_profiler_non_default_profile_options(spark, ws): - inp_schema = T.StructType( + input_schema = T.StructType( [ T.StructField("t1", T.IntegerType()), T.StructField("t2", T.StringType()), @@ -221,8 +221,16 @@ def test_profiler_non_default_profile_options(spark, ws): ), ] ) - inp_df = spark.createDataFrame( + input_df = spark.createDataFrame( [ + [ + 0, + " test ", + { + "ns1": datetime.fromisoformat("2023-01-08T10:00:11+00:00"), + "s2": {"ns2": "test", "ns3": date.fromisoformat("2023-01-08")}, + }, + ], [ 1, " test ", @@ -248,7 +256,7 @@ def test_profiler_non_default_profile_options(spark, ws): }, ], ], - schema=inp_schema, + schema=input_schema, ) profiler = DQProfiler(ws) @@ -265,30 +273,35 @@ def test_profiler_non_default_profile_options(spark, ws): "sample_fraction": 1.0, # fraction of data to sample "sample_seed": None, # seed for sampling "limit": 1000, # limit the number of samples + "filter": "t1 > 0" # filter out the first row } - stats, rules = profiler.profile(inp_df, columns=inp_df.columns, options=profile_options) + stats, rules = profiler.profile(input_df, columns=input_df.columns, options=profile_options) expected_rules = [ - DQProfile(name="is_not_null", column="t1", description=None, parameters=None), + DQProfile(name="is_not_null", column="t1", description=None, parameters=None, filter="t1 > 0"), DQProfile( - name="min_max", column="t1", description="Real min/max values were used", parameters={"min": 1, "max": 3} + name="min_max", column="t1", description="Real min/max values were used", parameters={"min": 1, "max": 3}, + filter="t1 > 0" ), - DQProfile(name='is_not_null_or_empty', column='t2', description=None, parameters={'trim_strings': False}), - DQProfile(name="is_not_null", column="s1.ns1", description=None, parameters=None), + DQProfile(name='is_not_null_or_empty', column='t2', description=None, parameters={'trim_strings': False}, + filter="t1 > 0"), + DQProfile(name="is_not_null", column="s1.ns1", description=None, parameters=None, filter="t1 > 0"), DQProfile( name="min_max", column="s1.ns1", description="Real min/max values were used", parameters={'max': datetime(2023, 1, 8, 10, 0, 11), 'min': datetime(2023, 1, 6, 10, 0, 11)}, + filter="t1 > 0" ), - DQProfile(name="is_not_null", column="s1.s2.ns2", description=None, parameters=None), - DQProfile(name="is_not_null", column="s1.s2.ns3", description=None, parameters=None), + DQProfile(name="is_not_null", column="s1.s2.ns2", description=None, parameters=None, filter="t1 > 0"), + DQProfile(name="is_not_null", column="s1.s2.ns3", description=None, parameters=None, filter="t1 > 0"), DQProfile( name="min_max", column="s1.s2.ns3", description="Real min/max values were used", parameters={"min": date(2023, 1, 6), "max": date(2023, 1, 8)}, + filter="t1 > 0" ), ] print(stats) diff --git a/tests/integration/test_profiler_runner.py b/tests/integration/test_profiler_runner.py index 4441ed89b..175edbc90 100644 --- a/tests/integration/test_profiler_runner.py +++ b/tests/integration/test_profiler_runner.py @@ -67,10 +67,10 @@ def test_profiler_runner(ws, spark, installation_ctx, make_schema, make_table): table = make_table( catalog_name=catalog_name, schema_name=schema.name, - ctas="SELECT * FROM VALUES (1, 'a'), (2, 'b'), (3, NULL) AS data(id, name)", + ctas="SELECT * FROM VALUES (0, 'a'),(1, 'a'), (2, 'b'), (3, NULL) AS data(id, name)", ) input_config = InputConfig(location=table.full_name) - profiler_config = ProfilerConfig(sample_fraction=1.0) + profiler_config = ProfilerConfig(sample_fraction=1.0, filter="id > 0") checks, summary_stats = runner.run(input_config=input_config, profiler_config=profiler_config) From 92187ca8c3ce3c5081346e77b95e483b4d1c5fcc Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Thu, 2 Oct 2025 20:53:35 +0200 Subject: [PATCH 33/34] fmt --- tests/integration/test_profiler.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index 26922159e..c523c57ef 100644 --- a/tests/integration/test_profiler.py +++ b/tests/integration/test_profiler.py @@ -273,7 +273,7 @@ def test_profiler_non_default_profile_options(spark, ws): "sample_fraction": 1.0, # fraction of data to sample "sample_seed": None, # seed for sampling "limit": 1000, # limit the number of samples - "filter": "t1 > 0" # filter out the first row + "filter": "t1 > 0", # filter out the first row } stats, rules = profiler.profile(input_df, columns=input_df.columns, options=profile_options) @@ -281,18 +281,26 @@ def test_profiler_non_default_profile_options(spark, ws): expected_rules = [ DQProfile(name="is_not_null", column="t1", description=None, parameters=None, filter="t1 > 0"), DQProfile( - name="min_max", column="t1", description="Real min/max values were used", parameters={"min": 1, "max": 3}, - filter="t1 > 0" + name="min_max", + column="t1", + description="Real min/max values were used", + parameters={"min": 1, "max": 3}, + filter="t1 > 0", + ), + DQProfile( + name='is_not_null_or_empty', + column='t2', + description=None, + parameters={'trim_strings': False}, + filter="t1 > 0", ), - DQProfile(name='is_not_null_or_empty', column='t2', description=None, parameters={'trim_strings': False}, - filter="t1 > 0"), DQProfile(name="is_not_null", column="s1.ns1", description=None, parameters=None, filter="t1 > 0"), DQProfile( name="min_max", column="s1.ns1", description="Real min/max values were used", parameters={'max': datetime(2023, 1, 8, 10, 0, 11), 'min': datetime(2023, 1, 6, 10, 0, 11)}, - filter="t1 > 0" + filter="t1 > 0", ), DQProfile(name="is_not_null", column="s1.s2.ns2", description=None, parameters=None, filter="t1 > 0"), DQProfile(name="is_not_null", column="s1.s2.ns3", description=None, parameters=None, filter="t1 > 0"), @@ -301,7 +309,7 @@ def test_profiler_non_default_profile_options(spark, ws): column="s1.s2.ns3", description="Real min/max values were used", parameters={"min": date(2023, 1, 6), "max": date(2023, 1, 8)}, - filter="t1 > 0" + filter="t1 > 0", ), ] print(stats) From 2c4576c06ee77af044933aac70dce21e737b7d25 Mon Sep 17 00:00:00 2001 From: Marcin Wojtyczka Date: Thu, 2 Oct 2025 21:07:00 +0200 Subject: [PATCH 34/34] fixed tests --- tests/conftest.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 5766c0846..98d2ba6f3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -361,8 +361,8 @@ def expected_checks(): @pytest.fixture -def make_local_check_file_as_yaml(checks_yaml_content): - file_path = "checks.yml" +def make_local_check_file_as_yaml(checks_yaml_content, make_random): + file_path = f"checks_{make_random(10).lower()}.yml" with open(file_path, "w", encoding="utf-8") as f: f.write(checks_yaml_content) yield file_path @@ -371,8 +371,8 @@ def make_local_check_file_as_yaml(checks_yaml_content): @pytest.fixture -def make_local_check_file_as_yaml_diff_ext(checks_yaml_content): - file_path = "checks.yaml" +def make_local_check_file_as_yaml_diff_ext(checks_yaml_content, make_random): + file_path = f"checks_{make_random(10).lower()}.yaml" with open(file_path, "w", encoding="utf-8") as f: f.write(checks_yaml_content) yield file_path @@ -381,8 +381,8 @@ def make_local_check_file_as_yaml_diff_ext(checks_yaml_content): @pytest.fixture -def make_local_check_file_as_json(checks_json_content): - file_path = "checks.json" +def make_local_check_file_as_json(checks_json_content, make_random): + file_path = f"checks_{make_random(10).lower()}.json" with open(file_path, "w", encoding="utf-8") as f: f.write(checks_json_content) yield file_path @@ -391,8 +391,8 @@ def make_local_check_file_as_json(checks_json_content): @pytest.fixture -def make_invalid_local_check_file_as_yaml(checks_yaml_invalid_content): - file_path = "invalid_checks.yml" +def make_invalid_local_check_file_as_yaml(checks_yaml_invalid_content, make_random): + file_path = f"invalid_checks_{make_random(10).lower()}.yml" with open(file_path, "w", encoding="utf-8") as f: f.write(checks_yaml_invalid_content) yield file_path @@ -401,8 +401,8 @@ def make_invalid_local_check_file_as_yaml(checks_yaml_invalid_content): @pytest.fixture -def make_empty_local_yaml_file(): - file_path = "empty.yml" +def make_empty_local_yaml_file(make_random): + file_path = f"empty_{make_random(10).lower()}.yml" with open(file_path, "w", encoding="utf-8") as f: f.write("") yield file_path @@ -411,8 +411,8 @@ def make_empty_local_yaml_file(): @pytest.fixture -def make_empty_local_json_file(): - file_path = "empty.json" +def make_empty_local_json_file(make_random): + file_path = f"empty_{make_random(10).lower()}.json" with open(file_path, "w", encoding="utf-8") as f: f.write("{}") yield file_path @@ -421,8 +421,8 @@ def make_empty_local_json_file(): @pytest.fixture -def make_invalid_local_check_file_as_json(checks_json_invalid_content): - file_path = "invalid_checks.json" +def make_invalid_local_check_file_as_json(checks_json_invalid_content, make_random): + file_path = f"invalid_checks_{make_random(10).lower()}.json" with open(file_path, "w", encoding="utf-8") as f: f.write(checks_json_invalid_content) yield file_path