diff --git a/.gitignore b/.gitignore index 16a704bde..1f6e39ffe 100644 --- a/.gitignore +++ b/.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 + diff --git a/docs/dqx/docs/guide/data_profiling.mdx b/docs/dqx/docs/guide/data_profiling.mdx index 0118c9371..0e3d71fb9 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 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'. @@ -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, # No filter (SQL expression) } ws = WorkspaceClient() @@ -322,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 @@ -399,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/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..e6dd07120 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 for the input data 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 diff --git a/src/databricks/labs/dqx/profiler/generator.py b/src/databricks/labs/dqx/profiler/generator.py index 5ca43681d..f0849b60b 100644 --- a/src/databricks/labs/dqx/profiler/generator.py +++ b/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", diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index ceccaa000..9be22b861 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/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 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/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 diff --git a/tests/integration/test_load_checks_from_uc_volume.py b/tests/integration/test_load_checks_from_uc_volume.py index 841496e73..96ea6cbb0 100644 --- a/tests/integration/test_load_checks_from_uc_volume.py +++ b/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." diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index 7e920ddfc..c523c57ef 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,43 @@ 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) @@ -1255,3 +1276,318 @@ def test_profile_tables_no_tables_or_patterns(ws): with pytest.raises(MissingParameterError, match="Either 'tables' or 'patterns' must be provided"): profiler.profile_tables() + + +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("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), + Decimal("450.00"), + date(2025, 7, 1), + True, + ), + ( + "MCH-002", + "corrective", + date(2025, 4, 15), + Decimal("1200.50"), + date(2026, 4, 1), + False, + ), + ( + "MCH-003", + None, + date(2025, 4, 20), + Decimal("-500.00"), + date(2024, 4, 20), + None, + ), + ( + "MCH-001", + "predictive", + date(2025, 4, 25), + Decimal("800.00"), + date(2025, 10, 1), + True, + ), + ( + "MCH-002", + "preventive", + date(2025, 4, 29), + Decimal("300.50"), + date(2025, 7, 15), + True, + ), + ( + "MCH-003", + "corrective", + date(2025, 4, 30), + Decimal("150.00"), + 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, + ), + ] + + input_df = spark.createDataFrame(maintenance_data, schema=schema) + + custom_options = { + "sample_fraction": None, + "round": False, + "limit": None, + "filter": "machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + } + + profiler = DQProfiler(ws) + + 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'", + ), + DQProfile( + name="is_not_null", + column="maintenance_type", + description=None, + 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'", + ), + 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'", + ), + 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", + parameters={ + "min": Decimal('100.00'), + "max": Decimal('300.50'), + }, + 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, + 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'", + ), + DQProfile( + name="is_not_null", + column="safety_check_passed", + description=None, + filter="machine_id IN ('MCH-002', 'MCH-003') AND maintenance_type = 'preventive'", + ), + ] + + assert len(stats.keys()) > 0 + assert rules == expected_rules + + +def test_profile_with_no_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("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), + Decimal("450.00"), + date(2025, 7, 1), + True, + ), + ( + "MCH-002", + "corrective", + date(2025, 4, 15), + Decimal("1200.50"), + date(2026, 4, 1), + False, + ), + ( + "MCH-003", + None, + date(2025, 4, 20), + Decimal("-500.00"), + date(2024, 4, 20), + False, + ), + ( + "MCH-001", + "predictive", + date(2025, 4, 25), + Decimal("800.00"), + date(2025, 10, 1), + True, + ), + ( + "MCH-002", + "preventive", + date(2025, 4, 29), + Decimal("300.50"), + date(2025, 7, 15), + True, + ), + ( + "MCH-003", + "corrective", + date(2025, 4, 30), + Decimal("150.00"), + 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, + ), + ] + + input_df = spark.createDataFrame(maintenance_data, schema=schema) + + profiler = DQProfiler(ws) + custom_options = { + "sample_fraction": None, + "round": True, + "limit": None, + "filter": None, + } + stats, rules = profiler.profile(input_df, options=custom_options) + + expected_rules = [ + DQProfile( + name="is_not_null", + column="machine_id", + description=None, + filter=None, + ), + DQProfile( + name="is_not_null_or_empty", + column="maintenance_type", + description=None, + parameters={"trim_strings": True}, + filter=None, + ), + DQProfile( + name="is_not_null", + column="maintenance_date", + description=None, + filter=None, + ), + DQProfile( + name="min_max", + column="maintenance_date", + description="Real min/max values were used", + parameters={"min": date(2025, 4, 1), "max": date(2025, 7, 30)}, + filter=None, + ), + DQProfile( + name="is_not_null", + column="cost", + description=None, + filter=None, + ), + DQProfile( + name="min_max", + column="cost", + parameters={ + "min": Decimal('-500.00'), + "max": Decimal('1200.50'), + }, + filter=None, + description="Real min/max values were used", + ), + DQProfile( + name="is_not_null", + column="next_scheduled_date", + description=None, + filter=None, + ), + DQProfile( + name="min_max", + column="next_scheduled_date", + description="Real min/max values were used", + 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=None, + ), + ] + + assert len(stats.keys()) > 0 + assert rules == expected_rules 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) diff --git a/tests/integration/test_rules_generator.py b/tests/integration/test_rules_generator.py index aec7b81d6..8fa49459c 100644 --- a/tests/integration/test_rules_generator.py +++ b/tests/integration/test_rules_generator.py @@ -122,3 +122,130 @@ 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_filter = [ + 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_filter) + + 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"}}, + "name": "cost_is_null", + "criticality": "error", + }, + { + "check": {"function": "is_not_null", "arguments": {"column": "next_scheduled_date"}}, + "name": "next_scheduled_date_is_null", + "criticality": "error", + }, + { + "check": {"function": "is_not_null", "arguments": {"column": "safety_check_passed"}}, + "name": "safety_check_passed_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 + + +def test_generate_dq_rules_dataframe_filter_none(ws): + generator = DQGenerator(ws) + test_rules_no_filter = [ + 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_no_filter) + + 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 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 49475307a..f27dd3ce9 100644 --- a/tests/integration/test_save_and_load_checks_from_table.py +++ b/tests/integration/test_save_and_load_checks_from_table.py @@ -26,6 +26,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"}, }, { @@ -51,10 +52,8 @@ { "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"}, }, { @@ -129,7 +128,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 797790538..8746600a9 100644 --- a/tests/integration/test_save_checks_to_workspace_file.py +++ b/tests/integration/test_save_checks_to_workspace_file.py @@ -18,7 +18,21 @@ { "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'", + }, + "name": "next_scheduled_date_is_null", + "criticality": "error", + }, + { + "check": {"function": "is_not_null", "arguments": {"column": "cost"}, "filter": None}, + "name": "cost_is_null", + "criticality": "error", + }, ]