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
+

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()

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
 

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

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

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,49 @@ 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",
+        "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"}},
+        "name": "next_scheduled_date_is_null",
+    },
+]
+
+EXPECTED_CHECKS_FILTER = [
+    {
+        "criticality": "error",
+        "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"},
+        },
+        "name": "next_scheduled_date_is_null",
+    },
+]
+
+
+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."