Profile subset dataframe

.gitignore

@@ -50,6 +50,7 @@ htmlcov/
 .cache
 nosetests.xml
 coverage.xml
+coverage-unit.xml
 *.cover
 *.py,cover
 .hypothesis/
@@ -168,3 +169,7 @@ dev/cleanup.py
 # docgen
 docs/dqx/docs/reference/api
 !docs/dqx/docs/reference/api/index.mdx
+
+.databricks
+databricks.yml
+

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

@@ -29,11 +29,15 @@ 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
             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)
 
         status = DQEngine.validate_checks(dq_rules)
@@ -97,10 +101,7 @@ def dq_generate_min_max(column: str, level: str = "error", **params: dict):
             return {
                 "check": {
                     "function": "is_not_greater_than",
-                    "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)},
                 },
                 "name": f"{column}_not_greater_than",
                 "criticality": level,
@@ -110,10 +111,7 @@ def dq_generate_min_max(column: str, level: str = "error", **params: dict):
             return {
                 "check": {
                     "function": "is_not_less_than",
-                    "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)},
                 },
                 "name": f"{column}_not_less_than",
                 "criticality": level,
@@ -135,6 +133,7 @@ def dq_generate_is_not_null(column: str, level: str = "error", **params: dict):
                 A dictionary representing the data quality rule.
         """
         params = params or {}
+
         return {
             "check": {"function": "is_not_null", "arguments": {"column": column}},
             "name": f"{column}_is_null",
@@ -154,6 +153,7 @@ 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",

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

@@ -31,6 +31,7 @@ class DQProfile:
     column: str
     description: str | None = None
     parameters: dict[str, Any] | None = None
+    filter: str | None = None
 
 
 class DQProfiler(DQEngineBase):
@@ -53,6 +54,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": None,  # filter to apply to the dataset
     }
 
     @staticmethod
@@ -92,6 +94,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])
@@ -109,6 +112,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:
+            for rule in dq_rules:
+                rule.filter = filter
 
         return summary_stats, dq_rules
 
@@ -315,6 +322,10 @@ def _sample(df: DataFrame, opts: dict[str, Any]) -> DataFrame:
         sample_seed = opts.get("sample_seed", None)
         limit = opts.get("limit", None)
 
+        filter = opts.get("filter", None)
+
+        if filter:
+            df = df.filter(filter)
         if sample_fraction:
             df = df.sample(withReplacement=False, fraction=sample_fraction, seed=sample_seed)
         if limit:
@@ -395,6 +406,7 @@ def _calculate_metrics(
                 dq_rules.append(
                     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
@@ -405,7 +417,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})
+                    DQProfile(
+                        name="is_not_null_or_empty",
+                        column=field_name,
+                        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)
@@ -559,7 +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
+                name="min_max",
+                column=col_name,
+                parameters={"min": min_limit, "max": max_limit},
+                description=descr,
             )
 
         return None

tests/integration/test_load_checks_from_uc_volume.py

@@ -9,15 +9,33 @@
 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",
+    },
 ]