Skip to content

Commit 66a05a8

Browse files
Jgprog117Jgprog117
committed
Add support for Decimal to checks and change to single-line function calls
1 parent 0d63ce8 commit 66a05a8

3 files changed

Lines changed: 33 additions & 39 deletions

File tree

‎src/databricks/labs/dqx/check_funcs.py‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import warnings
44
import ipaddress
55
import uuid
6+
from decimal import Decimal
67
from collections.abc import Callable, Sequence
78
from enum import Enum
89
from itertools import zip_longest
@@ -609,7 +610,7 @@ def is_not_in_near_future(column: str | Column, offset: int = 0, curr_timestamp:
609610
@register_rule("row")
610611
def is_equal_to(
611612
column: str | Column,
612-
value: int | float | str | datetime.date | datetime.datetime | Column | None = None,
613+
value: int | float | Decimal | str | datetime.date | datetime.datetime | Column | None = None,
613614
abs_tolerance: float | None = None,
614615
rel_tolerance: float | None = None,
615616
) -> Column:
@@ -671,7 +672,7 @@ def is_equal_to(
671672
@register_rule("row")
672673
def is_not_equal_to(
673674
column: str | Column,
674-
value: int | float | str | datetime.date | datetime.datetime | Column | None = None,
675+
value: int | float | Decimal | str | datetime.date | datetime.datetime | Column | None = None,
675676
abs_tolerance: float | None = None,
676677
rel_tolerance: float | None = None,
677678
) -> Column:
@@ -733,7 +734,7 @@ def is_not_equal_to(
733734

734735
@register_rule("row")
735736
def is_not_less_than(
736-
column: str | Column, limit: int | float | datetime.date | datetime.datetime | str | Column | None = None
737+
column: str | Column, limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None
737738
) -> Column:
738739
"""Checks whether the values in the input column are not less than the provided limit.
739740
@@ -763,7 +764,7 @@ def is_not_less_than(
763764

764765
@register_rule("row")
765766
def is_not_greater_than(
766-
column: str | Column, limit: int | float | datetime.date | datetime.datetime | str | Column | None = None
767+
column: str | Column, limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None
767768
) -> Column:
768769
"""Checks whether the values in the input column are not greater than the provided limit.
769770
@@ -794,8 +795,8 @@ def is_not_greater_than(
794795
@register_rule("row")
795796
def is_in_range(
796797
column: str | Column,
797-
min_limit: int | float | datetime.date | datetime.datetime | str | Column | None = None,
798-
max_limit: int | float | datetime.date | datetime.datetime | str | Column | None = None,
798+
min_limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None,
799+
max_limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None,
799800
) -> Column:
800801
"""Checks whether the values in the input column are in the provided limits (inclusive of both boundaries).
801802
@@ -832,8 +833,8 @@ def is_in_range(
832833
@register_rule("row")
833834
def is_not_in_range(
834835
column: str | Column,
835-
min_limit: int | float | datetime.date | datetime.datetime | str | Column | None = None,
836-
max_limit: int | float | datetime.date | datetime.datetime | str | Column | None = None,
836+
min_limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None,
837+
max_limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None,
837838
) -> Column:
838839
"""Checks whether the values in the input column are outside the provided limits (inclusive of both boundaries).
839840
@@ -1585,7 +1586,7 @@ def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) ->
15851586
@register_rule("dataset")
15861587
def is_aggr_not_greater_than(
15871588
column: str | Column,
1588-
limit: int | float | str | Column,
1589+
limit: int | float | Decimal | str | Column,
15891590
aggr_type: str = "count",
15901591
group_by: list[str | Column] | None = None,
15911592
row_filter: str | None = None,
@@ -1630,7 +1631,7 @@ def is_aggr_not_greater_than(
16301631
@register_rule("dataset")
16311632
def is_aggr_not_less_than(
16321633
column: str | Column,
1633-
limit: int | float | str | Column,
1634+
limit: int | float | Decimal | str | Column,
16341635
aggr_type: str = "count",
16351636
group_by: list[str | Column] | None = None,
16361637
row_filter: str | None = None,
@@ -1675,7 +1676,7 @@ def is_aggr_not_less_than(
16751676
@register_rule("dataset")
16761677
def is_aggr_equal(
16771678
column: str | Column,
1678-
limit: int | float | str | Column,
1679+
limit: int | float | Decimal | str | Column,
16791680
aggr_type: str = "count",
16801681
group_by: list[str | Column] | None = None,
16811682
row_filter: str | None = None,
@@ -1726,7 +1727,7 @@ def is_aggr_equal(
17261727
@register_rule("dataset")
17271728
def is_aggr_not_equal(
17281729
column: str | Column,
1729-
limit: int | float | str | Column,
1730+
limit: int | float | Decimal | str | Column,
17301731
aggr_type: str = "count",
17311732
group_by: list[str | Column] | None = None,
17321733
row_filter: str | None = None,
@@ -2780,7 +2781,6 @@ def _add_column_diffs(
27802781
"""
27812782
columns_changed = []
27822783
if compare_columns:
2783-
27842784
for col_name in compare_columns:
27852785
is_numeric = isinstance(df.schema[col_name].dataType, types.NumericType)
27862786

@@ -2991,7 +2991,7 @@ def _build_aggregate_check_metadata(
29912991

29922992
def _is_aggr_compare(
29932993
column: str | Column,
2994-
limit: int | float | str | Column,
2994+
limit: int | float | Decimal | str | Column,
29952995
aggr_type: str,
29962996
aggr_params: dict[str, Any] | None,
29972997
group_by: list[str | Column] | None,
@@ -3218,7 +3218,7 @@ def _cleanup_alias_name(column: str) -> str:
32183218

32193219

32203220
def get_limit_expr(
3221-
limit: int | float | datetime.date | datetime.datetime | str | Column | None = None,
3221+
limit: int | float | Decimal | datetime.date | datetime.datetime | str | Column | None = None,
32223222
) -> Column:
32233223
"""
32243224
Generate a Spark Column expression for a limit value.
@@ -3247,7 +3247,7 @@ def get_limit_expr(
32473247
parsed_dt = datetime.datetime.fromisoformat(limit)
32483248

32493249
# Check if the string contains time component
3250-
has_time = ':' in limit
3250+
has_time = ":" in limit
32513251

32523252
if has_time:
32533253
return F.to_timestamp(F.lit(parsed_dt))

‎tests/integration/test_generator.py‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -31,16 +31,16 @@
3131
),
3232
DQProfile(name="is_random", column="vendor_id", parameters={"in": ["1", "4", "2"]}),
3333
DQProfile(
34-
name='min_max',
35-
column='d1',
36-
description='Real min/max values were used',
37-
parameters={'max': 333323.00, 'min': 1.23},
34+
name="min_max",
35+
column="d1",
36+
description="Real min/max values were used",
37+
parameters={"max": 333323.00, "min": 1.23},
3838
),
3939
DQProfile(
40-
name='min_max',
41-
column='price',
42-
description='Real min/max values were used',
43-
parameters={'min': Decimal("0.01"), 'max': Decimal("999.99")},
40+
name="min_max",
41+
column="price",
42+
description="Real min/max values were used",
43+
parameters={"min": Decimal("0.01"), "max": Decimal("999.99")},
4444
),
4545
]
4646

@@ -332,9 +332,9 @@ def test_generate_is_unique_dq_rule(ws, spark):
332332
generator = DQGenerator(ws, spark)
333333
test_is_unique_rules = [
334334
DQProfile(
335-
name='is_unique',
336-
column='col1,col2',
337-
description='LLM-detected primary key columns: col1, col2',
335+
name="is_unique",
336+
column="col1,col2",
337+
description="LLM-detected primary key columns: col1, col2",
338338
parameters={"nulls_distinct": False, "confidence": "high"},
339339
),
340340
]
@@ -355,9 +355,9 @@ def test_generate_is_unique_dq_rule_default_criticality(ws, spark):
355355
generator = DQGenerator(ws, spark)
356356
test_is_unique_rules = [
357357
DQProfile(
358-
name='is_unique',
359-
column='col1',
360-
description='LLM-detected primary key columns: col1, col2',
358+
name="is_unique",
359+
column="col1",
360+
description="LLM-detected primary key columns: col1, col2",
361361
parameters={"nulls_distinct": True, "confidence": "low"},
362362
),
363363
]

‎tests/unit/test_generator_numeric.py‎

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,7 @@
55

66
def test_decimal_both_bounds_is_in_range():
77
"""Test min_max generator with Decimal type for both bounds."""
8-
result = DQGenerator.dq_generate_min_max(
9-
"price_col", **{"min": Decimal("0.01"), "max": Decimal("999.99")}
10-
)
8+
result = DQGenerator.dq_generate_min_max("price_col", **{"min": Decimal("0.01"), "max": Decimal("999.99")})
119
assert result["check"]["function"] == "is_in_range"
1210
args = result["check"]["arguments"]
1311
assert args["column"] == "price_col"
@@ -17,9 +15,7 @@ def test_decimal_both_bounds_is_in_range():
1715

1816
def test_decimal_only_min_is_not_less_than():
1917
"""Test min_max generator with Decimal type for minimum bound only."""
20-
result = DQGenerator.dq_generate_min_max(
21-
"amount_col", **{"min": Decimal("10.50"), "max": None}
22-
)
18+
result = DQGenerator.dq_generate_min_max("amount_col", **{"min": Decimal("10.50"), "max": None})
2319
assert result["check"]["function"] == "is_not_less_than"
2420
args = result["check"]["arguments"]
2521
assert args["column"] == "amount_col"
@@ -57,9 +53,7 @@ def test_float_both_bounds_is_in_range():
5753

5854
def test_mixed_int_and_decimal_is_in_range():
5955
"""Test that mixing int and Decimal produces is_in_range since both are numeric."""
60-
result = DQGenerator.dq_generate_min_max(
61-
"mixed_col", **{"min": 10, "max": Decimal("100.00")}
62-
)
56+
result = DQGenerator.dq_generate_min_max("mixed_col", **{"min": 10, "max": Decimal("100.00")})
6357
assert result is not None
6458
assert result["check"]["function"] == "is_in_range"
6559
args = result["check"]["arguments"]

0 commit comments

Comments
 (0)