Skip to content

Add is_valid_national_id row-level check #1346

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mwojtyczka merged 5 commits into from
Jul 24, 2026
+271 −4
Merged

Add is_valid_national_id row-level check #1346

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
9f52a5d
Add is_valid_national_id row-level check
SreeramaYeshwanthGowd Jul 19, 2026
b813907
Add performance benchmark for is_valid_national_id
SreeramaYeshwanthGowd Jul 20, 2026
8378f7e
Add is_valid_national_id to test_apply_checks_all_checks_using_classes
SreeramaYeshwanthGowd Jul 22, 2026
39885b7
Merge branch 'main' into feature/1265-is-valid-national-id
mwojtyczka Jul 22, 2026
c71780c
Merge branch 'main' into feature/1265-is-valid-national-id
mwojtyczka Jul 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions docs/dqx/docs/reference/quality_checks.mdx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) |
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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",
Expand Down
57 changes: 57 additions & 0 deletions src/databricks/labs/dqx/check_funcs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,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.
Expand Down Expand Up @@ -1017,6 +1030,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:
"""
Expand Down
16 changes: 14 additions & 2 deletions tests/integration/test_apply_checks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<int>, col5: date, col6: timestamp, "
"col7: map<string, int>, col8: struct<field1: int>, 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(
[
Expand All @@ -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",
Expand All @@ -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",
Expand All @@ -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,
Expand Down Expand Up @@ -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,
],
Expand All @@ -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,
],
Expand All @@ -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,
],
Expand Down Expand Up @@ -6192,7 +6198,7 @@ def test_apply_checks_all_checks_as_yaml(ws, spark):
"col1: string, col2: int, col3: int, col4 array<int>, col5: date, col6: timestamp, "
"col7: map<string, int>, col8: struct<field1: int>, 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(
[
Expand All @@ -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",
Expand All @@ -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",
Expand All @@ -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,
Expand Down Expand Up @@ -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,
],
Expand All @@ -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,
],
Expand All @@ -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,
],
Expand Down
81 changes: 80 additions & 1 deletion tests/integration/test_row_checks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
)
Expand Down Expand Up @@ -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"

Expand Down
22 changes: 22 additions & 0 deletions tests/perf/conftest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Loading