diff --git a/src/databricks/labs/dqx/profiler/generator.py b/src/databricks/labs/dqx/profiler/generator.py index ca1b7d325..e26e1a139 100644 --- a/src/databricks/labs/dqx/profiler/generator.py +++ b/src/databricks/labs/dqx/profiler/generator.py @@ -1,12 +1,12 @@ from __future__ import annotations import logging +import datetime import json from collections.abc import Callable - from pyspark.sql import SparkSession - from databricks.sdk import WorkspaceClient + from databricks.labs.dqx.base import DQEngineBase from databricks.labs.dqx.config import LLMModelConfig, InputConfig from databricks.labs.dqx.engine import DQEngine @@ -239,25 +239,43 @@ def dq_generate_min_max(column: str, criticality: str = "error", **params: dict) Generates a data quality rule to check if a column's value is within a specified range. Args: - column: The name of the column to check. - criticality: The criticality of the rule as "warn" or "error" (default is "error"). - params: Additional parameters, including the minimum and maximum values. + column: The name of the column to check. + criticality: The criticality of the rule as "warn" or "error" (default is "error"). + params: Additional parameters, including the minimum and maximum values. Returns: - A dictionary representing the data quality rule, or None if no limits are provided. + A dictionary representing the data quality rule, or None if no limits are provided. """ 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 + if min_limit is None and max_limit is None: + return None + + def _is_num(value): + return isinstance(value, (int, float)) + + def _is_temporal(value): + return isinstance(value, (datetime.date, datetime.datetime)) + + def _same_family(value_a, value_b): + # numeric with numeric OR temporal with temporal + return any( + [ + _is_num(value_a) and _is_num(value_b), + _is_temporal(value_a) and _is_temporal(value_b), + ] + ) - if min_limit is not None and max_limit is not None: + # Both bounds + if min_limit is not None and max_limit is not None and _same_family(min_limit, max_limit): return { "check": { "function": "is_in_range", "arguments": { "column": column, + # pass through Python numeric types (int, float) without stringification + # except for temporal types (datetime/date) which are not Json serializable "min_limit": val_maybe_to_str(min_limit, include_sql_quotes=False), "max_limit": val_maybe_to_str(max_limit, include_sql_quotes=False), }, @@ -266,7 +284,8 @@ def dq_generate_min_max(column: str, criticality: str = "error", **params: dict) "criticality": criticality, } - if max_limit is not None: + # Only max + if max_limit is not None and (_is_num(max_limit) or _is_temporal(max_limit)): return { "check": { "function": "is_not_greater_than", @@ -276,7 +295,8 @@ def dq_generate_min_max(column: str, criticality: str = "error", **params: dict) "criticality": criticality, } - if min_limit is not None: + # Only min + if min_limit is not None and (_is_num(min_limit) or _is_temporal(min_limit)): return { "check": { "function": "is_not_less_than", diff --git a/tests/integration/test_generator.py b/tests/integration/test_generator.py index 5611cdcd6..3b20995a6 100644 --- a/tests/integration/test_generator.py +++ b/tests/integration/test_generator.py @@ -1,5 +1,5 @@ +import logging import datetime -from decimal import Decimal from databricks.labs.dqx.profiler.generator import DQGenerator from databricks.labs.dqx.profiler.profiler import DQProfile @@ -19,13 +19,13 @@ DQProfile( name="min_max", column="product_launch_date", - parameters={"min": datetime.date(2020, 1, 1), "max": None}, + parameters={"min": datetime.date(2020, 1, 2), "max": None}, description="Real min/max values were used", ), DQProfile( name="min_max", column="product_expiry_ts", - parameters={"min": None, "max": datetime.datetime(2020, 1, 1)}, + parameters={"min": None, "max": datetime.datetime(2020, 1, 2, 3, 4, 5)}, description="Real min/max values were used", ), DQProfile(name="is_random", column="vendor_id", parameters={"in": ["1", "4", "2"]}), @@ -33,13 +33,13 @@ name='min_max', column='d1', description='Real min/max values were used', - parameters={'max': Decimal('333323.00'), 'min': Decimal('1.23')}, + parameters={'max': 333323.00, 'min': 1.23}, ), ] -def test_generate_dq_rules(ws): - generator = DQGenerator(ws) +def test_generate_dq_rules(ws, spark): + generator = DQGenerator(ws, spark) expectations = generator.generate_dq_rules(test_rules) expected = [ { @@ -71,12 +71,36 @@ def test_generate_dq_rules(ws): "name": "rate_code_id_isnt_in_range", "criticality": "error", }, + { + "check": { + "function": "is_not_less_than", + "arguments": {"column": "product_launch_date", "limit": "2020-01-02"}, + }, + "name": "product_launch_date_not_less_than", + "criticality": "error", + }, + { + "check": { + "function": "is_not_greater_than", + "arguments": {"column": "product_expiry_ts", "limit": "2020-01-02T03:04:05.000000"}, + }, + "name": "product_expiry_ts_not_greater_than", + "criticality": "error", + }, + { + "check": { + "function": "is_in_range", + "arguments": {"column": "d1", "min_limit": 1.23, "max_limit": 333323.00}, + }, + "name": "d1_isnt_in_range", + "criticality": "error", + }, ] - assert expectations == expected + assert expectations == expected, f"Actual expectations: {expectations}" -def test_generate_dq_rules_warn(ws): - generator = DQGenerator(ws) +def test_generate_dq_rules_warn(ws, spark): + generator = DQGenerator(ws, spark) expectations = generator.generate_dq_rules(test_rules, criticality="warn") expected = [ { @@ -108,24 +132,54 @@ def test_generate_dq_rules_warn(ws): "name": "rate_code_id_isnt_in_range", "criticality": "warn", }, + { + "check": { + "function": "is_not_less_than", + "arguments": {"column": "product_launch_date", "limit": "2020-01-02"}, + }, + "name": "product_launch_date_not_less_than", + "criticality": "warn", + }, + { + "check": { + "function": "is_not_greater_than", + "arguments": {"column": "product_expiry_ts", "limit": "2020-01-02T03:04:05.000000"}, + }, + "name": "product_expiry_ts_not_greater_than", + "criticality": "warn", + }, + { + "check": { + "function": "is_in_range", + "arguments": {"column": "d1", "min_limit": 1.23, "max_limit": 333323.00}, + }, + "name": "d1_isnt_in_range", + "criticality": "warn", + }, ] - assert expectations == expected + assert expectations == expected, f"Actual expectations: {expectations}" + +def test_generate_dq_rules_logging(ws, spark, caplog): + # capture INFO from the generator module where the skip log is emitted + caplog.set_level(logging.INFO, logger="databricks.labs.dqx.profiler.generator") + + generator = DQGenerator(ws, spark) + # add an unknown rule to trigger the "skipping..." log + unknown_rule = DQProfile(name="is_random", column="vendor_id") + generator.generate_dq_rules(test_rules + [unknown_rule]) -def test_generate_dq_rules_logging(ws, caplog): - generator = DQGenerator(ws) - generator.generate_dq_rules(test_rules) assert "No rule 'is_random' for column 'vendor_id'. skipping..." in caplog.text -def test_generate_dq_no_rules(ws): - generator = DQGenerator(ws) +def test_generate_dq_no_rules(ws, spark): + generator = DQGenerator(ws, spark) expectations = generator.generate_dq_rules(None, criticality="warn") assert not expectations -def test_generate_dq_rules_dataframe_filter(ws): - generator = DQGenerator(ws) +def test_generate_dq_rules_dataframe_filter(ws, spark): + generator = DQGenerator(ws, spark) test_rules_filter = [ DQProfile( name="is_not_null", @@ -198,8 +252,8 @@ def test_generate_dq_rules_dataframe_filter(ws): assert expectations == expected -def test_generate_dq_rules_dataframe_filter_none(ws): - generator = DQGenerator(ws) +def test_generate_dq_rules_dataframe_filter_none(ws, spark): + generator = DQGenerator(ws, spark) test_rules_no_filter = [ DQProfile( name="is_not_null", diff --git a/tests/unit/test_generator_temporal.py b/tests/unit/test_generator_temporal.py new file mode 100644 index 000000000..a67ca7c98 --- /dev/null +++ b/tests/unit/test_generator_temporal.py @@ -0,0 +1,32 @@ +import datetime + +from databricks.labs.dqx.profiler.generator import DQGenerator + + +def test_date_both_bounds_is_in_range(set_utc_timezone): + result = DQGenerator.dq_generate_min_max( + "dcol", **{"min": datetime.date(2020, 1, 1), "max": datetime.date(2020, 12, 31)} + ) + assert result["check"]["function"] == "is_in_range" + args = result["check"]["arguments"] + assert args["column"] == "dcol" + assert args["min_limit"] == "2020-01-01" + assert args["max_limit"] == "2020-12-31" + + +def test_timestamp_only_min_is_not_less_than(set_utc_timezone): + timestamp = datetime.datetime(2024, 6, 1, 12, 0, 0) + result = DQGenerator.dq_generate_min_max("tscol", **{"min": timestamp, "max": None}) + assert result["check"]["function"] == "is_not_less_than" + args = result["check"]["arguments"] + assert args["column"] == "tscol" + assert args["limit"] == "2024-06-01T12:00:00.000000" + + +def test_timestamp_only_max_is_not_greater_than(set_utc_timezone): + timestamp = datetime.datetime(2024, 6, 30, 23, 59, 59) + result = DQGenerator.dq_generate_min_max("tscol", **{"min": None, "max": timestamp}) + assert result["check"]["function"] == "is_not_greater_than" + args = result["check"]["arguments"] + assert args["column"] == "tscol" + assert args["limit"] == "2024-06-30T23:59:59.000000"