diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx index e448158d3..4e999e512 100644 --- a/docs/dqx/docs/reference/quality_checks.mdx +++ b/docs/dqx/docs/reference/quality_checks.mdx @@ -60,6 +60,7 @@ You can also define your own custom checks in Python (see [Creating custom check | `is_older_than_col2_for_n_days` | Checks whether the values in one input column are at least N days older than the values in another column. | `column1`: first column to check (can be a string column name or a column expression); `column2`: second column to check (can be a string column name or a column expression); `days`: number of days; `negate`: if the condition should be negated | | `regex_match` | Checks whether the values in the input column match a given regex. | `column`: column to check (can be a string column name or a column expression); regex: regex to check; `negate`: if the condition should be negated (true) or not | | `is_valid_email` | Checks whether the values in the input column have valid email address format. | `column`: column to check (can be a string column name or a column expression) | +| `is_valid_national_id` | Checks whether the values in the input column are valid national identification numbers (e.g., US Social Security Numbers) for the given country. | `column`: column to check (can be a string column name or a column expression); `country`: ISO 3166 alpha-2 country code (optional, default: `US`) | | `is_valid_ipv4_address` | Checks whether the values in the input column have valid IPv4 address format. | `column` to check (can be a string column name or a column expression) | | `is_ipv4_address_in_cidr` | Checks whether the values in the input column have valid IPv4 address format and fall within the given CIDR block. | `column`: column to check (can be a string column name or a column expression); `cidr_block`: CIDR block string | | `is_valid_ipv6_address` | Checks whether the values in the input column have valid IPv6 address format. | `column` to check (can be a string column name or a column expression) | @@ -532,6 +533,14 @@ For brevity, the `name` field in the examples is omitted and it will be auto-gen arguments: column: col1 +# is_valid_national_id check +- criticality: error + check: + function: is_valid_national_id + arguments: + column: col1 + country: US + # is_valid_ipv4_address check - criticality: error check: @@ -1279,6 +1288,16 @@ checks = [ column="col1" ), + # is_valid_national_id check + DQRowRule( + criticality="error", + check_func=check_funcs.is_valid_national_id, + column="col1", # or as expr: F.col("col1") + check_func_kwargs={ + "country": "US" + } + ), + # is_valid_ipv4_address check DQRowRule( criticality="error", diff --git a/src/databricks/labs/dqx/check_funcs.py b/src/databricks/labs/dqx/check_funcs.py index 62f32db78..87a09e788 100644 --- a/src/databricks/labs/dqx/check_funcs.py +++ b/src/databricks/labs/dqx/check_funcs.py @@ -96,6 +96,19 @@ class DQPattern(Enum): rf"$" ) + # US Social Security Number AAA-GG-SSSS: the separator (hyphen, single space, or + # none) must be consistent via backreference \1. Excludes invalid ranges - area + # 000/666/9xx (9xx covers ITINs), group 00, serial 0000. Anchored, fixed-width; ReDoS-safe. + SSN_US = r"^(?!000|666|9\d{2})\d{3}([- ]?)(?!00)\d{2}\1(?!0000)\d{4}$" + + +# ISO 3166 alpha-2 country code -> SSN / national-id validation pattern. Extension +# point for additional countries: add a new DQPattern member and map its ISO 3166 +# alpha-2 code here. +_NATIONAL_ID_PATTERNS_BY_COUNTRY: dict[str, DQPattern] = { + "US": DQPattern.SSN_US, +} + def make_condition(condition: Column, message: Column | str, alias: str) -> Column: """Helper function to create a condition column. @@ -1022,6 +1035,50 @@ def is_valid_email(column: str | Column) -> Column: return _matches_pattern(column, DQPattern.EMAIL_ADDRESS) +@register_rule("row") +def is_valid_national_id(column: str | Column, country: str = "US") -> Column: + """Checks whether the values in the input column are valid national identification + numbers (for example, US Social Security Numbers) for the given country. + + Validation is limited to *format* and *number ranges*; it does not verify that a + number was actually issued. + + Supported countries are keyed by ISO 3166 alpha-2 code. Currently only *US* is + supported: the *AAA-GG-SSSS* form is required, where the separators may be all + hyphens, all single spaces, or omitted entirely (e.g. *123-45-6789*, *123 45 6789* + or *123456789*), but must be used consistently. Structurally invalid ranges are + rejected (area *000*, *666* and *900-999* - the latter covering ITINs; group *00*; + serial *0000*). + + Null values will pass the check with no violation reported. + + Args: + column: column to check; can be a string column name or a column expression + country: ISO 3166 alpha-2 country code selecting the validation pattern (default: *US*) + + Returns: + Column object for condition + + Raises: + MissingParameterError: if *country* is None. + InvalidParameterError: if *country* is not a string, or is not a supported country code. + """ + if country is None: + raise MissingParameterError("'country' is not provided.") + + if not isinstance(country, str): + raise InvalidParameterError(f"'country' must be a string, got {type(country)} instead.") + + normalized_country = country.upper() + pattern = _NATIONAL_ID_PATTERNS_BY_COUNTRY.get(normalized_country) + if pattern is None: + supported = ", ".join(sorted(_NATIONAL_ID_PATTERNS_BY_COUNTRY)) + raise InvalidParameterError( + f"Unsupported country code for national ID validation: '{country}'. Supported: [{supported}]." + ) + return _matches_pattern(column, pattern) + + @register_rule("row") def is_ipv4_address_in_cidr(column: str | Column, cidr_block: str) -> Column: """ diff --git a/tests/integration/test_apply_checks.py b/tests/integration/test_apply_checks.py index a55209234..98a30a12f 100755 --- a/tests/integration/test_apply_checks.py +++ b/tests/integration/test_apply_checks.py @@ -5905,7 +5905,7 @@ def test_apply_checks_all_row_checks_as_yaml_with_streaming(ws, make_schema, mak "col1: string, col2: int, col3: int, col4 array, col5: date, col6: timestamp, " "col7: map, col8: struct, col10: int, col11: string, " "col_ipv4: string, col_ipv6: string, col_json_str: string, col_json_str2: string, " - "col_email: string" + "col_email: string, col_ssn: string" ) test_df = spark.createDataFrame( [ @@ -5925,6 +5925,7 @@ def test_apply_checks_all_row_checks_as_yaml_with_streaming(ws, make_schema, mak '{"key1": "1"}', '{"a" : 1, "b": 2}', "user@example.com", + "123-45-6789", ], [ "val2", @@ -5942,6 +5943,7 @@ def test_apply_checks_all_row_checks_as_yaml_with_streaming(ws, make_schema, mak '{"key1": "1", "key2": "2"}', '{ "a" : 1, "b": 1000, "c": {"1": 8}}', '"quoted_user"@example.co.uk', + "223-45-6789", ], [ "val3", @@ -5959,6 +5961,7 @@ def test_apply_checks_all_row_checks_as_yaml_with_streaming(ws, make_schema, mak '{"key1": "[1, 2, 3]"}', '{ "a" : 1, "b": 1023455, "c": null }', "user@[12.96.144.202]", + "323-45-6789", ], ], schema, @@ -6000,6 +6003,7 @@ def test_apply_checks_all_row_checks_as_yaml_with_streaming(ws, make_schema, mak '{"key1": "1"}', '{"a" : 1, "b": 2}', "user@example.com", + "123-45-6789", None, None, ], @@ -6019,6 +6023,7 @@ def test_apply_checks_all_row_checks_as_yaml_with_streaming(ws, make_schema, mak '{"key1": "1", "key2": "2"}', '{ "a" : 1, "b": 1000, "c": {"1": 8}}', '"quoted_user"@example.co.uk', + "223-45-6789", None, None, ], @@ -6038,6 +6043,7 @@ def test_apply_checks_all_row_checks_as_yaml_with_streaming(ws, make_schema, mak '{"key1": "[1, 2, 3]"}', '{ "a" : 1, "b": 1023455, "c": null }', "user@[12.96.144.202]", + "323-45-6789", None, None, ], @@ -6192,7 +6198,7 @@ def test_apply_checks_all_checks_as_yaml(ws, spark): "col1: string, col2: int, col3: int, col4 array, col5: date, col6: timestamp, " "col7: map, col8: struct, col10: int, col11: string, " "col_ipv4: string, col_ipv6: string, col_json_str: string, col_json_str2: string, " - "col_email: string" + "col_email: string, col_ssn: string" ) test_df = spark.createDataFrame( [ @@ -6212,6 +6218,7 @@ def test_apply_checks_all_checks_as_yaml(ws, spark): '{"key1": "1"}', '{"a" : 1, "b": 2}', "user@example.com", + "123-45-6789", ], [ "val2", @@ -6229,6 +6236,7 @@ def test_apply_checks_all_checks_as_yaml(ws, spark): '{"key1": "1", "key2": "2"}', '{ "a" : 1, "b": 1000, "c": {"1": 8}}', '"quoted_user"@example.co.uk', + "223-45-6789", ], [ "val3", @@ -6246,6 +6254,7 @@ def test_apply_checks_all_checks_as_yaml(ws, spark): '{"key1": "[1, 2, 3]"}', '{ "a" : 1, "b": 1023455, "c": null }', "user@[12.96.144.202]", + "323-45-6789", ], ], schema, @@ -6275,6 +6284,7 @@ def test_apply_checks_all_checks_as_yaml(ws, spark): '{"key1": "1"}', '{"a" : 1, "b": 2}', "user@example.com", + "123-45-6789", None, None, ], @@ -6294,6 +6304,7 @@ def test_apply_checks_all_checks_as_yaml(ws, spark): '{"key1": "1", "key2": "2"}', '{ "a" : 1, "b": 1000, "c": {"1": 8}}', '"quoted_user"@example.co.uk', + "223-45-6789", None, None, ], @@ -6313,6 +6324,7 @@ def test_apply_checks_all_checks_as_yaml(ws, spark): '{"key1": "[1, 2, 3]"}', '{ "a" : 1, "b": 1023455, "c": null }', "user@[12.96.144.202]", + "323-45-6789", None, None, ], @@ -7097,6 +7109,13 @@ def test_apply_checks_all_checks_using_classes(ws, spark): column="col_json_str2", check_func_kwargs={"schema": "STRUCT"}, ), + # is_valid_national_id check + DQRowRule( + criticality="error", + check_func=check_funcs.is_valid_national_id, + column="col_ssn", + check_func_kwargs={"country": "US"}, + ), ] dq_engine = DQEngine(ws) @@ -7104,7 +7123,7 @@ def test_apply_checks_all_checks_using_classes(ws, spark): schema = ( "col1: string, col2: int, col3: int, col4 array, col5: date, col6: timestamp, " "col7: map, col8: struct, col10: int, col11: string, " - "col_ipv4: string, col_ipv6: string, col_json_str: string, col_json_str2: string" + "col_ipv4: string, col_ipv6: string, col_json_str: string, col_json_str2: string, col_ssn: string" ) test_df = spark.createDataFrame( [ @@ -7123,6 +7142,7 @@ def test_apply_checks_all_checks_using_classes(ws, spark): "2001:0db8:85a3:08d3:1319:8a2e:0370:7344", '{"key1": "1"}', '{"a" : 1, "b": 2}', + "123-45-6789", ], [ "val2", @@ -7139,6 +7159,7 @@ def test_apply_checks_all_checks_using_classes(ws, spark): "2001:0db8:85a3:08d3:ffff:ffff:ffff:ffff", '{"key1": "1", "key2": "2"}', '{ "a" : 1, "b": 1000, "c": {"1": 8}}', + "223-45-6789", ], [ "val3", @@ -7155,6 +7176,7 @@ def test_apply_checks_all_checks_using_classes(ws, spark): "2001:db8:85a3:8d3:1319:8a2e:3.112.115.68", '{"key1": "[1, 2, 3]"}', '{ "a" : 1, "b": 1023455, "c": null }', + "323-45-6789", ], ], schema, @@ -7183,6 +7205,7 @@ def test_apply_checks_all_checks_using_classes(ws, spark): "2001:0db8:85a3:08d3:1319:8a2e:0370:7344", '{"key1": "1"}', '{"a" : 1, "b": 2}', + "123-45-6789", None, None, ], @@ -7201,6 +7224,7 @@ def test_apply_checks_all_checks_using_classes(ws, spark): "2001:0db8:85a3:08d3:ffff:ffff:ffff:ffff", '{"key1": "1", "key2": "2"}', '{ "a" : 1, "b": 1000, "c": {"1": 8}}', + "223-45-6789", None, None, ], @@ -7219,6 +7243,7 @@ def test_apply_checks_all_checks_using_classes(ws, spark): "2001:db8:85a3:8d3:1319:8a2e:3.112.115.68", '{"key1": "[1, 2, 3]"}', '{ "a" : 1, "b": 1023455, "c": null }', + "323-45-6789", None, None, ], diff --git a/tests/integration/test_row_checks.py b/tests/integration/test_row_checks.py index 2d80cbbe3..052589c43 100644 --- a/tests/integration/test_row_checks.py +++ b/tests/integration/test_row_checks.py @@ -32,6 +32,7 @@ is_valid_timestamp, is_valid_ipv4_address, is_valid_email, + is_valid_national_id, is_ipv4_address_in_cidr, is_valid_ipv6_address, is_ipv6_address_in_cidr, @@ -1888,7 +1889,7 @@ def test_col_is_valid_email(spark): ["user@[192.0.2.1"], # missing closing bracket ["user@localhost"], # no TLD ["missing@tld"], # no domain "." - [None], # null - passes (no violation reported) + [None], # Null - passes (no violation reported) ], schema_email, ) @@ -1952,6 +1953,84 @@ def violation(value: str) -> str: assertDataFrameEqual(actual, expected) +def test_col_is_valid_national_id(spark): + schema_ssn = "a: string" + test_df = spark.createDataFrame( + [ + # Valid - separators must be consistent (all '-', all ' ', or none) + ["123-45-6789"], + ["123456789"], + ["123 45 6789"], + ["899-45-6789"], # area boundary just below 900 + ["667-45-6789"], # area just above 666 + ["001-01-0001"], # minimal valid area / group / serial + # Invalid - excluded number ranges + ["000-45-6789"], # area 000 + ["666-45-6789"], # area 666 + ["900-45-6789"], # area 9xx (ITIN range, rejected) + ["123-00-6789"], # group 00 + ["123-45-0000"], # serial 0000 + # Invalid - separator / structure + ["123-45 6789"], # mixed separators + ["12-45-6789"], # area too short + ["1234-45-6789"], # area too long + ["abc-de-fghi"], # non-numeric + [""], # empty string + [None], # Null - passes (no violation reported) + ], + schema_ssn, + ) + + actual = test_df.select(is_valid_national_id("a", country="US")) + + def violation(value: str) -> str: + return f"Value '{value}' in Column 'a' does not match pattern 'SSN_US'" + + checked_schema = "a_does_not_match_pattern_ssn_us: string" + checked_data = [ + # Valid (no violation reported) + [None], + [None], + [None], + [None], + [None], + [None], + # Invalid - excluded number ranges + [violation("000-45-6789")], + [violation("666-45-6789")], + [violation("900-45-6789")], + [violation("123-00-6789")], + [violation("123-45-0000")], + # Invalid - separator / structure + [violation("123-45 6789")], + [violation("12-45-6789")], + [violation("1234-45-6789")], + [violation("abc-de-fghi")], + [violation("")], + # Null passes + [None], + ] + expected = spark.createDataFrame(checked_data, checked_schema) + + assertDataFrameEqual(actual, expected) + + +def test_col_is_valid_national_id_column_expr_and_lowercase_country(spark): + schema_ssn = "a: string" + test_df = spark.createDataFrame([["123-45-6789"], ["000-45-6789"]], schema_ssn) + + # Column-expression input + case-insensitive country normalization end-to-end + actual = test_df.select(is_valid_national_id(F.col("a"), country="us")) + + checked_schema = "a_does_not_match_pattern_ssn_us: string" + expected = spark.createDataFrame( + [[None], ["Value '000-45-6789' in Column 'a' does not match pattern 'SSN_US'"]], + checked_schema, + ) + + assertDataFrameEqual(actual, expected) + + def test_is_ipv4_address_in_cidr(spark): schema_ipv4 = "a: string, b: string" diff --git a/tests/perf/conftest.py b/tests/perf/conftest.py index cbfb5bff7..0e002d00e 100644 --- a/tests/perf/conftest.py +++ b/tests/perf/conftest.py @@ -319,3 +319,25 @@ def generated_email_df(spark): gen = gen.withColumnSpec(col, template=template) return gen.build() + + +@pytest.fixture +def generated_national_id_df(spark): + ssn_schema_str = ( + "col1_ssn_dashed: string, " "col2_ssn_plain: string, " "col3_ssn_valid_area: string, " "col4_ssn_spaced: string" + ) + schema = _parse_datatype_string(ssn_schema_str) + + ssn_templates = { + "col1_ssn_dashed": r"\n\n\n-\n\n-\n\n\n\n", + "col2_ssn_plain": r"\n\n\n\n\n\n\n\n\n", + "col3_ssn_valid_area": r"1\n\n-\n\n-\n\n\n\n", + "col4_ssn_spaced": r"\n\n\n \n\n \n\n\n\n", + } + + _, gen = make_data_gen(spark, n_rows=DEFAULT_ROWS, n_columns=len(ssn_templates), partitions=DEFAULT_PARTITIONS) + gen = gen.withSchema(schema) + for col, template in ssn_templates.items(): + gen = gen.withColumnSpec(col, template=template) + + return gen.build() diff --git a/tests/perf/test_apply_checks.py b/tests/perf/test_apply_checks.py index 532c22679..0c025f1b5 100644 --- a/tests/perf/test_apply_checks.py +++ b/tests/perf/test_apply_checks.py @@ -2208,3 +2208,29 @@ def test_benchmark_is_valid_email(benchmark, ws, generated_email_df, column): checked = dq_engine.apply_checks(generated_email_df, checks) actual_count = benchmark(lambda: checked.count()) assert actual_count == EXPECTED_ROWS + + +@pytest.mark.parametrize( + "column", + [ + "col1_ssn_dashed", + "col2_ssn_plain", + "col3_ssn_valid_area", + "col4_ssn_spaced", + ], +) +@pytest.mark.benchmark(group="test_benchmark_is_valid_national_id") +def test_benchmark_is_valid_national_id(benchmark, ws, generated_national_id_df, column): + dq_engine = DQEngine(workspace_client=ws, extra_params=EXTRA_PARAMS) + checks = [ + DQRowRule( + name=f"{column}_is_valid_national_id", + criticality="warn", + check_func=check_funcs.is_valid_national_id, + column=column, + ), + ] + benchmark.group += f" {column}" + checked = dq_engine.apply_checks(generated_national_id_df, checks) + actual_count = benchmark(lambda: checked.count()) + assert actual_count == EXPECTED_ROWS diff --git a/tests/resources/all_row_checks.yaml b/tests/resources/all_row_checks.yaml index c9ea9dd99..04984d132 100644 --- a/tests/resources/all_row_checks.yaml +++ b/tests/resources/all_row_checks.yaml @@ -483,3 +483,11 @@ function: is_valid_email arguments: column: col_email + +# is_valid_national_id check +- criticality: error + check: + function: is_valid_national_id + arguments: + column: col_ssn + country: US diff --git a/tests/unit/test_row_checks.py b/tests/unit/test_row_checks.py index 8e46caf9f..87ebd9c5c 100644 --- a/tests/unit/test_row_checks.py +++ b/tests/unit/test_row_checks.py @@ -12,6 +12,7 @@ is_aggr_not_greater_than, is_ipv4_address_in_cidr, is_ipv6_address_in_cidr, + is_valid_national_id, sql_expression, ) from databricks.labs.dqx.pii.pii_detection_funcs import does_not_contain_pii @@ -141,3 +142,33 @@ def test_sql_expression_complex_exists_negate_auto_name(): expression = "EXISTS (SELECT 1 FROM cfg WHERE cfg.val = STATUS)" result = sql_expression(expression, negate=True) assert get_column_name_or_alias(result) == "exists_select_1_from_cfg_where_cfg_val_status" + + +def test_is_valid_national_id_default_country_auto_name(): + result = is_valid_national_id("a") + assert get_column_name_or_alias(result) == "a_does_not_match_pattern_ssn_us" + + +def test_is_valid_national_id_explicit_us_auto_name(): + result = is_valid_national_id("a", country="US") + assert get_column_name_or_alias(result) == "a_does_not_match_pattern_ssn_us" + + +def test_is_valid_national_id_country_is_case_insensitive(): + result = is_valid_national_id("a", country="us") + assert get_column_name_or_alias(result) == "a_does_not_match_pattern_ssn_us" + + +def test_is_valid_national_id_missing_country(): + with pytest.raises(MissingParameterError, match="'country' is not provided."): + is_valid_national_id("a", country=None) + + +def test_is_valid_national_id_non_string_country(): + with pytest.raises(InvalidParameterError, match="'country' must be a string"): + is_valid_national_id("a", country=123) + + +def test_is_valid_national_id_unsupported_country(): + with pytest.raises(InvalidParameterError, match="Unsupported country code for national ID validation"): + is_valid_national_id("a", country="ZZ")