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

src/databricks/labs/dqx/check_funcs.py

@@ -4,6 +4,7 @@
 import ipaddress
 import uuid
 from collections.abc import Callable, Sequence
+from decimal import Decimal
 from enum import Enum
 from itertools import zip_longest
 import operator as py_operator
@@ -488,7 +489,7 @@ def is_not_in_near_future(column: str | Column, offset: int = 0, curr_timestamp:
 
 @register_rule("row")
 def is_equal_to(
-    column: str | Column, value: int | float | str | datetime.date | datetime.datetime | Column | None = None
+    column: str | Column, value: int | float | Decimal | str | datetime.date | datetime.datetime | Column | None = None
 ) -> Column:
     """Check whether the values in the input column are equal to the given value.
 
@@ -519,7 +520,7 @@ def is_equal_to(
 
 @register_rule("row")
 def is_not_equal_to(
-    column: str | Column, value: int | float | str | datetime.date | datetime.datetime | Column | None = None
+    column: str | Column, value: int | float | Decimal | str | datetime.date | datetime.datetime | Column | None = None
 ) -> Column:
     """Check whether the values in the input column are not equal to the given value.
 
@@ -550,13 +551,13 @@ def is_not_equal_to(
 
 @register_rule("row")
 def is_not_less_than(
-    column: str | Column, limit: int | float | datetime.date | datetime.datetime | str | Column | None = None
+    column: str | Column, limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None
 ) -> Column:
     """Checks whether the values in the input column are not less than the provided limit.
 
     Args:
         column: column to check; can be a string column name or a column expression
-        limit: limit to use in the condition as number, date, timestamp, column name or sql expression
+        limit: limit to use in the condition as number, date, timestamp, Decimal, column name or sql expression
 
     Returns:
         new Column
@@ -580,13 +581,13 @@ def is_not_less_than(
 
 @register_rule("row")
 def is_not_greater_than(
-    column: str | Column, limit: int | float | datetime.date | datetime.datetime | str | Column | None = None
+    column: str | Column, limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None
 ) -> Column:
     """Checks whether the values in the input column are not greater than the provided limit.
 
     Args:
         column: column to check; can be a string column name or a column expression
-        limit: limit to use in the condition as number, date, timestamp, column name or sql expression
+        limit: limit to use in the condition as number, date, timestamp, Decimal, column name or sql expression
 
     Returns:
         new Column
@@ -611,15 +612,15 @@ def is_not_greater_than(
 @register_rule("row")
 def is_in_range(
     column: str | Column,
-    min_limit: int | float | datetime.date | datetime.datetime | str | Column | None = None,
-    max_limit: int | float | datetime.date | datetime.datetime | str | Column | None = None,
+    min_limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None,
+    max_limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None,
 ) -> Column:
     """Checks whether the values in the input column are in the provided limits (inclusive of both boundaries).
 
     Args:
         column: column to check; can be a string column name or a column expression
-        min_limit: min limit to use in the condition as number, date, timestamp, column name or sql expression
-        max_limit: max limit to use in the condition as number, date, timestamp, column name or sql expression
+        min_limit: min limit to use in the condition as number, date, timestamp, Decimal, column name or sql expression
+        max_limit: max limit to use in the condition as number, date, timestamp, Decimal, column name or sql expression
 
     Returns:
         new Column
@@ -649,8 +650,8 @@ def is_in_range(
 @register_rule("row")
 def is_not_in_range(
     column: str | Column,
-    min_limit: int | float | datetime.date | datetime.datetime | str | Column | None = None,
-    max_limit: int | float | datetime.date | datetime.datetime | str | Column | None = None,
+    min_limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None,
+    max_limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None,
 ) -> Column:
     """Checks whether the values in the input column are outside the provided limits (inclusive of both boundaries).
 
@@ -1402,7 +1403,7 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) ->
 @register_rule("dataset")
 def is_aggr_not_greater_than(
     column: str | Column,
-    limit: int | float | str | Column,
+    limit: int | float | Decimal | str | Column,
     aggr_type: str = "count",
     group_by: list[str | Column] | None = None,
     row_filter: str | None = None,
@@ -1447,7 +1448,7 @@ def is_aggr_not_greater_than(
 @register_rule("dataset")
 def is_aggr_not_less_than(
     column: str | Column,
-    limit: int | float | str | Column,
+    limit: int | float | Decimal | str | Column,
     aggr_type: str = "count",
     group_by: list[str | Column] | None = None,
     row_filter: str | None = None,
@@ -1492,7 +1493,7 @@ def is_aggr_not_less_than(
 @register_rule("dataset")
 def is_aggr_equal(
     column: str | Column,
-    limit: int | float | str | Column,
+    limit: int | float | Decimal | str | Column,
     aggr_type: str = "count",
     group_by: list[str | Column] | None = None,
     row_filter: str | None = None,
@@ -1537,7 +1538,7 @@ def is_aggr_equal(
 @register_rule("dataset")
 def is_aggr_not_equal(
     column: str | Column,
-    limit: int | float | str | Column,
+    limit: int | float | Decimal | str | Column,
     aggr_type: str = "count",
     group_by: list[str | Column] | None = None,
     row_filter: str | None = None,
@@ -2693,7 +2694,7 @@ def _validate_aggregate_return_type(
 
 def _is_aggr_compare(
     column: str | Column,
-    limit: int | float | str | Column,
+    limit: int | float | Decimal | str | Column,
     aggr_type: str,
     aggr_params: dict[str, Any] | None,
     group_by: list[str | Column] | None,
@@ -2912,7 +2913,7 @@ def _cleanup_alias_name(column: str) -> str:
 
 
 def get_limit_expr(
-    limit: int | float | datetime.date | datetime.datetime | str | Column | None = None,
+    limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None,
 ) -> Column:
     """
     Generate a Spark Column expression for a limit value.

src/databricks/labs/dqx/profiler/generator.py

@@ -1,16 +1,16 @@
 from __future__ import annotations
 
 import logging
+import datetime
 import json
+from decimal import Decimal
 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
-from databricks.labs.dqx.profiler.common import val_maybe_to_str
 from databricks.labs.dqx.profiler.profiler import DQProfile
 from databricks.labs.dqx.telemetry import telemetry_logger
 from databricks.labs.dqx.errors import MissingParameterError
@@ -230,49 +230,64 @@ 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, Decimal))
+
+        def _is_temporal(value):
+            return isinstance(value, (datetime.date, datetime.datetime))
 
-        if min_limit is not None and max_limit is not None:
+        def _same_family(value_a, value_b):
+            # numeric with numeric OR temporal with temporal
+            if value_a is None or value_b is None:
+                return True
+            return any(
+                [
+                    _is_num(value_a) and _is_num(value_b),
+                    _is_temporal(value_a) and _is_temporal(value_b),
+                ]
+            )
+
+        # 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,
-                        "min_limit": val_maybe_to_str(min_limit, include_sql_quotes=False),
-                        "max_limit": val_maybe_to_str(max_limit, include_sql_quotes=False),
+                        # pass through Python num or datetime/date without stringification
+                        "min_limit": min_limit,
+                        "max_limit": max_limit,
                     },
                 },
                 "name": f"{column}_isnt_in_range",
                 "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",
-                    "arguments": {"column": column, "limit": val_maybe_to_str(max_limit, include_sql_quotes=False)},
-                },
+                "check": {"function": "is_not_greater_than", "arguments": {"column": column, "limit": max_limit}},
                 "name": f"{column}_not_greater_than",
                 "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",
-                    "arguments": {"column": column, "limit": val_maybe_to_str(min_limit, include_sql_quotes=False)},
-                },
+                "check": {"function": "is_not_less_than", "arguments": {"column": column, "limit": min_limit}},
                 "name": f"{column}_not_less_than",
                 "criticality": criticality,
             }

tests/integration/test_generator.py

@@ -1,3 +1,4 @@
+import logging
 import datetime
 from decimal import Decimal
 
@@ -38,8 +39,8 @@
 ]
 
 
-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 +72,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": datetime.date(2020, 1, 1)},
+            },
+            "name": "product_launch_date_not_less_than",
+            "criticality": "error",
+        },
+        {
+            "check": {
+                "function": "is_not_greater_than",
+                "arguments": {"column": "product_expiry_ts", "limit": datetime.datetime(2020, 1, 1)},
+            },
+            "name": "product_expiry_ts_not_greater_than",
+            "criticality": "error",
+        },
+        {
+            "check": {
+                "function": "is_in_range",
+                "arguments": {"column": "d1", "min_limit": Decimal('1.23'), "max_limit": Decimal('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 +133,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": datetime.date(2020, 1, 1)},
+            },
+            "name": "product_launch_date_not_less_than",
+            "criticality": "warn",
+        },
+        {
+            "check": {
+                "function": "is_not_greater_than",
+                "arguments": {"column": "product_expiry_ts", "limit": datetime.datetime(2020, 1, 1)},
+            },
+            "name": "product_expiry_ts_not_greater_than",
+            "criticality": "warn",
+        },
+        {
+            "check": {
+                "function": "is_in_range",
+                "arguments": {"column": "d1", "min_limit": Decimal('1.23'), "max_limit": Decimal('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 +253,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",