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 {}
+            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 dataset_filter is not None:
+                    expr["filter"] = dataset_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

@@ -32,6 +32,7 @@ class DQProfile:
     column: str
     description: str | None = None
     parameters: dict[str, Any] | None = None
+    filter: str | None = None
 
 
 class DQProfiler(DQEngineBase):
@@ -54,6 +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
+        "filter": None,  # filter to apply to the dataset
     }
 
     @staticmethod
@@ -93,6 +95,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])
@@ -318,7 +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)
 
+        if filter_dataset:
+            df = df.filter(filter_dataset)
         if sample_fraction:
             df = df.sample(withReplacement=False, fraction=sample_fraction, seed=sample_seed)
         if limit:
@@ -388,17 +394,30 @@ 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 (
             typ == T.StringType()
             and not any(  # does not make sense to add is_not_null_or_empty if is_not_null already exists
@@ -409,7 +428,12 @@ 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},
+                        filter=opts.get("filter", None),
+                    )
                 )
         if metrics["count_non_null"] > 0 and self._type_supports_min_max(typ):
             rule = self._extract_min_max(dst, field_name, typ, metrics, opts)
@@ -563,7 +587,11 @@ 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,
+                filter=opts.get("filter", None),
             )
 
         return None

tests/integration/test_load_checks_from_uc_volume.py

@@ -216,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."