Improve Generator to emit temporal checks for min/max date & datetime

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",

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,27 +19,27 @@
     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"]}),
     DQProfile(
         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",

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"