diff --git a/packages/overture-schema-pyspark/tests/expressions/test_column_patterns.py b/packages/overture-schema-pyspark/tests/expressions/test_column_patterns.py index 3753c9f29..a43be04a6 100644 --- a/packages/overture-schema-pyspark/tests/expressions/test_column_patterns.py +++ b/packages/overture-schema-pyspark/tests/expressions/test_column_patterns.py @@ -1,5 +1,20 @@ -"""Tests for column_patterns — structural PySpark composition helpers.""" - +"""Tests for column_patterns -- structural PySpark composition helpers. + +Each helper composition is exercised as a `_Case`: a uniquely-named input +column, the check built over it, and a predicate on the collected result. The +`results` fixture packs every case's input into one wide single-row DataFrame, +applies every check in one `select`, and collects once -- so the whole file +pays for a single `createDataFrame` + `collect` instead of one pair per test +(the same batch-once pattern the generated conformance harness uses). Cases +needing both a violating and a clean input carry two entries (`*_invalid` / +`*_valid`). +""" + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import pytest from overture.schema.pyspark.expressions.column_patterns import ( array_check, check_struct_unique, @@ -16,497 +31,411 @@ check_require_any_of, check_string_min_length, ) -from pyspark.sql import Row, SparkSession +from pyspark.sql import Column, SparkSession from pyspark.sql import functions as F -from pyspark.sql.types import ( - ArrayType, - IntegerType, - MapType, - StringType, - StructField, - StructType, -) - - -def test_error_msg_concatenates(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="bad")]) - result = df.select(error_msg("field: got ", F.col("val")).alias("msg")).collect() - assert result[0]["msg"] == "field: got bad" - - -def test_error_msg_multiple_values(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(a="x", b="y")]) - result = df.select( - error_msg("prefix ", F.col("a"), F.lit(" and "), F.col("b")).alias("msg") - ).collect() - assert result[0]["msg"] == "prefix x and y" - - -def test_error_msg_null_value_does_not_nullify_message(spark: SparkSession) -> None: - # A NULL interpolated value must not make the whole message NULL: F.concat - # would, and a NULL message is dropped by array_compact, silently swallowing - # the violation (e.g. an out-of-bounds linear-reference range [null, 1.5]). - # The null must render as a literal instead. - df = spark.createDataFrame([Row(val=None)], schema="val double") - result = df.select(error_msg("got ", F.col("val")).alias("msg")).collect() - assert result[0]["msg"] == "got null" - - -def test_array_check_null_column_returns_null(spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(items=None)], - schema="items array>", - ) - result = df.select( - array_check("items", lambda el: F.lit("err")).alias("errs") - ).collect() - assert result[0]["errs"] is None - - -def test_array_check_filters_nulls(spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(items=[Row(val="ok"), Row(val="bad")])], - schema="items array>", - ) - result = df.select( - array_check( - "items", - lambda el: F.when(el["val"] == "bad", F.lit("error")), - ).alias("errs") - ).collect() - assert result[0]["errs"] == ["error"] +# PySpark 3.4's collect() leaves its result socket for the GC to finalize; under +# -W error that ResourceWarning fails the batched `results` fixture. conftest's +# unraisablehook catches the finalizer path, but this fixture emits it +# synchronously -- and a filterwarnings mark is the only filter outranking the +# command-line -W error. +pytestmark = pytest.mark.filterwarnings("ignore::ResourceWarning") -def test_array_check_empty_when_all_valid(spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(items=[Row(val="ok")])], - schema="items array>", - ) - result = df.select( - array_check( - "items", - lambda el: F.when(el["val"] == "bad", F.lit("error")), - ).alias("errs") - ).collect() - assert result[0]["errs"] == [] +@dataclass(frozen=True) +class _Case: + """One column_patterns assertion driven off the shared wide row. -def test_struct_unique_no_duplicates(spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(items=[Row(id="a"), Row(id="b")])], - schema="items array>", - ) - result = df.select(check_struct_unique("items").alias("err")).collect() - assert result[0]["err"] is None - - -def test_struct_unique_with_duplicates(spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(items=[Row(id="a"), Row(id="a")])], - schema="items array>", - ) - result = df.select(check_struct_unique("items").alias("err")).collect() - assert result[0]["err"] is not None - assert "duplicate" in result[0]["err"] - - -def test_struct_unique_null_column(spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(items=None)], - schema="items array>", - ) - result = df.select(check_struct_unique("items").alias("err")).collect() - assert result[0]["err"] is None - + `check(field)` builds the composition over the case's input column; + `expect(result)` is a predicate on that column's single collected value. + """ -def test_struct_unique_repeated_value_different_fields(spark: SparkSession) -> None: - """Structs with same value subfield but different other fields are unique.""" - df = spark.createDataFrame( + id: str + ddl: str + value: Any + check: Callable[[str], Column] + expect: Callable[[Any], bool] + + +_CASES: list[_Case] = [ + # --- error_msg: builds a string message, not an error array --------------- + _Case( + "em_concat", + "string", + "bad", + lambda f: error_msg("field: got ", F.col(f)), + lambda r: r == "field: got bad", + ), + _Case( + "em_multi", + "struct", + {"a": "x", "b": "y"}, + lambda f: error_msg("prefix ", F.col(f)["a"], F.lit(" and "), F.col(f)["b"]), + lambda r: r == "prefix x and y", + ), + # A NULL interpolated value must not nullify the whole message (F.concat + # would): a NULL message is dropped by array_compact, silently swallowing the + # violation (e.g. an out-of-bounds range [null, 1.5]). It renders literally. + _Case( + "em_null", + "double", + None, + lambda f: error_msg("got ", F.col(f)), + lambda r: r == "got null", + ), + # --- array_check ---------------------------------------------------------- + _Case( + "ac_null", + "array>", + None, + lambda f: array_check(f, lambda el: F.lit("err")), + lambda r: r is None, + ), + _Case( + "ac_filter", + "array>", + [{"val": "ok"}, {"val": "bad"}], + lambda f: array_check(f, lambda el: F.when(el["val"] == "bad", F.lit("error"))), + lambda r: r == ["error"], + ), + _Case( + "ac_empty", + "array>", + [{"val": "ok"}], + lambda f: array_check(f, lambda el: F.when(el["val"] == "bad", F.lit("error"))), + lambda r: r == [], + ), + # array_check / check_struct_unique also accept a Column, not just a name. + _Case( + "ac_col", + "array>", + [{"val": "ok"}, {"val": "bad"}], + lambda f: array_check( + F.col(f), lambda el: F.when(el["val"] == "bad", F.lit("error")) + ), + lambda r: r == ["error"], + ), + # --- check_struct_unique -------------------------------------------------- + _Case( + "su_nodup", + "array>", + [{"id": "a"}, {"id": "b"}], + lambda f: check_struct_unique(f), + lambda r: r is None, + ), + _Case( + "su_dup", + "array>", + [{"id": "a"}, {"id": "a"}], + lambda f: check_struct_unique(f), + lambda r: r is not None and "duplicate" in r, + ), + _Case( + "su_null", + "array>", + None, + lambda f: check_struct_unique(f), + lambda r: r is None, + ), + # Same value subfield but different other fields is not a duplicate. + _Case( + "su_repeat", + "array>", [ - Row( - items=[ - Row(value="a", pos=0.0), - Row(value="b", pos=0.5), - Row(value="a", pos=0.7), - ] + {"value": "a", "pos": 0.0}, + {"value": "b", "pos": 0.5}, + {"value": "a", "pos": 0.7}, + ], + lambda f: check_struct_unique(f), + lambda r: r is None, + ), + _Case( + "su_single", + "array>", + [{"id": "a"}], + lambda f: check_struct_unique(f), + lambda r: r is None, + ), + _Case( + "csu_col", + "array>", + [{"id": "a"}, {"id": "a"}], + lambda f: check_struct_unique(F.col(f)), + lambda r: r is not None and "duplicate" in r, + ), + _Case( + "csu_colnull", + "array>", + None, + lambda f: check_struct_unique(F.col(f)), + lambda r: r is None, + ), + # --- nested_array_check --------------------------------------------------- + _Case( + "na_flat", + "array>>", + [{"tags": ["good", "bad"]}, {"tags": ["worse"]}], + lambda f: coalesce_errors( + nested_array_check( + f, + lambda el: array_check( + el["tags"], + lambda tag: F.when(tag != "good", F.concat(F.lit("bad: "), tag)), + ), ) - ] - ) - result = df.select(check_struct_unique("items").alias("err")).collect() - assert result[0]["err"] is None - - -def test_struct_unique_single_element(spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(items=[Row(id="a")])], - schema="items array>", - ) - result = df.select(check_struct_unique("items").alias("err")).collect() - assert result[0]["err"] is None - - -def test_array_check_accepts_column(spark: SparkSession) -> None: - """array_check works when passed a Column instead of a string name.""" - df = spark.createDataFrame( - [Row(items=[Row(val="ok"), Row(val="bad")])], - schema="items array>", - ) - result = df.select( - array_check( - F.col("items"), - lambda el: F.when(el["val"] == "bad", F.lit("error")), - ).alias("errs") - ).collect() - assert result[0]["errs"] == ["error"] - - -def test_check_struct_unique_accepts_column(spark: SparkSession) -> None: - """check_struct_unique works when passed a Column instead of a string name.""" - df = spark.createDataFrame( - [Row(items=[Row(id="a"), Row(id="a")])], - schema="items array>", - ) - result = df.select(check_struct_unique(F.col("items")).alias("err")).collect() - assert result[0]["err"] is not None - assert "duplicate" in result[0]["err"] - - -def test_check_struct_unique_column_null(spark: SparkSession) -> None: - """check_struct_unique with Column input handles null.""" - df = spark.createDataFrame( - [Row(items=None)], schema="items array>" - ) - result = df.select(check_struct_unique(F.col("items")).alias("err")).collect() - assert result[0]["err"] is None - - -def test_nested_array_check_flattens(spark: SparkSession) -> None: - """Inner array_check per outer element produces flat error list.""" - schema = "items array>>" - df = spark.createDataFrame( - [ - Row( - items=[ - Row(tags=["good", "bad"]), - Row(tags=["worse"]), - ] + ), + lambda r: len(r) == 2 and all(isinstance(e, str) for e in r), + ), + _Case( + "na_null", + "array>>", + None, + lambda f: coalesce_errors( + nested_array_check( + f, + lambda el: array_check( + el["tags"], lambda tag: F.when(tag != "good", F.lit("bad")) + ), ) - ], - schema=schema, - ) - result_col = nested_array_check( - "items", - lambda el: array_check( - el["tags"], - lambda tag: F.when(tag != "good", F.concat(F.lit("bad: "), tag)), ), - ) - result = df.select(coalesce_errors(result_col).alias("errs")).collect() - errors = result[0]["errs"] - assert len(errors) == 2 - assert all(isinstance(e, str) for e in errors) - - -def test_nested_array_check_null_outer(spark: SparkSession) -> None: - schema = "items array>>" - df = spark.createDataFrame([Row(items=None)], schema=schema) - result_col = nested_array_check( - "items", - lambda el: array_check( - el["tags"], - lambda tag: F.when(tag != "good", F.lit("bad")), + lambda r: r == [], + ), + # A null inner array must not nullify sibling errors during flatten: + # F.flatten returns NULL whenever any sub-array is NULL, which would drop + # every sibling error unless inner nulls are guarded. + _Case( + "na_mixed", + "array>>", + [{"tags": ["good"]}, {"tags": None}, {"tags": ["bad"]}], + lambda f: coalesce_errors( + nested_array_check( + f, + lambda el: array_check( + el["tags"], + lambda tag: F.when(tag != "good", F.concat(F.lit("bad: "), tag)), + ), + ) ), - ) - result = df.select(coalesce_errors(result_col).alias("errs")).collect() - assert result[0]["errs"] == [] - - -def test_nested_array_check_mixed_null_inner_with_sibling_errors( - spark: SparkSession, -) -> None: - """A null inner array must not nullify sibling errors during flatten. - - `F.flatten` returns NULL whenever any sub-array is NULL. Without - guarding inner nulls, the outer transform produces NULL and every - sibling error is silently dropped. - """ - schema = "items array>>" - df = spark.createDataFrame( - [ - Row( - items=[ - Row(tags=["good"]), - Row(tags=None), - Row(tags=["bad"]), - ] + lambda r: r == ["bad: bad"], + ), + _Case( + "na_noerr", + "array>>", + [{"tags": ["good"]}], + lambda f: coalesce_errors( + nested_array_check( + f, + lambda el: array_check( + el["tags"], lambda tag: F.when(tag != "good", F.lit("bad")) + ), ) - ], - schema=schema, - ) - result_col = nested_array_check( - "items", - lambda el: array_check( - el["tags"], - lambda tag: F.when(tag != "good", F.concat(F.lit("bad: "), tag)), ), - ) - result = df.select(coalesce_errors(result_col).alias("errs")).collect() - assert result[0]["errs"] == ["bad: bad"] - - -def test_nested_array_check_no_errors(spark: SparkSession) -> None: - schema = "items array>>" - df = spark.createDataFrame( - [Row(items=[Row(tags=["good"])])], - schema=schema, - ) - result_col = nested_array_check( - "items", - lambda el: array_check( - el["tags"], - lambda tag: F.when(tag != "good", F.lit("bad")), + lambda r: r == [], + ), + # --- map_keys_check / map_values_check ------------------------------------ + _Case( + "mk_bad", + "map", + {"good": "v", "bad": "v"}, + lambda f: map_keys_check(f, lambda k: F.when(k == "bad", F.lit("bad key"))), + lambda r: r == ["bad key"], + ), + _Case( + "mv_bad", + "map", + {"a": "ok", "b": "bad"}, + lambda f: map_values_check(f, lambda v: F.when(v == "bad", F.lit("bad value"))), + lambda r: r == ["bad value"], + ), + # A field check on a dict[str, Model] value navigates into the value struct + # -- the exact composition the renderer emits for a map-of-model value field. + _Case( + "mv_struct", + "map>", + {"a": {"label": ""}}, + lambda f: map_values_check(f, lambda v: check_string_min_length(v["label"], 1)), + lambda r: r == ["minimum length 1, got 0"], + ), + _Case( + "mv_struct_ok", + "map>", + {"a": {"label": "ok"}}, + lambda f: map_values_check(f, lambda v: check_string_min_length(v["label"], 1)), + lambda r: r == [], + ), + # A model-level constraint on a dict[str, Model] value -- the composition the + # renderer emits for a map-of-model value-model constraint. + _Case( + "mv_model", + "map>", + {"a": {"foo": None, "bar": None}}, + lambda f: map_values_check( + f, lambda v: check_require_any_of([v["foo"], v["bar"]], ["foo", "bar"]) ), - ) - result = df.select(coalesce_errors(result_col).alias("errs")).collect() - assert result[0]["errs"] == [] - - -def test_map_keys_check_flags_bad_key(spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(tags={"good": "v", "bad": "v"})], - schema="tags map", - ) - result = df.select( - map_keys_check("tags", lambda k: F.when(k == "bad", F.lit("bad key"))).alias( - "errs" - ) - ).collect() - assert result[0]["errs"] == ["bad key"] - - -def test_map_values_check_flags_bad_value(spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(tags={"a": "ok", "b": "bad"})], - schema="tags map", - ) - result = df.select( - map_values_check( - "tags", lambda v: F.when(v == "bad", F.lit("bad value")) - ).alias("errs") - ).collect() - assert result[0]["errs"] == ["bad value"] - - -def test_map_values_check_descends_into_value_struct_field( - spark: SparkSession, -) -> None: - # A field check on a `dict[str, Model]` value navigates into the value - # struct: map_values_check over a struct-navigating lambda, the exact - # composition the renderer emits for a map-of-model value field. - df = spark.createDataFrame( - [Row(items={"a": Row(label="")})], - schema="items map>", - ) - result = df.select( - map_values_check( - "items", lambda v: check_string_min_length(v["label"], 1) - ).alias("errs") - ).collect() - assert result[0]["errs"] == ["minimum length 1, got 0"] - - -def test_map_values_check_passes_valid_value_struct_field( - spark: SparkSession, -) -> None: - df = spark.createDataFrame( - [Row(items={"a": Row(label="ok")})], - schema="items map>", - ) - result = df.select( - map_values_check( - "items", lambda v: check_string_min_length(v["label"], 1) - ).alias("errs") - ).collect() - assert result[0]["errs"] == [] - - -def test_map_values_check_enforces_value_model_constraint( - spark: SparkSession, -) -> None: - # A model-level constraint on a `dict[str, Model]` value: map_values_check - # wrapping check_require_any_of over the value struct's fields, the exact - # composition the renderer emits for a map-of-model value-model constraint. - df = spark.createDataFrame( - [Row(subs={"a": Row(foo=None, bar=None)})], - schema="subs map>", - ) - result = df.select( - map_values_check( - "subs", - lambda v: check_require_any_of([v["foo"], v["bar"]], ["foo", "bar"]), - ).alias("errs") - ).collect() - assert result[0]["errs"] == ["requires at least one of foo, bar"] - - -def test_map_values_check_passes_satisfied_value_model_constraint( - spark: SparkSession, -) -> None: - df = spark.createDataFrame( - [Row(subs={"a": Row(foo=1, bar=None)})], - schema="subs map>", - ) - result = df.select( - map_values_check( - "subs", - lambda v: check_require_any_of([v["foo"], v["bar"]], ["foo", "bar"]), - ).alias("errs") - ).collect() - assert result[0]["errs"] == [] - - -def test_map_keys_check_null_column_returns_null(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(tags=None)], schema="tags map") - result = df.select( - map_keys_check("tags", lambda k: F.lit("err")).alias("errs") - ).collect() - assert result[0]["errs"] is None - - -def test_map_values_check_all_valid_empty(spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(tags={"a": "ok"})], schema="tags map" - ) - result = df.select( - map_values_check("tags", lambda v: F.when(v == "bad", F.lit("err"))).alias( - "errs" - ) - ).collect() - assert result[0]["errs"] == [] - - -def test_map_keys_check_accepts_column(spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(tags={"bad": "v"})], schema="tags map" - ) - result = df.select( - map_keys_check(F.col("tags"), lambda k: F.when(k == "bad", F.lit("err"))).alias( - "errs" - ) - ).collect() - assert result[0]["errs"] == ["err"] - - -def test_coalesce_errors_null_becomes_empty(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(x=1)]) - result = df.select( - coalesce_errors(F.lit(None).cast("array")).alias("errs") - ).collect() - assert result[0]["errs"] == [] - - -def test_coalesce_errors_preserves_array(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(x=1)]) - result = df.select(coalesce_errors(F.array(F.lit("err"))).alias("errs")).collect() - assert result[0]["errs"] == ["err"] - - -def test_nested_map_values_check_flattens_inner_array_errors( - spark: SparkSession, -) -> None: - # column `m`: map>; flag any element == 0 in any value list. - schema = StructType( - [StructField("m", MapType(StringType(), ArrayType(IntegerType())))] - ) - df = spark.createDataFrame( - [({"a": [1, 0], "b": [2]},), ({"a": [1], "b": [3]},)], - schema, - ) - check = nested_map_values_check( - "m", - lambda v: array_check(v, lambda e: F.when(e == 0, F.lit("zero not allowed"))), - ) - rows = df.select(check.alias("errors")).collect() - # Row 0 has one zero -> one flattened error; row 1 has none -> empty. - assert rows[0]["errors"] == ["zero not allowed"] - assert rows[1]["errors"] == [] - - -def test_nested_map_keys_check_flattens_inner_array_errors( - spark: SparkSession, -) -> None: - # column `m`: map; split each key into char array and flag 'x'. - df = spark.createDataFrame( - [({"ax": 1, "b": 2},), ({"cd": 1},)], - schema="m map", - ) - check = nested_map_keys_check( - "m", - lambda k: array_check( - F.split(k, ""), lambda ch: F.when(ch == "x", F.lit("x not allowed")) + lambda r: r == ["requires at least one of foo, bar"], + ), + _Case( + "mv_model_ok", + "map>", + {"a": {"foo": 1, "bar": None}}, + lambda f: map_values_check( + f, lambda v: check_require_any_of([v["foo"], v["bar"]], ["foo", "bar"]) ), - ) - rows = df.select(check.alias("errors")).collect() - # Row 0 has key "ax" with 'x' -> one flattened error; row 1 has no 'x' -> empty. - assert rows[0]["errors"] == ["x not allowed"] - assert rows[1]["errors"] == [] - - -def test_nested_array_check_wraps_map_values_check(spark: SparkSession) -> None: - # column `items`: array>>; flag any map value - # shorter than 3 chars in any element -- the `items[].tags{value}` pairing the - # generated `nested_array_check(... map_values_check ...)` folds. - schema = StructType( - [ - StructField( - "items", - ArrayType( - StructType( - [StructField("tags", MapType(StringType(), StringType()))] - ) - ), - ) - ] - ) - df = spark.createDataFrame( - [ - ([{"tags": {"k": "ab"}}],), - ([{"tags": {"k": "abc"}}, {"tags": {"k2": "wxyz"}}],), - ], - schema, - ) - check = nested_array_check( - "items", - lambda el: map_values_check( - el["tags"], lambda v: check_string_min_length(v, 3) + lambda r: r == [], + ), + _Case( + "mk_null", + "map", + None, + lambda f: map_keys_check(f, lambda k: F.lit("err")), + lambda r: r is None, + ), + _Case( + "mv_valid", + "map", + {"a": "ok"}, + lambda f: map_values_check(f, lambda v: F.when(v == "bad", F.lit("err"))), + lambda r: r == [], + ), + _Case( + "mk_col", + "map", + {"bad": "v"}, + lambda f: map_keys_check(F.col(f), lambda k: F.when(k == "bad", F.lit("err"))), + lambda r: r == ["err"], + ), + # --- coalesce_errors (input-independent literals) ------------------------- + _Case( + "coal_null", + "int", + 1, + lambda f: coalesce_errors(F.lit(None).cast("array")), + lambda r: r == [], + ), + _Case( + "coal_array", + "int", + 1, + lambda f: coalesce_errors(F.array(F.lit("err"))), + lambda r: r == ["err"], + ), + # --- nested_map_{values,keys}_check flatten an inner array -------- + _Case( + "nmv_flat_invalid", + "map>", + {"a": [1, 0], "b": [2]}, + lambda f: nested_map_values_check( + f, + lambda v: array_check( + v, lambda e: F.when(e == 0, F.lit("zero not allowed")) + ), ), - ) - rows = df.select(check.alias("errors")).collect() - # Row 0's "ab" is too short -> one flattened error; row 1 is clean -> empty. - assert rows[0]["errors"] != [] - assert rows[1]["errors"] == [] - - -def test_nested_map_values_check_wraps_map_values_check(spark: SparkSession) -> None: - # column `subs`: map>; flag any inner value < 0 -- the - # `subs{value}{value}` pairing the generated - # `nested_map_values_check(... map_values_check ...)` folds. - schema = StructType( - [ - StructField( - "subs", MapType(StringType(), MapType(StringType(), IntegerType())) - ) - ] - ) - df = spark.createDataFrame( - [({"k": {"a": -2}},), ({"k": {"a": 0, "b": 1}},)], - schema, - ) - check = nested_map_values_check( - "subs", - lambda v: map_values_check(v, lambda w: check_bounds(w, ge=0, check_nan=False)), - ) - rows = df.select(check.alias("errors")).collect() - # Row 0's -2 violates ge=0 -> one flattened error; row 1 is clean -> empty. - assert rows[0]["errors"] != [] - assert rows[1]["errors"] == [] + lambda r: r == ["zero not allowed"], + ), + _Case( + "nmv_flat_valid", + "map>", + {"a": [1], "b": [3]}, + lambda f: nested_map_values_check( + f, + lambda v: array_check( + v, lambda e: F.when(e == 0, F.lit("zero not allowed")) + ), + ), + lambda r: r == [], + ), + _Case( + "nmk_flat_invalid", + "map", + {"ax": 1, "b": 2}, + lambda f: nested_map_keys_check( + f, + lambda k: array_check( + F.split(k, ""), lambda ch: F.when(ch == "x", F.lit("x not allowed")) + ), + ), + lambda r: r == ["x not allowed"], + ), + _Case( + "nmk_flat_valid", + "map", + {"cd": 1}, + lambda f: nested_map_keys_check( + f, + lambda k: array_check( + F.split(k, ""), lambda ch: F.when(ch == "x", F.lit("x not allowed")) + ), + ), + lambda r: r == [], + ), + # --- flattening helper wrapped around map_values_check -------------------- + # `items[].tags{value}` and `subs{value}{value}`: the pairings the generated + # nested_array_check / nested_map_values_check fold around map_values_check. + _Case( + "nawmv_invalid", + "array>>", + [{"tags": {"k": "ab"}}], + lambda f: nested_array_check( + f, + lambda el: map_values_check( + el["tags"], lambda v: check_string_min_length(v, 3) + ), + ), + lambda r: r != [], + ), + _Case( + "nawmv_valid", + "array>>", + [{"tags": {"k": "abc"}}, {"tags": {"k2": "wxyz"}}], + lambda f: nested_array_check( + f, + lambda el: map_values_check( + el["tags"], lambda v: check_string_min_length(v, 3) + ), + ), + lambda r: r == [], + ), + _Case( + "nmvwmv_invalid", + "map>", + {"k": {"a": -2}}, + lambda f: nested_map_values_check( + f, + lambda v: map_values_check( + v, lambda w: check_bounds(w, ge=0, check_nan=False) + ), + ), + lambda r: r != [], + ), + _Case( + "nmvwmv_valid", + "map>", + {"k": {"a": 0, "b": 1}}, + lambda f: nested_map_values_check( + f, + lambda v: map_values_check( + v, lambda w: check_bounds(w, ge=0, check_nan=False) + ), + ), + lambda r: r == [], + ), +] + + +@pytest.fixture(scope="module") +def results(spark: SparkSession) -> Any: + """Pack every case's input into one row, apply every check, collect once.""" + # A DDL schema string, not StructType.fromDDL -- fromDDL landed in PySpark + # 3.5, and createDataFrame parses the string itself on the >=3.4 floor. + schema = ", ".join(f"`{c.id}` {c.ddl}" for c in _CASES) + row = {c.id: c.value for c in _CASES} + # dict rows are read by field name against the explicit schema, a form the + # createDataFrame stubs don't model (they want tuple/Row for RowLike). + df = spark.createDataFrame([row], schema=schema, verifySchema=False) # type: ignore[call-overload] + return df.select(*[c.check(c.id).alias(c.id) for c in _CASES]).collect()[0] + + +@pytest.mark.parametrize("case", _CASES, ids=lambda c: c.id) +def test_column_pattern(case: _Case, results: Any) -> None: + value = results[case.id] + assert case.expect(value), f"{case.id}: got {value!r}" diff --git a/packages/overture-schema-pyspark/tests/expressions/test_constraint_expressions.py b/packages/overture-schema-pyspark/tests/expressions/test_constraint_expressions.py index 248db0cf2..41ceaec77 100644 --- a/packages/overture-schema-pyspark/tests/expressions/test_constraint_expressions.py +++ b/packages/overture-schema-pyspark/tests/expressions/test_constraint_expressions.py @@ -1,6 +1,23 @@ -"""Tests for constraint_expressions — constraint type to Column translation.""" +"""Tests for constraint_expressions -- constraint type to Column translation. + +Each constraint is exercised as a `_Case`: a uniquely-named input column, the +check built over it, and a predicate on the collected result. The `results` +fixture packs every case's input into one wide single-row DataFrame, applies +every check in one `select`, and collects once -- so the whole file pays for a +single `createDataFrame` + `collect` instead of one pair per test (the same +batch-once pattern the generated conformance harness uses). Cases needing both +a violating and a clean input (or a valid/invalid pair of geometries) carry +separate entries. + +Column DDL types are load-bearing: `double` vs `int` changes NaN behavior, and +multi-field checks (require_if / forbid_if / radio_group / min_fields_set) pack +their inputs into a struct column so the check can reference each field. +""" import struct +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any import pytest from overture.schema.pyspark.expressions.constraint_expressions import ( @@ -31,1498 +48,1398 @@ except_literals, ) from overture.schema.system.primitive import GeometryType -from pyspark.sql import Row, SparkSession +from pyspark.sql import Column, SparkSession from pyspark.sql import functions as F -from pyspark.sql.types import DoubleType, StringType, StructField, StructType from shapely.geometry import LineString, MultiPolygon, Point, Polygon +# PySpark 3.4's collect() leaves its result socket for the GC to finalize; under +# -W error that ResourceWarning fails the batched `results` fixture. conftest's +# unraisablehook catches the finalizer path, but this fixture emits it +# synchronously -- and a filterwarnings mark is the only filter outranking the +# command-line -W error. +pytestmark = pytest.mark.filterwarnings("ignore::ResourceWarning") -def _except_literals_error(spark: SparkSession, value: str | None) -> str | None: - """Run `except_literals` over `check_url_format` for one string value.""" - df = spark.createDataFrame( - [Row(val=value)], schema=StructType([StructField("val", StringType(), True)]) - ) - col = F.col("val") - expr = except_literals(col, check_url_format(col), [""]) - # Spark Row field access is untyped (Any); the column holds an error string. - return df.select(expr.alias("err")).collect()[0]["err"] # type: ignore[no-any-return] - - -def test_except_literals_suppresses_allowed_literal(spark: SparkSession) -> None: - # "" is an allowed literal alternative -> the url_format error is suppressed. - assert _except_literals_error(spark, "") is None - - -def test_except_literals_passes_through_real_violation(spark: SparkSession) -> None: - # A non-literal invalid value still surfaces the inner check's error. - assert _except_literals_error(spark, "not a url") is not None - - -def test_except_literals_passes_through_valid_value(spark: SparkSession) -> None: - assert _except_literals_error(spark, "https://example.com/x") is None - - -def test_except_literals_null_is_not_an_error(spark: SparkSession) -> None: - assert _except_literals_error(spark, None) is None - - -def test_check_bounds_ge_le_valid(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val=5)]) - result = df.select(check_bounds(F.col("val"), ge=1, le=10).alias("err")).collect() - assert result[0]["err"] is None - - -def test_check_bounds_ge_violation(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val=0)]) - result = df.select(check_bounds(F.col("val"), ge=1).alias("err")).collect() - assert result[0]["err"] is not None - assert ">= 1" in result[0]["err"] - - -def test_check_bounds_gt_violation(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val=0)]) - result = df.select(check_bounds(F.col("val"), gt=0).alias("err")).collect() - assert result[0]["err"] is not None - assert "> 0" in result[0]["err"] - - -def test_check_bounds_le_violation(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val=100)]) - result = df.select(check_bounds(F.col("val"), le=50).alias("err")).collect() - assert result[0]["err"] is not None - - -def test_check_bounds_null_passthrough(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val=None)], schema="val int") - result = df.select(check_bounds(F.col("val"), ge=1).alias("err")).collect() - assert result[0]["err"] is None - - -def test_check_bounds_nan_lower_bound_violation(spark: SparkSession) -> None: - """NaN satisfies no Pydantic bound, but Spark sorts NaN above all values, - so a lower bound (NaN < v) never fires. check_bounds must reject it.""" - df = spark.createDataFrame([Row(val=float("nan"))], schema="val double") - result = df.select(check_bounds(F.col("val"), ge=0).alias("err")).collect() - assert result[0]["err"] is not None - assert "NaN" in result[0]["err"] - - -def test_check_bounds_nan_gt_violation(spark: SparkSession) -> None: - """Same lower-bound leak as ge, via the strict-greater comparison.""" - df = spark.createDataFrame([Row(val=float("nan"))], schema="val double") - result = df.select(check_bounds(F.col("val"), gt=0).alias("err")).collect() - assert result[0]["err"] is not None - assert "NaN" in result[0]["err"] - - -def test_check_bounds_nan_upper_bound_violation(spark: SparkSession) -> None: - """An upper bound already rejects NaN in Spark (NaN > v is true); the - explicit NaN check keeps that behavior.""" - df = spark.createDataFrame([Row(val=float("nan"))], schema="val double") - result = df.select(check_bounds(F.col("val"), le=1).alias("err")).collect() - assert result[0]["err"] is not None - - -def test_check_bounds_nan_no_bounds_passes(spark: SparkSession) -> None: - """With no bounds there is nothing to violate; NaN passes, matching - Pydantic's allow_inf_nan default for unconstrained floats.""" - df = spark.createDataFrame([Row(val=float("nan"))], schema="val double") - result = df.select(check_bounds(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - -def test_check_bounds_valid_float_passes(spark: SparkSession) -> None: - """A finite in-range float is unaffected by the NaN guard.""" - df = spark.createDataFrame([Row(val=0.5)], schema="val double") - result = df.select(check_bounds(F.col("val"), ge=0, le=1).alias("err")).collect() - assert result[0]["err"] is None - - -def test_check_bounds_nan_guard_off_passes_nan(spark: SparkSession) -> None: - """With check_nan=False the NaN guard is absent; NaN slips past a lower bound.""" - df = spark.createDataFrame([Row(val=float("nan"))], schema="val double") - result = df.select( - check_bounds(F.col("val"), ge=0, check_nan=False).alias("err") - ).collect() - assert result[0]["err"] is None - - -def test_check_bounds_nan_guard_on_rejects_nan(spark: SparkSession) -> None: - """With check_nan=True (default) NaN is rejected even with a lower bound.""" - df = spark.createDataFrame([Row(val=float("nan"))], schema="val double") - result = df.select( - check_bounds(F.col("val"), ge=0, check_nan=True).alias("err") - ).collect() - assert result[0]["err"] is not None - assert "NaN" in result[0]["err"] - - -def test_check_bounds_integer_column_rejects_violation(spark: SparkSession) -> None: - """check_nan=False is safe for integer columns; bound violations still fire.""" - df = spark.createDataFrame([Row(val=0)], schema="val int") - result = df.select( - check_bounds(F.col("val"), ge=1, check_nan=False).alias("err") - ).collect() - assert result[0]["err"] is not None - - -def test_check_bounds_integer_column_accepts_valid(spark: SparkSession) -> None: - """check_nan=False on an integer column: in-bound values pass.""" - df = spark.createDataFrame([Row(val=5)], schema="val int") - result = df.select( - check_bounds(F.col("val"), ge=1, le=10, check_nan=False).alias("err") - ).collect() - assert result[0]["err"] is None - - -def test_check_enum_valid(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="road")]) - result = df.select( - check_enum(F.col("val"), ["road", "rail", "water"]).alias("err") - ).collect() - assert result[0]["err"] is None - - -def test_check_enum_invalid(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="sky")]) - result = df.select( - check_enum(F.col("val"), ["road", "rail", "water"]).alias("err") - ).collect() - assert result[0]["err"] is not None - assert "sky" in result[0]["err"] - - -class TestCheckPattern: - def test_valid(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("AB",)], ["v"]) - result = df.select( - check_pattern(F.col("v"), r"^[A-Z]{2}$", label="test pattern").alias("e") - ) - assert result.collect()[0]["e"] is None - - def test_invalid(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("abc",)], ["v"]) - result = df.select( - check_pattern(F.col("v"), r"^[A-Z]{2}$", label="test pattern").alias("e") - ) - err = result.collect()[0]["e"] - assert "invalid test pattern" in err - assert "abc" in err - - def test_null_passes(self, spark: SparkSession) -> None: - df = spark.createDataFrame([(None,)], schema="v string") - result = df.select( - check_pattern(F.col("v"), r"^[A-Z]{2}$", label="test pattern").alias("e") - ) - assert result.collect()[0]["e"] is None - - -class TestCheckMinLength: - def test_at_limit(self, spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(items=["a", "b"])], schema="items array" - ) - result = df.select( - check_array_min_length(F.col("items"), 2).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_below_limit(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(items=["a"])], schema="items array") - result = df.select( - check_array_min_length(F.col("items"), 2).alias("err") - ).collect() - assert result[0]["err"] is not None - assert "minimum length 2" in result[0]["err"] - - def test_null_passthrough(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(items=None)], schema="items array") - result = df.select( - check_array_min_length(F.col("items"), 2).alias("err") - ).collect() - assert result[0]["err"] is None - - -class TestCheckMaxLength: - def test_within_limit(self, spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(items=["a", "b"])], schema="items array" - ) - result = df.select( - check_array_max_length(F.col("items"), 3).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_at_limit(self, spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(items=["a", "b"])], schema="items array" - ) - result = df.select( - check_array_max_length(F.col("items"), 2).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_exceeds_limit(self, spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(items=["a", "b", "c"])], schema="items array" - ) - result = df.select( - check_array_max_length(F.col("items"), 2).alias("err") - ).collect() - assert result[0]["err"] is not None - assert "maximum length 2" in result[0]["err"] - - def test_null_passthrough(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(items=None)], schema="items array") - result = df.select( - check_array_max_length(F.col("items"), 2).alias("err") - ).collect() - assert result[0]["err"] is None - - -def test_check_require_any_of_satisfied(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(a=1, b=None)], schema="a int, b int") - result = df.select( - check_require_any_of([F.col("a"), F.col("b")], ["a", "b"]).alias("err") - ).collect() - assert result[0]["err"] is None - - -def test_check_require_any_of_all_null(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(a=None, b=None)], schema="a int, b int") - result = df.select( - check_require_any_of([F.col("a"), F.col("b")], ["a", "b"]).alias("err") - ).collect() - assert result[0]["err"] is not None - assert "a" in result[0]["err"] - assert "b" in result[0]["err"] - - -class TestCheckRequireAnyTrue: - _NAMES = ["is_land", "is_territorial"] - - def _conds(self) -> list: - return [ - F.col("is_land") == F.lit(True), - F.col("is_territorial") == F.lit(True), - ] - - def test_one_condition_true(self, spark: SparkSession) -> None: - """At least one condition true -> no error.""" - df = spark.createDataFrame( - [Row(is_land=True, is_territorial=False)], - schema="is_land boolean, is_territorial boolean", - ) - result = df.select( - check_require_any_true(self._conds(), self._NAMES).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_all_conditions_false(self, spark: SparkSession) -> None: - """No condition true -> error naming the fields.""" - df = spark.createDataFrame( - [Row(is_land=False, is_territorial=False)], - schema="is_land boolean, is_territorial boolean", - ) - result = df.select( - check_require_any_true(self._conds(), self._NAMES).alias("err") - ).collect() - assert result[0]["err"] is not None - assert "is_land" in result[0]["err"] - assert "is_territorial" in result[0]["err"] - - def test_all_conditions_null_is_violation(self, spark: SparkSession) -> None: - """Null fields -> conditions not true -> error. - - Mirrors Python's `None == True` -> `False`: a null column value - does not satisfy the condition, so an all-null row violates. - """ - df = spark.createDataFrame( - [Row(is_land=None, is_territorial=None)], - schema="is_land boolean, is_territorial boolean", - ) - result = df.select( - check_require_any_true(self._conds(), self._NAMES).alias("err") - ).collect() - assert result[0]["err"] is not None - - -class TestCheckRequireIf: - def test_required_present(self, spark: SparkSession) -> None: - """Target is present when condition is true -> no error.""" - df = spark.createDataFrame( - [("road", "primary")], schema="subtype string, road_class string" - ) - result = df.select( - check_require_if( - F.col("road_class"), - F.col("subtype").isin(["road", "rail"]), - "subtype in [road, rail]", - ).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_required_absent(self, spark: SparkSession) -> None: - """Target is null when condition is true -> error.""" - df = spark.createDataFrame( - [("road", None)], schema="subtype string, road_class string" - ) - result = df.select( - check_require_if( - F.col("road_class"), - F.col("subtype").isin(["road", "rail"]), - "subtype in [road, rail]", - ).alias("err") - ).collect() - assert result[0]["err"] is not None - assert "required" in result[0]["err"] - - def test_condition_false_skips(self, spark: SparkSession) -> None: - """Target is null but condition is false -> no error.""" - df = spark.createDataFrame( - [("water", None)], schema="subtype string, road_class string" - ) - result = df.select( - check_require_if( - F.col("road_class"), - F.col("subtype").isin(["road", "rail"]), - "subtype in [road, rail]", - ).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_with_value_cols(self, spark: SparkSession) -> None: - """Error message includes actual discriminator value.""" - df = spark.createDataFrame( - [("road", None)], schema="subtype string, road_class string" - ) - result = df.select( - check_require_if( - F.col("road_class"), - F.col("subtype").isin(["road", "rail"]), - "subtype in [road, rail]", - F.col("subtype"), - ).alias("err") - ).collect() - assert "road" in result[0]["err"] - - -class TestCheckForbidIf: - def test_forbidden_absent(self, spark: SparkSession) -> None: - """Target is null when condition is true -> no error.""" - df = spark.createDataFrame( - [Row(subtype="country", parent=None)], - schema="subtype string, parent string", - ) - result = df.select( - check_forbid_if( - F.col("parent"), - F.col("subtype") == "country", - "subtype = country", - ).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_forbidden_present(self, spark: SparkSession) -> None: - """Target is present when condition is true -> error.""" - df = spark.createDataFrame([Row(subtype="country", parent="abc")]) - result = df.select( - check_forbid_if( - F.col("parent"), - F.col("subtype") == "country", - "subtype = country", - ).alias("err") - ).collect() - assert result[0]["err"] is not None - assert "forbidden" in result[0]["err"] - - def test_condition_false_skips(self, spark: SparkSession) -> None: - """Target is present but condition is false -> no error.""" - df = spark.createDataFrame([Row(subtype="region", parent="abc")]) - result = df.select( - check_forbid_if( - F.col("parent"), - F.col("subtype") == "country", - "subtype = country", - ).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_with_value_cols(self, spark: SparkSession) -> None: - """Error message includes actual discriminator value.""" - df = spark.createDataFrame([Row(subtype="country", parent="abc")]) - result = df.select( - check_forbid_if( - F.col("parent"), - F.col("subtype") == "country", - "subtype = country", - F.col("subtype"), - ).alias("err") - ).collect() - assert "country" in result[0]["err"] - - -class TestCheckStringMinLength: - def test_valid_length(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="abc")]) - result = df.select( - check_string_min_length(F.col("val"), 1).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_empty_string_violation(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="")]) - result = df.select( - check_string_min_length(F.col("val"), 1).alias("err") - ).collect() - assert result[0]["err"] is not None - assert "minimum length" in result[0]["err"] - - def test_null_passthrough(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val=None)], schema="val string") - result = df.select( - check_string_min_length(F.col("val"), 1).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_exact_min_length(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="ab")]) - result = df.select( - check_string_min_length(F.col("val"), 2).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_below_min_length(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="a")]) - result = df.select( - check_string_min_length(F.col("val"), 2).alias("err") - ).collect() - assert result[0]["err"] is not None - - -class TestCheckStringMaxLength: - def test_valid_length(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="abc")]) - result = df.select( - check_string_max_length(F.col("val"), 5).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_above_max_length(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="abcdef")]) - result = df.select( - check_string_max_length(F.col("val"), 5).alias("err") - ).collect() - assert result[0]["err"] is not None - assert "maximum length" in result[0]["err"] - - def test_null_passthrough(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val=None)], schema="val string") - result = df.select( - check_string_max_length(F.col("val"), 5).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_exact_max_length(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="abcde")]) - result = df.select( - check_string_max_length(F.col("val"), 5).alias("err") - ).collect() - assert result[0]["err"] is None - - -class TestCheckRadioGroup: - def test_exactly_one_true(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(is_land=True, is_territorial=False)]) - result = df.select( - check_radio_group( - [F.col("is_land"), F.col("is_territorial")], - ["is_land", "is_territorial"], - ).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_none_true(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(is_land=False, is_territorial=False)]) - result = df.select( - check_radio_group( - [F.col("is_land"), F.col("is_territorial")], - ["is_land", "is_territorial"], - ).alias("err") - ).collect() - assert result[0]["err"] is not None - assert "exactly one" in result[0]["err"] - assert "0" in result[0]["err"] - - def test_both_true(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(is_land=True, is_territorial=True)]) - result = df.select( - check_radio_group( - [F.col("is_land"), F.col("is_territorial")], - ["is_land", "is_territorial"], - ).alias("err") - ).collect() - assert result[0]["err"] is not None - assert "2" in result[0]["err"] - - def test_null_treated_as_false(self, spark: SparkSession) -> None: - """Null booleans count as not-true (0 toward the count).""" - df = spark.createDataFrame( - [Row(is_land=True, is_territorial=None)], - schema="is_land boolean, is_territorial boolean", - ) - result = df.select( - check_radio_group( - [F.col("is_land"), F.col("is_territorial")], - ["is_land", "is_territorial"], - ).alias("err") - ).collect() - assert result[0]["err"] is None - - -class TestCheckGeometryType: - def test_point_matches(self, spark: SparkSession) -> None: - wkb_bytes = Point(0, 0).wkb - df = spark.createDataFrame( - [Row(geometry=bytearray(wkb_bytes))], schema="geometry binary" - ) - result = df.select( - check_geometry_type(F.col("geometry"), GeometryType.POINT).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_point_rejects_linestring(self, spark: SparkSession) -> None: - wkb_bytes = LineString([(0, 0), (1, 1)]).wkb - df = spark.createDataFrame( - [Row(geometry=bytearray(wkb_bytes))], schema="geometry binary" - ) - result = df.select( - check_geometry_type(F.col("geometry"), GeometryType.POINT).alias("err") - ).collect() - assert result[0]["err"] is not None - assert "Point" in result[0]["err"] - - def test_multiple_allowed_types(self, spark: SparkSession) -> None: - wkb_polygon = Polygon([(0, 0), (1, 0), (1, 1), (0, 0)]).wkb - wkb_multi = MultiPolygon([Polygon([(0, 0), (1, 0), (1, 1), (0, 0)])]).wkb - df = spark.createDataFrame( - [ - Row(geometry=bytearray(wkb_polygon)), - Row(geometry=bytearray(wkb_multi)), - ], - schema="geometry binary", - ) - result = df.select( - check_geometry_type( - F.col("geometry"), - GeometryType.POLYGON, - GeometryType.MULTI_POLYGON, - ).alias("err") - ).collect() - assert all(r["err"] is None for r in result) - - def test_multiple_allowed_rejects_wrong_type(self, spark: SparkSession) -> None: - wkb_point = Point(0, 0).wkb - df = spark.createDataFrame( - [Row(geometry=bytearray(wkb_point))], schema="geometry binary" - ) - result = df.select( - check_geometry_type( - F.col("geometry"), - GeometryType.POLYGON, - GeometryType.MULTI_POLYGON, - ).alias("err") - ).collect() - assert result[0]["err"] is not None - - def test_null_passthrough(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(geometry=None)], schema="geometry binary") - result = df.select( - check_geometry_type(F.col("geometry"), GeometryType.POINT).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_big_endian_wkb(self, spark: SparkSession) -> None: - """Verify BE byte order handling. - - Shapely writes LE by default. Construct BE WKB for a Point - manually: byte_order=0x00, type=0x00000001, x=0.0, y=0.0. - """ - be_point = struct.pack(">bIdd", 0, 1, 0.0, 0.0) - df = spark.createDataFrame( - [Row(geometry=bytearray(be_point))], schema="geometry binary" - ) - result = df.select( - check_geometry_type(F.col("geometry"), GeometryType.POINT).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_iso_wkb_z_point_accepted(self, spark: SparkSession) -> None: - """ISO WKB encodes Z by offsetting the type (PointZ=1001), shifting - the low byte to 0xE9. GeoParquet mandates ISO WKB, so 3D geometries - reach the check this way and must still validate by base type.""" - iso_point_z = struct.pack(" None: - iso_point_z_be = struct.pack(">bIddd", 0, 1001, 0.0, 0.0, 5.0) - df = spark.createDataFrame( - [Row(geometry=bytearray(iso_point_z_be))], schema="geometry binary" - ) - result = df.select( - check_geometry_type(F.col("geometry"), GeometryType.POINT).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_ewkb_z_point_accepted(self, spark: SparkSession) -> None: - """EWKB encodes Z as a high flag bit (0x80000001), leaving the low - byte at 0x01 -- shapely's `.wkb` default. Must keep validating.""" - ewkb_point_z = struct.pack(" None: - """shapely's native 3D WKB output validates as POINT.""" - df = spark.createDataFrame( - [Row(geometry=bytearray(Point(0, 0, 5).wkb))], schema="geometry binary" - ) - result = df.select( - check_geometry_type(F.col("geometry"), GeometryType.POINT).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_iso_wkb_z_wrong_type_rejected(self, spark: SparkSession) -> None: - """A 3D LineString (ISO 1002) is still rejected when POINT is expected - -- normalization strips the dimension offset, not the base type.""" - iso_linestring_z = struct.pack(" None: - """A non-null WKB blob too short to contain a type word is flagged as a violation.""" - truncated = bytearray(b"\x01") - df = spark.createDataFrame([Row(geometry=truncated)], schema="geometry binary") - result = df.select( - check_geometry_type(F.col("geometry"), GeometryType.POINT).alias("err") - ).collect() - assert result[0]["err"] is not None - - @pytest.mark.parametrize("nbytes", [1, 2, 3, 4]) - def test_partial_header_wkb_flagged(self, spark: SparkSession, nbytes: int) -> None: - """A blob with a partial WKB header is flagged, even when conv() yields a non-null type. - - A little-endian order flag followed by a partial type word (2-4 bytes) - parses to a non-null but bogus base type -- e.g. `b"\\x01\\x01"` reads as - type 1, the Point code, and would silently validate as a Point. Only a - 0-1 byte blob makes conv() return NULL, so a length gate (not a null - check) is what closes the truncation hole. - """ - partial = bytearray(b"\x01" * nbytes) - df = spark.createDataFrame([Row(geometry=partial)], schema="geometry binary") - result = df.select( - check_geometry_type(F.col("geometry"), GeometryType.POINT).alias("err") - ).collect() - assert result[0]["err"] is not None - - -class TestCheckStripped: - def test_clean_string(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="hello world")]) - result = df.select(check_stripped(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - def test_single_char(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="x")]) - result = df.select(check_stripped(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - def test_leading_space(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val=" hello")]) - result = df.select(check_stripped(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - assert "whitespace" in result[0]["err"] - - def test_trailing_space(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="hello ")]) - result = df.select(check_stripped(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - assert "whitespace" in result[0]["err"] - - def test_leading_tab(self, spark: SparkSession) -> None: - """Tab is Unicode whitespace -- must be caught (not just ASCII space).""" - df = spark.createDataFrame([Row(val="\thello")]) - result = df.select(check_stripped(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - - def test_trailing_newline(self, spark: SparkSession) -> None: - """Trailing newline requires \\z anchor -- $ matches before it in Java regex.""" - df = spark.createDataFrame([Row(val="hello\n")]) - result = df.select(check_stripped(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - - def test_null_passthrough(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val=None)], schema="val string") - result = df.select(check_stripped(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - def test_empty_string(self, spark: SparkSession) -> None: - """Empty string has no leading/trailing whitespace -- passes.""" - df = spark.createDataFrame([Row(val="")]) - result = df.select(check_stripped(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - def test_trailing_unit_separator(self, spark: SparkSession) -> None: - """U+001F (unit separator) -- Python strips it, Java \\S with (?U) does not.""" - df = spark.createDataFrame([Row(val="Main St \x1f")]) - result = df.select(check_stripped(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - - def test_leading_file_separator(self, spark: SparkSession) -> None: - """U+001C (file separator) -- C0 control char Python treats as whitespace.""" - df = spark.createDataFrame([Row(val="\x1chello")]) - result = df.select(check_stripped(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - - def test_trailing_soh(self, spark: SparkSession) -> None: - """U+0001 (SOH) -- C0 control char that even Python's strip() misses.""" - df = spark.createDataFrame([Row(val="hello\x01")]) - result = df.select(check_stripped(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - - def test_trailing_del(self, spark: SparkSession) -> None: - """U+007F (DEL) -- control char outside C0 range.""" - df = spark.createDataFrame([Row(val="hello\x7f")]) - result = df.select(check_stripped(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - - def test_trailing_c1_control(self, spark: SparkSession) -> None: - """U+009F (APC) -- C1 control char.""" - df = spark.createDataFrame([Row(val="hello\x9f")]) - result = df.select(check_stripped(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - - def test_control_char_in_middle_passes(self, spark: SparkSession) -> None: - """Control chars in the middle of a string are not a stripped concern.""" - df = spark.createDataFrame([Row(val="hel\x1flo")]) - result = df.select(check_stripped(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - -class TestCheckJsonPointer: - def test_valid_pointer(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="/properties/name")]) - result = df.select(check_json_pointer(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - def test_root_pointer(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="/")]) - result = df.select(check_json_pointer(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - def test_empty_string_valid(self, spark: SparkSession) -> None: - """Empty string is valid per RFC 6901 (references whole document).""" - df = spark.createDataFrame([Row(val="")]) - result = df.select(check_json_pointer(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - def test_missing_leading_slash(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="properties/name")]) - result = df.select(check_json_pointer(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - assert "JSON pointer" in result[0]["err"] - assert "properties/name" in result[0]["err"] - - def test_null_passthrough(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val=None)], schema="val string") - result = df.select(check_json_pointer(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - -class TestCheckLinearRangeLength: - def test_valid_length(self, spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(between=[0.0, 1.0])], schema="between array" - ) - result = df.select( - check_linear_range_length(F.col("between")).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_wrong_length_one(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(between=[0.5])], schema="between array") - result = df.select( - check_linear_range_length(F.col("between")).alias("err") - ).collect() - assert result[0]["err"] is not None - assert "2 elements" in result[0]["err"] - - def test_wrong_length_three(self, spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(between=[0.0, 0.5, 1.0])], schema="between array" - ) - result = df.select( - check_linear_range_length(F.col("between")).alias("err") - ).collect() - assert result[0]["err"] is not None - assert "2 elements" in result[0]["err"] - - def test_empty_array(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(between=[])], schema="between array") - result = df.select( - check_linear_range_length(F.col("between")).alias("err") - ).collect() - assert result[0]["err"] is not None - assert "2 elements" in result[0]["err"] - - def test_null_passthrough(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(between=None)], schema="between array") - result = df.select( - check_linear_range_length(F.col("between")).alias("err") - ).collect() - assert result[0]["err"] is None - - -class TestCheckLinearRangeBounds: - def test_valid_bounds(self, spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(between=[0.2, 0.8])], schema="between array" - ) - result = df.select( - check_linear_range_bounds(F.col("between")).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_value_below_zero(self, spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(between=[-0.1, 0.5])], schema="between array" - ) - result = df.select( - check_linear_range_bounds(F.col("between")).alias("err") - ).collect() - assert result[0]["err"] is not None - assert "[0.0, 1.0]" in result[0]["err"] - - def test_value_above_one(self, spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(between=[0.0, 1.1])], schema="between array" - ) - result = df.select( - check_linear_range_bounds(F.col("between")).alias("err") - ).collect() - assert result[0]["err"] is not None - assert "[0.0, 1.0]" in result[0]["err"] - - def test_wrong_length_passthrough(self, spark: SparkSession) -> None: - """Wrong-length arrays are not this function's concern.""" - df = spark.createDataFrame([Row(between=[0.5])], schema="between array") - result = df.select( - check_linear_range_bounds(F.col("between")).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_null_passthrough(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(between=None)], schema="between array") - result = df.select( - check_linear_range_bounds(F.col("between")).alias("err") - ).collect() - assert result[0]["err"] is None - - -class TestCheckLinearRangeOrder: - def test_valid_order(self, spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(between=[0.2, 0.8])], schema="between array" - ) - result = df.select( - check_linear_range_order(F.col("between")).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_start_equals_end(self, spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(between=[0.5, 0.5])], schema="between array" - ) - result = df.select( - check_linear_range_order(F.col("between")).alias("err") - ).collect() - assert result[0]["err"] is not None - assert "start must be < end" in result[0]["err"] - - def test_start_after_end(self, spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(between=[0.8, 0.2])], schema="between array" - ) - result = df.select( - check_linear_range_order(F.col("between")).alias("err") - ).collect() - assert result[0]["err"] is not None - assert "start must be < end" in result[0]["err"] - - def test_wrong_length_passthrough(self, spark: SparkSession) -> None: - """Wrong-length arrays are not this function's concern.""" - df = spark.createDataFrame([Row(between=[0.5])], schema="between array") - result = df.select( - check_linear_range_order(F.col("between")).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_null_passthrough(self, spark: SparkSession) -> None: - df = spark.createDataFrame([Row(between=None)], schema="between array") - result = df.select( - check_linear_range_order(F.col("between")).alias("err") - ).collect() - assert result[0]["err"] is None - - -def test_check_required_null_is_error(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val=None)], schema="val string") - result = df.select(check_required(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - assert "missing" in result[0]["err"] - - -def test_check_required_non_null_passes(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="hello")]) - result = df.select(check_required(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - -def test_check_required_composes_with_enum(spark: SparkSession) -> None: - """check_required + check_enum via F.coalesce catches both null and invalid.""" - df = spark.createDataFrame([Row(val=None)], schema="val string") - expr = F.coalesce( - check_required(F.col("val")), - check_enum(F.col("val"), ["a", "b"]), - ) - result = df.select(expr.alias("err")).collect() - assert result[0]["err"] is not None - assert "missing" in result[0]["err"] - - +# Pattern/label pairs shared between a case's check and its expect predicate. _COUNTRY_CODE_PATTERN = r"^[A-Z]{2}\z" _COUNTRY_CODE_LABEL = "ISO 3166-1 alpha-2 country code" - - -class TestCheckCountryCodeViaPattern: - """Country code validation through check_pattern with label.""" - - def test_valid(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("US",)], ["v"]) - result = df.select( - check_pattern( - F.col("v"), _COUNTRY_CODE_PATTERN, label=_COUNTRY_CODE_LABEL - ).alias("e") - ) - assert result.collect()[0]["e"] is None - - def test_lowercase_invalid(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("us",)], ["v"]) - result = df.select( - check_pattern( - F.col("v"), _COUNTRY_CODE_PATTERN, label=_COUNTRY_CODE_LABEL - ).alias("e") - ) - err = result.collect()[0]["e"] - assert f"invalid {_COUNTRY_CODE_LABEL}" in err - assert "us" in err - - def test_three_chars_invalid(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("USA",)], ["v"]) - result = df.select( - check_pattern( - F.col("v"), _COUNTRY_CODE_PATTERN, label=_COUNTRY_CODE_LABEL - ).alias("e") - ) - assert result.collect()[0]["e"] is not None - - def test_null_passes(self, spark: SparkSession) -> None: - df = spark.createDataFrame([(None,)], schema="v string") - result = df.select( - check_pattern( - F.col("v"), _COUNTRY_CODE_PATTERN, label=_COUNTRY_CODE_LABEL - ).alias("e") - ) - assert result.collect()[0]["e"] is None - - _REGION_CODE_PATTERN = r"^[A-Z]{2}-[A-Z0-9]{1,3}\z" _REGION_CODE_LABEL = "ISO 3166-2 subdivision code" - - -class TestCheckRegionCodeViaPattern: - """Region code validation through check_pattern with label.""" - - def test_valid(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("US-NY",)], ["v"]) - result = df.select( - check_pattern( - F.col("v"), _REGION_CODE_PATTERN, label=_REGION_CODE_LABEL - ).alias("e") - ) - assert result.collect()[0]["e"] is None - - def test_valid_numeric(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("CN-11",)], ["v"]) - result = df.select( - check_pattern( - F.col("v"), _REGION_CODE_PATTERN, label=_REGION_CODE_LABEL - ).alias("e") - ) - assert result.collect()[0]["e"] is None - - def test_no_dash_invalid(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("USNY",)], ["v"]) - result = df.select( - check_pattern( - F.col("v"), _REGION_CODE_PATTERN, label=_REGION_CODE_LABEL - ).alias("e") - ) - err = result.collect()[0]["e"] - assert f"invalid {_REGION_CODE_LABEL}" in err - assert "USNY" in err - - def test_null_passes(self, spark: SparkSession) -> None: - df = spark.createDataFrame([(None,)], schema="v string") - result = df.select( - check_pattern( - F.col("v"), _REGION_CODE_PATTERN, label=_REGION_CODE_LABEL - ).alias("e") - ) - assert result.collect()[0]["e"] is None - - _SNAKE_CASE_PATTERN = r"^[a-z0-9]+(_[a-z0-9]+)*\z" _SNAKE_CASE_LABEL = "Category in snake_case format" - - -class TestCheckSnakeCaseViaPattern: - """Snake_case validation through check_pattern with label.""" - - def test_valid(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("hello_world",)], ["v"]) - result = df.select( - check_pattern( - F.col("v"), _SNAKE_CASE_PATTERN, label=_SNAKE_CASE_LABEL - ).alias("e") - ) - assert result.collect()[0]["e"] is None - - def test_single_word(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("hello",)], ["v"]) - result = df.select( - check_pattern( - F.col("v"), _SNAKE_CASE_PATTERN, label=_SNAKE_CASE_LABEL - ).alias("e") - ) - assert result.collect()[0]["e"] is None - - def test_with_numbers(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("hello_123",)], ["v"]) - result = df.select( - check_pattern( - F.col("v"), _SNAKE_CASE_PATTERN, label=_SNAKE_CASE_LABEL - ).alias("e") - ) - assert result.collect()[0]["e"] is None - - def test_uppercase_invalid(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("Hello_World",)], ["v"]) - result = df.select( - check_pattern( - F.col("v"), _SNAKE_CASE_PATTERN, label=_SNAKE_CASE_LABEL - ).alias("e") - ) - err = result.collect()[0]["e"] - assert f"invalid {_SNAKE_CASE_LABEL}" in err - - def test_spaces_invalid(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("hello world",)], ["v"]) - result = df.select( - check_pattern( - F.col("v"), _SNAKE_CASE_PATTERN, label=_SNAKE_CASE_LABEL - ).alias("e") - ) - assert result.collect()[0]["e"] is not None - - def test_null_passes(self, spark: SparkSession) -> None: - df = spark.createDataFrame([(None,)], schema="v string") - result = df.select( - check_pattern( - F.col("v"), _SNAKE_CASE_PATTERN, label=_SNAKE_CASE_LABEL - ).alias("e") - ) - assert result.collect()[0]["e"] is None - - -def test_check_url_format_http_valid(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="http://example.com")]) - result = df.select(check_url_format(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - -def test_check_url_format_https_valid(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="https://example.com/path?q=1")]) - result = df.select(check_url_format(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - -def test_check_url_format_uppercase_scheme_valid(spark: SparkSession) -> None: - """Pydantic HttpUrl lowercases the scheme, so HTTP:// is accepted.""" - df = spark.createDataFrame([Row(val="HTTP://example.com")]) - result = df.select(check_url_format(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - -def test_check_url_format_mixed_case_scheme_valid(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="HtTpS://example.com/path")]) - result = df.select(check_url_format(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - -def test_check_url_format_no_scheme_invalid(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="example.com")]) - result = df.select(check_url_format(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - - -def test_check_url_format_ftp_scheme_invalid(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="ftp://example.com")]) - result = df.select(check_url_format(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - - -def test_check_url_format_null_passes(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val=None)], schema="val string") - result = df.select(check_url_format(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - -def test_check_url_length_exceeds_2083_chars_invalid(spark: SparkSession) -> None: - long_url = "https://example.com/" + "a" * 2064 # 2084 chars - df = spark.createDataFrame([Row(val=long_url)]) - result = df.select(check_url_length(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - - -def test_check_url_length_exactly_2083_chars_valid(spark: SparkSession) -> None: - url = "https://example.com/" + "a" * 2063 # 2083 chars - df = spark.createDataFrame([Row(val=url)]) - result = df.select(check_url_length(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - -def test_check_url_length_null_passes(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val=None)], schema="val string") - result = df.select(check_url_length(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - -def test_check_email_valid(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="user@example.com")]) - result = df.select(check_email(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - -def test_check_email_no_at_invalid(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="userexample.com")]) - result = df.select(check_email(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - - -def test_check_email_no_domain_invalid(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="user@")]) - result = df.select(check_email(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - - -def test_check_email_spaces_invalid(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="user @example.com")]) - result = df.select(check_email(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - - -def test_check_email_null_passes(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val=None)], schema="val string") - result = df.select(check_email(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - -def test_check_email_trailing_period_invalid(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="user@example.com.")]) - result = df.select(check_email(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - - -def test_check_email_leading_period_invalid(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val=".user@example.com")]) - result = df.select(check_email(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - - -def test_check_email_period_before_at_invalid(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="user.@example.com")]) - result = df.select(check_email(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - - -def test_check_email_period_after_at_invalid(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="user@.example.com")]) - result = df.select(check_email(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - - -def test_check_email_double_period_domain_invalid(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="user@example..com")]) - result = df.select(check_email(F.col("val")).alias("err")).collect() - assert result[0]["err"] is not None - - -def test_check_email_dotted_local_valid(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="user.name@example.com")]) - result = df.select(check_email(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - -def test_check_email_subdomain_valid(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(val="user@mail.example.co.uk")]) - result = df.select(check_email(F.col("val")).alias("err")).collect() - assert result[0]["err"] is None - - _PHONE_PATTERN = r"^\+\d{1,3}[\s\-\(\)0-9]+\z" _PHONE_LABEL = "International phone number (+ followed by country code and number)" - - -class TestCheckPhoneViaPattern: - """Phone number validation through check_pattern with label.""" - - def test_valid_us(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("+1 555-555-5555",)], ["v"]) - result = df.select( - check_pattern(F.col("v"), _PHONE_PATTERN, label=_PHONE_LABEL).alias("e") - ) - assert result.collect()[0]["e"] is None - - def test_valid_international(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("+44 20 7946 0958",)], ["v"]) - result = df.select( - check_pattern(F.col("v"), _PHONE_PATTERN, label=_PHONE_LABEL).alias("e") - ) - assert result.collect()[0]["e"] is None - - def test_no_plus_invalid(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("555-555-5555",)], ["v"]) - result = df.select( - check_pattern(F.col("v"), _PHONE_PATTERN, label=_PHONE_LABEL).alias("e") - ) - err = result.collect()[0]["e"] - assert f"invalid {_PHONE_LABEL}" in err - - def test_letters_invalid(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("+1 abc-defg",)], ["v"]) - result = df.select( - check_pattern(F.col("v"), _PHONE_PATTERN, label=_PHONE_LABEL).alias("e") - ) - assert result.collect()[0]["e"] is not None - - def test_null_passes(self, spark: SparkSession) -> None: - df = spark.createDataFrame([(None,)], schema="v string") - result = df.select( - check_pattern(F.col("v"), _PHONE_PATTERN, label=_PHONE_LABEL).alias("e") - ) - assert result.collect()[0]["e"] is None - - _WIKIDATA_PATTERN = r"^Q\d+\z" _WIKIDATA_LABEL = "Wikidata identifier (Q followed by digits)" -class TestCheckWikidataIdViaPattern: - """Wikidata ID validation through check_pattern with label.""" - - def test_valid(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("Q42",)], ["v"]) - result = df.select( - check_pattern(F.col("v"), _WIKIDATA_PATTERN, label=_WIKIDATA_LABEL).alias( - "e" - ) - ) - assert result.collect()[0]["e"] is None +@dataclass(frozen=True) +class _Case: + """One constraint_expressions assertion driven off the shared wide row. - def test_large_number(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("Q123456789",)], ["v"]) - result = df.select( - check_pattern(F.col("v"), _WIKIDATA_PATTERN, label=_WIKIDATA_LABEL).alias( - "e" - ) - ) - assert result.collect()[0]["e"] is None + `check(field)` builds the constraint check over the case's input column; + `expect(result)` is a predicate on that column's single collected value. + """ - def test_lowercase_q_invalid(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("q42",)], ["v"]) - result = df.select( - check_pattern(F.col("v"), _WIKIDATA_PATTERN, label=_WIKIDATA_LABEL).alias( - "e" - ) - ) - err = result.collect()[0]["e"] - assert f"invalid {_WIKIDATA_LABEL}" in err + id: str + ddl: str + value: Any + check: Callable[[str], Column] + expect: Callable[[Any], bool] - def test_no_digits_invalid(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("Q",)], ["v"]) - result = df.select( - check_pattern(F.col("v"), _WIKIDATA_PATTERN, label=_WIKIDATA_LABEL).alias( - "e" - ) - ) - assert result.collect()[0]["e"] is not None - def test_p_prefix_invalid(self, spark: SparkSession) -> None: - df = spark.createDataFrame([("P42",)], ["v"]) - result = df.select( - check_pattern(F.col("v"), _WIKIDATA_PATTERN, label=_WIKIDATA_LABEL).alias( - "e" - ) - ) - assert result.collect()[0]["e"] is not None - - def test_null_passes(self, spark: SparkSession) -> None: - df = spark.createDataFrame([(None,)], schema="v string") - result = df.select( - check_pattern(F.col("v"), _WIKIDATA_PATTERN, label=_WIKIDATA_LABEL).alias( - "e" - ) - ) - assert result.collect()[0]["e"] is None - - -class TestCheckMinFieldsSet: - def test_meets_threshold(self, spark: SparkSession) -> None: - """Count at threshold -> no error.""" - df = spark.createDataFrame( - [Row(a=1, b=2, c=None)], schema="a int, b int, c int" - ) - result = df.select( - check_min_fields_set( - [F.col("a"), F.col("b"), F.col("c")], - ["a", "b", "c"], - 2, - ).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_exceeds_threshold(self, spark: SparkSession) -> None: - """Count above threshold -> no error.""" - df = spark.createDataFrame([Row(a=1, b=2, c=3)], schema="a int, b int, c int") - result = df.select( - check_min_fields_set( - [F.col("a"), F.col("b"), F.col("c")], - ["a", "b", "c"], - 2, - ).alias("err") - ).collect() - assert result[0]["err"] is None - - def test_below_threshold(self, spark: SparkSession) -> None: - """Count below threshold -> error with field names and actual count.""" - df = spark.createDataFrame( - [Row(a=1, b=None, c=None)], schema="a int, b int, c int" - ) - result = df.select( - check_min_fields_set( - [F.col("a"), F.col("b"), F.col("c")], - ["a", "b", "c"], - 2, - ).alias("err") - ).collect() - err = result[0]["err"] - assert err is not None - assert "at least 2" in err - assert "a, b, c" in err - assert "1" in err - - def test_all_null_below_threshold(self, spark: SparkSession) -> None: - """All null -> error showing 0 non-null.""" - df = spark.createDataFrame([Row(a=None, b=None)], schema="a int, b int") - result = df.select( - check_min_fields_set( - [F.col("a"), F.col("b")], - ["a", "b"], - 1, - ).alias("err") - ).collect() - err = result[0]["err"] - assert err is not None - assert "0" in err - - def test_error_message_format(self, spark: SparkSession) -> None: - """Error message matches expected format exactly.""" - df = spark.createDataFrame([Row(x=None, y=None)], schema="x int, y int") - result = df.select( - check_min_fields_set( - [F.col("x"), F.col("y")], - ["x", "y"], - 1, - ).alias("err") - ).collect() - err = result[0]["err"] - assert err == "at least 1 of x, y required, got 0 non-null" - - -_BBOX_SCHEMA = StructType( - [ - StructField( - "bbox", - StructType( - [ - StructField("xmin", DoubleType(), True), - StructField("xmax", DoubleType(), True), - StructField("ymin", DoubleType(), True), - StructField("ymax", DoubleType(), True), - ] - ), - True, +_CASES: list[_Case] = [ + # --- except_literals: suppress an allowed literal, pass real violations ---- + # "" is an allowed literal alternative -> the url_format error is suppressed. + _Case( + "el_suppress", + "string", + "", + lambda f: except_literals(F.col(f), check_url_format(F.col(f)), [""]), + lambda r: r is None, + ), + # A non-literal invalid value still surfaces the inner check's error. + _Case( + "el_violation", + "string", + "not a url", + lambda f: except_literals(F.col(f), check_url_format(F.col(f)), [""]), + lambda r: r is not None, + ), + _Case( + "el_valid", + "string", + "https://example.com/x", + lambda f: except_literals(F.col(f), check_url_format(F.col(f)), [""]), + lambda r: r is None, + ), + _Case( + "el_null", + "string", + None, + lambda f: except_literals(F.col(f), check_url_format(F.col(f)), [""]), + lambda r: r is None, + ), + # --- check_bounds --------------------------------------------------------- + _Case( + "bounds_ge_le_valid", + "int", + 5, + lambda f: check_bounds(F.col(f), ge=1, le=10), + lambda r: r is None, + ), + _Case( + "bounds_ge_violation", + "int", + 0, + lambda f: check_bounds(F.col(f), ge=1), + lambda r: r is not None and ">= 1" in r, + ), + _Case( + "bounds_gt_violation", + "int", + 0, + lambda f: check_bounds(F.col(f), gt=0), + lambda r: r is not None and "> 0" in r, + ), + _Case( + "bounds_le_violation", + "int", + 100, + lambda f: check_bounds(F.col(f), le=50), + lambda r: r is not None, + ), + _Case( + "bounds_null_passthrough", + "int", + None, + lambda f: check_bounds(F.col(f), ge=1), + lambda r: r is None, + ), + # NaN satisfies no Pydantic bound, but Spark sorts NaN above all values, so a + # lower bound (NaN < v) never fires. check_bounds must reject it explicitly. + _Case( + "bounds_nan_ge_violation", + "double", + float("nan"), + lambda f: check_bounds(F.col(f), ge=0), + lambda r: r is not None and "NaN" in r, + ), + # Same lower-bound leak, via the strict-greater comparison. + _Case( + "bounds_nan_gt_violation", + "double", + float("nan"), + lambda f: check_bounds(F.col(f), gt=0), + lambda r: r is not None and "NaN" in r, + ), + # An upper bound already rejects NaN in Spark (NaN > v is true); the explicit + # NaN check keeps that behavior. + _Case( + "bounds_nan_le_violation", + "double", + float("nan"), + lambda f: check_bounds(F.col(f), le=1), + lambda r: r is not None, + ), + # With no bounds there is nothing to violate; NaN passes, matching Pydantic's + # allow_inf_nan default for unconstrained floats. + _Case( + "bounds_nan_no_bounds_passes", + "double", + float("nan"), + lambda f: check_bounds(F.col(f)), + lambda r: r is None, + ), + # A finite in-range float is unaffected by the NaN guard. + _Case( + "bounds_valid_float_passes", + "double", + 0.5, + lambda f: check_bounds(F.col(f), ge=0, le=1), + lambda r: r is None, + ), + # With check_nan=False the NaN guard is absent; NaN slips past a lower bound. + _Case( + "bounds_nan_guard_off_passes", + "double", + float("nan"), + lambda f: check_bounds(F.col(f), ge=0, check_nan=False), + lambda r: r is None, + ), + # With check_nan=True (default) NaN is rejected even with a lower bound. + _Case( + "bounds_nan_guard_on_rejects", + "double", + float("nan"), + lambda f: check_bounds(F.col(f), ge=0, check_nan=True), + lambda r: r is not None and "NaN" in r, + ), + # check_nan=False is safe for integer columns; bound violations still fire. + _Case( + "bounds_int_rejects_violation", + "int", + 0, + lambda f: check_bounds(F.col(f), ge=1, check_nan=False), + lambda r: r is not None, + ), + _Case( + "bounds_int_accepts_valid", + "int", + 5, + lambda f: check_bounds(F.col(f), ge=1, le=10, check_nan=False), + lambda r: r is None, + ), + # --- check_enum ----------------------------------------------------------- + _Case( + "enum_valid", + "string", + "road", + lambda f: check_enum(F.col(f), ["road", "rail", "water"]), + lambda r: r is None, + ), + _Case( + "enum_invalid", + "string", + "sky", + lambda f: check_enum(F.col(f), ["road", "rail", "water"]), + lambda r: r is not None and "sky" in r, + ), + # --- check_pattern (generic) ---------------------------------------------- + _Case( + "pat_valid", + "string", + "AB", + lambda f: check_pattern(F.col(f), r"^[A-Z]{2}$", label="test pattern"), + lambda r: r is None, + ), + _Case( + "pat_invalid", + "string", + "abc", + lambda f: check_pattern(F.col(f), r"^[A-Z]{2}$", label="test pattern"), + lambda r: r is not None and "invalid test pattern" in r and "abc" in r, + ), + _Case( + "pat_null_passes", + "string", + None, + lambda f: check_pattern(F.col(f), r"^[A-Z]{2}$", label="test pattern"), + lambda r: r is None, + ), + # --- check_array_min_length ----------------------------------------------- + _Case( + "amin_at_limit", + "array", + ["a", "b"], + lambda f: check_array_min_length(F.col(f), 2), + lambda r: r is None, + ), + _Case( + "amin_below_limit", + "array", + ["a"], + lambda f: check_array_min_length(F.col(f), 2), + lambda r: r is not None and "minimum length 2" in r, + ), + _Case( + "amin_null_passthrough", + "array", + None, + lambda f: check_array_min_length(F.col(f), 2), + lambda r: r is None, + ), + # --- check_array_max_length ----------------------------------------------- + _Case( + "amax_within_limit", + "array", + ["a", "b"], + lambda f: check_array_max_length(F.col(f), 3), + lambda r: r is None, + ), + _Case( + "amax_at_limit", + "array", + ["a", "b"], + lambda f: check_array_max_length(F.col(f), 2), + lambda r: r is None, + ), + _Case( + "amax_exceeds_limit", + "array", + ["a", "b", "c"], + lambda f: check_array_max_length(F.col(f), 2), + lambda r: r is not None and "maximum length 2" in r, + ), + _Case( + "amax_null_passthrough", + "array", + None, + lambda f: check_array_max_length(F.col(f), 2), + lambda r: r is None, + ), + # --- check_require_any_of (multi-field -> struct column) ------------------ + _Case( + "rao_satisfied", + "struct", + {"a": 1, "b": None}, + lambda f: check_require_any_of([F.col(f)["a"], F.col(f)["b"]], ["a", "b"]), + lambda r: r is None, + ), + _Case( + "rao_all_null", + "struct", + {"a": None, "b": None}, + lambda f: check_require_any_of([F.col(f)["a"], F.col(f)["b"]], ["a", "b"]), + lambda r: r is not None and "a" in r and "b" in r, + ), + # --- check_require_any_true (multi-field -> struct column) ---------------- + # At least one condition true -> no error. + _Case( + "rat_one_true", + "struct", + {"is_land": True, "is_territorial": False}, + lambda f: check_require_any_true( + [ + F.col(f)["is_land"] == F.lit(True), + F.col(f)["is_territorial"] == F.lit(True), + ], + ["is_land", "is_territorial"], ), - ] -) - - -def test_check_bbox_completeness_valid(spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(bbox=Row(xmin=0.0, xmax=1.0, ymin=0.0, ymax=1.0))], - schema=_BBOX_SCHEMA, - ) - result = df.select(check_bbox_completeness(F.col("bbox")).alias("err")).collect() - assert result[0]["err"] is None - - -def test_check_bbox_completeness_null_bbox_passes(spark: SparkSession) -> None: - df = spark.createDataFrame([Row(bbox=None)], schema=_BBOX_SCHEMA) - result = df.select(check_bbox_completeness(F.col("bbox")).alias("err")).collect() - assert result[0]["err"] is None - - -def test_check_bbox_completeness_null_subfield_fails(spark: SparkSession) -> None: - df = spark.createDataFrame( - [Row(bbox=Row(xmin=None, xmax=1.0, ymin=0.0, ymax=1.0))], - schema=_BBOX_SCHEMA, - ) - result = df.select(check_bbox_completeness(F.col("bbox")).alias("err")).collect() - assert result[0]["err"] is not None + lambda r: r is None, + ), + # No condition true -> error naming the fields. + _Case( + "rat_all_false", + "struct", + {"is_land": False, "is_territorial": False}, + lambda f: check_require_any_true( + [ + F.col(f)["is_land"] == F.lit(True), + F.col(f)["is_territorial"] == F.lit(True), + ], + ["is_land", "is_territorial"], + ), + lambda r: r is not None and "is_land" in r and "is_territorial" in r, + ), + # Null fields mirror Python's `None == True` -> False: an all-null row violates. + _Case( + "rat_all_null", + "struct", + {"is_land": None, "is_territorial": None}, + lambda f: check_require_any_true( + [ + F.col(f)["is_land"] == F.lit(True), + F.col(f)["is_territorial"] == F.lit(True), + ], + ["is_land", "is_territorial"], + ), + lambda r: r is not None, + ), + # --- check_require_if (multi-field -> struct column) ---------------------- + # Target present when condition true -> no error. + _Case( + "rif_present", + "struct", + {"subtype": "road", "road_class": "primary"}, + lambda f: check_require_if( + F.col(f)["road_class"], + F.col(f)["subtype"].isin(["road", "rail"]), + "subtype in [road, rail]", + ), + lambda r: r is None, + ), + # Target null when condition true -> error. + _Case( + "rif_absent", + "struct", + {"subtype": "road", "road_class": None}, + lambda f: check_require_if( + F.col(f)["road_class"], + F.col(f)["subtype"].isin(["road", "rail"]), + "subtype in [road, rail]", + ), + lambda r: r is not None and "required" in r, + ), + # Target null but condition false -> no error. + _Case( + "rif_condition_false", + "struct", + {"subtype": "water", "road_class": None}, + lambda f: check_require_if( + F.col(f)["road_class"], + F.col(f)["subtype"].isin(["road", "rail"]), + "subtype in [road, rail]", + ), + lambda r: r is None, + ), + # Error message includes the actual discriminator value. + _Case( + "rif_value_cols", + "struct", + {"subtype": "road", "road_class": None}, + lambda f: check_require_if( + F.col(f)["road_class"], + F.col(f)["subtype"].isin(["road", "rail"]), + "subtype in [road, rail]", + F.col(f)["subtype"], + ), + lambda r: r is not None and "road" in r, + ), + # --- check_forbid_if (multi-field -> struct column) ----------------------- + # Target null when condition true -> no error. + _Case( + "fif_absent", + "struct", + {"subtype": "country", "parent": None}, + lambda f: check_forbid_if( + F.col(f)["parent"], + F.col(f)["subtype"] == "country", + "subtype = country", + ), + lambda r: r is None, + ), + # Target present when condition true -> error. + _Case( + "fif_present", + "struct", + {"subtype": "country", "parent": "abc"}, + lambda f: check_forbid_if( + F.col(f)["parent"], + F.col(f)["subtype"] == "country", + "subtype = country", + ), + lambda r: r is not None and "forbidden" in r, + ), + # Target present but condition false -> no error. + _Case( + "fif_condition_false", + "struct", + {"subtype": "region", "parent": "abc"}, + lambda f: check_forbid_if( + F.col(f)["parent"], + F.col(f)["subtype"] == "country", + "subtype = country", + ), + lambda r: r is None, + ), + # Error message includes the actual discriminator value. + _Case( + "fif_value_cols", + "struct", + {"subtype": "country", "parent": "abc"}, + lambda f: check_forbid_if( + F.col(f)["parent"], + F.col(f)["subtype"] == "country", + "subtype = country", + F.col(f)["subtype"], + ), + lambda r: r is not None and "country" in r, + ), + # --- check_string_min_length ---------------------------------------------- + _Case( + "smin_valid", + "string", + "abc", + lambda f: check_string_min_length(F.col(f), 1), + lambda r: r is None, + ), + _Case( + "smin_empty_violation", + "string", + "", + lambda f: check_string_min_length(F.col(f), 1), + lambda r: r is not None and "minimum length" in r, + ), + _Case( + "smin_null_passthrough", + "string", + None, + lambda f: check_string_min_length(F.col(f), 1), + lambda r: r is None, + ), + _Case( + "smin_exact", + "string", + "ab", + lambda f: check_string_min_length(F.col(f), 2), + lambda r: r is None, + ), + _Case( + "smin_below", + "string", + "a", + lambda f: check_string_min_length(F.col(f), 2), + lambda r: r is not None, + ), + # --- check_string_max_length ---------------------------------------------- + _Case( + "smax_valid", + "string", + "abc", + lambda f: check_string_max_length(F.col(f), 5), + lambda r: r is None, + ), + _Case( + "smax_above", + "string", + "abcdef", + lambda f: check_string_max_length(F.col(f), 5), + lambda r: r is not None and "maximum length" in r, + ), + _Case( + "smax_null_passthrough", + "string", + None, + lambda f: check_string_max_length(F.col(f), 5), + lambda r: r is None, + ), + _Case( + "smax_exact", + "string", + "abcde", + lambda f: check_string_max_length(F.col(f), 5), + lambda r: r is None, + ), + # --- check_radio_group (multi-field -> struct column) --------------------- + _Case( + "rg_exactly_one", + "struct", + {"is_land": True, "is_territorial": False}, + lambda f: check_radio_group( + [F.col(f)["is_land"], F.col(f)["is_territorial"]], + ["is_land", "is_territorial"], + ), + lambda r: r is None, + ), + _Case( + "rg_none_true", + "struct", + {"is_land": False, "is_territorial": False}, + lambda f: check_radio_group( + [F.col(f)["is_land"], F.col(f)["is_territorial"]], + ["is_land", "is_territorial"], + ), + lambda r: r is not None and "exactly one" in r and "0" in r, + ), + _Case( + "rg_both_true", + "struct", + {"is_land": True, "is_territorial": True}, + lambda f: check_radio_group( + [F.col(f)["is_land"], F.col(f)["is_territorial"]], + ["is_land", "is_territorial"], + ), + lambda r: r is not None and "2" in r, + ), + # Null booleans count as not-true (0 toward the count). + _Case( + "rg_null_as_false", + "struct", + {"is_land": True, "is_territorial": None}, + lambda f: check_radio_group( + [F.col(f)["is_land"], F.col(f)["is_territorial"]], + ["is_land", "is_territorial"], + ), + lambda r: r is None, + ), + # --- check_geometry_type (WKB in a binary column) ------------------------- + _Case( + "geom_point_matches", + "binary", + bytearray(Point(0, 0).wkb), + lambda f: check_geometry_type(F.col(f), GeometryType.POINT), + lambda r: r is None, + ), + _Case( + "geom_point_rejects_line", + "binary", + bytearray(LineString([(0, 0), (1, 1)]).wkb), + lambda f: check_geometry_type(F.col(f), GeometryType.POINT), + lambda r: r is not None and "Point" in r, + ), + # Multiple allowed types: a polygon and a multipolygon both pass (one row each). + _Case( + "geom_multi_polygon_ok", + "binary", + bytearray(Polygon([(0, 0), (1, 0), (1, 1), (0, 0)]).wkb), + lambda f: check_geometry_type( + F.col(f), GeometryType.POLYGON, GeometryType.MULTI_POLYGON + ), + lambda r: r is None, + ), + _Case( + "geom_multi_multipolygon_ok", + "binary", + bytearray(MultiPolygon([Polygon([(0, 0), (1, 0), (1, 1), (0, 0)])]).wkb), + lambda f: check_geometry_type( + F.col(f), GeometryType.POLYGON, GeometryType.MULTI_POLYGON + ), + lambda r: r is None, + ), + _Case( + "geom_multi_rejects_point", + "binary", + bytearray(Point(0, 0).wkb), + lambda f: check_geometry_type( + F.col(f), GeometryType.POLYGON, GeometryType.MULTI_POLYGON + ), + lambda r: r is not None, + ), + _Case( + "geom_null_passthrough", + "binary", + None, + lambda f: check_geometry_type(F.col(f), GeometryType.POINT), + lambda r: r is None, + ), + # BE byte order: byte_order=0x00, type=0x00000001, x=0.0, y=0.0. + _Case( + "geom_big_endian", + "binary", + bytearray(struct.pack(">bIdd", 0, 1, 0.0, 0.0)), + lambda f: check_geometry_type(F.col(f), GeometryType.POINT), + lambda r: r is None, + ), + # ISO WKB encodes Z by offsetting the type (PointZ=1001); must validate by base type. + _Case( + "geom_iso_z_point", + "binary", + bytearray(struct.pack("bIddd", 0, 1001, 0.0, 0.0, 5.0)), + lambda f: check_geometry_type(F.col(f), GeometryType.POINT), + lambda r: r is None, + ), + # EWKB encodes Z as a high flag bit (0x80000001), leaving the low byte at 0x01. + _Case( + "geom_ewkb_z_point", + "binary", + bytearray(struct.pack("", + [0.0, 1.0], + lambda f: check_linear_range_length(F.col(f)), + lambda r: r is None, + ), + _Case( + "lrl_one", + "array", + [0.5], + lambda f: check_linear_range_length(F.col(f)), + lambda r: r is not None and "2 elements" in r, + ), + _Case( + "lrl_three", + "array", + [0.0, 0.5, 1.0], + lambda f: check_linear_range_length(F.col(f)), + lambda r: r is not None and "2 elements" in r, + ), + _Case( + "lrl_empty", + "array", + [], + lambda f: check_linear_range_length(F.col(f)), + lambda r: r is not None and "2 elements" in r, + ), + _Case( + "lrl_null", + "array", + None, + lambda f: check_linear_range_length(F.col(f)), + lambda r: r is None, + ), + # --- check_linear_range_bounds -------------------------------------------- + _Case( + "lrb_valid", + "array", + [0.2, 0.8], + lambda f: check_linear_range_bounds(F.col(f)), + lambda r: r is None, + ), + _Case( + "lrb_below_zero", + "array", + [-0.1, 0.5], + lambda f: check_linear_range_bounds(F.col(f)), + lambda r: r is not None and "[0.0, 1.0]" in r, + ), + _Case( + "lrb_above_one", + "array", + [0.0, 1.1], + lambda f: check_linear_range_bounds(F.col(f)), + lambda r: r is not None and "[0.0, 1.0]" in r, + ), + # Wrong-length arrays are not this function's concern. + _Case( + "lrb_wrong_length_passthrough", + "array", + [0.5], + lambda f: check_linear_range_bounds(F.col(f)), + lambda r: r is None, + ), + _Case( + "lrb_null", + "array", + None, + lambda f: check_linear_range_bounds(F.col(f)), + lambda r: r is None, + ), + # --- check_linear_range_order --------------------------------------------- + _Case( + "lro_valid", + "array", + [0.2, 0.8], + lambda f: check_linear_range_order(F.col(f)), + lambda r: r is None, + ), + _Case( + "lro_equal", + "array", + [0.5, 0.5], + lambda f: check_linear_range_order(F.col(f)), + lambda r: r is not None and "start must be < end" in r, + ), + _Case( + "lro_after", + "array", + [0.8, 0.2], + lambda f: check_linear_range_order(F.col(f)), + lambda r: r is not None and "start must be < end" in r, + ), + # Wrong-length arrays are not this function's concern. + _Case( + "lro_wrong_length_passthrough", + "array", + [0.5], + lambda f: check_linear_range_order(F.col(f)), + lambda r: r is None, + ), + _Case( + "lro_null", + "array", + None, + lambda f: check_linear_range_order(F.col(f)), + lambda r: r is None, + ), + # --- check_required ------------------------------------------------------- + _Case( + "req_null_is_error", + "string", + None, + lambda f: check_required(F.col(f)), + lambda r: r is not None and "missing" in r, + ), + _Case( + "req_non_null_passes", + "string", + "hello", + lambda f: check_required(F.col(f)), + lambda r: r is None, + ), + # check_required + check_enum via F.coalesce catches both null and invalid. + _Case( + "req_composes_with_enum", + "string", + None, + lambda f: F.coalesce( + check_required(F.col(f)), check_enum(F.col(f), ["a", "b"]) + ), + lambda r: r is not None and "missing" in r, + ), + # --- country code via check_pattern --------------------------------------- + _Case( + "cc_valid", + "string", + "US", + lambda f: check_pattern( + F.col(f), _COUNTRY_CODE_PATTERN, label=_COUNTRY_CODE_LABEL + ), + lambda r: r is None, + ), + _Case( + "cc_lowercase_invalid", + "string", + "us", + lambda f: check_pattern( + F.col(f), _COUNTRY_CODE_PATTERN, label=_COUNTRY_CODE_LABEL + ), + lambda r: r is not None and f"invalid {_COUNTRY_CODE_LABEL}" in r and "us" in r, + ), + _Case( + "cc_three_chars_invalid", + "string", + "USA", + lambda f: check_pattern( + F.col(f), _COUNTRY_CODE_PATTERN, label=_COUNTRY_CODE_LABEL + ), + lambda r: r is not None, + ), + _Case( + "cc_null_passes", + "string", + None, + lambda f: check_pattern( + F.col(f), _COUNTRY_CODE_PATTERN, label=_COUNTRY_CODE_LABEL + ), + lambda r: r is None, + ), + # --- region code via check_pattern ---------------------------------------- + _Case( + "rc_valid", + "string", + "US-NY", + lambda f: check_pattern( + F.col(f), _REGION_CODE_PATTERN, label=_REGION_CODE_LABEL + ), + lambda r: r is None, + ), + _Case( + "rc_valid_numeric", + "string", + "CN-11", + lambda f: check_pattern( + F.col(f), _REGION_CODE_PATTERN, label=_REGION_CODE_LABEL + ), + lambda r: r is None, + ), + _Case( + "rc_no_dash_invalid", + "string", + "USNY", + lambda f: check_pattern( + F.col(f), _REGION_CODE_PATTERN, label=_REGION_CODE_LABEL + ), + lambda r: ( + r is not None and f"invalid {_REGION_CODE_LABEL}" in r and "USNY" in r + ), + ), + _Case( + "rc_null_passes", + "string", + None, + lambda f: check_pattern( + F.col(f), _REGION_CODE_PATTERN, label=_REGION_CODE_LABEL + ), + lambda r: r is None, + ), + # --- snake_case via check_pattern ----------------------------------------- + _Case( + "sc_valid", + "string", + "hello_world", + lambda f: check_pattern(F.col(f), _SNAKE_CASE_PATTERN, label=_SNAKE_CASE_LABEL), + lambda r: r is None, + ), + _Case( + "sc_single_word", + "string", + "hello", + lambda f: check_pattern(F.col(f), _SNAKE_CASE_PATTERN, label=_SNAKE_CASE_LABEL), + lambda r: r is None, + ), + _Case( + "sc_with_numbers", + "string", + "hello_123", + lambda f: check_pattern(F.col(f), _SNAKE_CASE_PATTERN, label=_SNAKE_CASE_LABEL), + lambda r: r is None, + ), + _Case( + "sc_uppercase_invalid", + "string", + "Hello_World", + lambda f: check_pattern(F.col(f), _SNAKE_CASE_PATTERN, label=_SNAKE_CASE_LABEL), + lambda r: r is not None and f"invalid {_SNAKE_CASE_LABEL}" in r, + ), + _Case( + "sc_spaces_invalid", + "string", + "hello world", + lambda f: check_pattern(F.col(f), _SNAKE_CASE_PATTERN, label=_SNAKE_CASE_LABEL), + lambda r: r is not None, + ), + _Case( + "sc_null_passes", + "string", + None, + lambda f: check_pattern(F.col(f), _SNAKE_CASE_PATTERN, label=_SNAKE_CASE_LABEL), + lambda r: r is None, + ), + # --- check_url_format ----------------------------------------------------- + _Case( + "url_http_valid", + "string", + "http://example.com", + lambda f: check_url_format(F.col(f)), + lambda r: r is None, + ), + _Case( + "url_https_valid", + "string", + "https://example.com/path?q=1", + lambda f: check_url_format(F.col(f)), + lambda r: r is None, + ), + # Pydantic HttpUrl lowercases the scheme, so HTTP:// is accepted. + _Case( + "url_uppercase_scheme", + "string", + "HTTP://example.com", + lambda f: check_url_format(F.col(f)), + lambda r: r is None, + ), + _Case( + "url_mixed_case_scheme", + "string", + "HtTpS://example.com/path", + lambda f: check_url_format(F.col(f)), + lambda r: r is None, + ), + _Case( + "url_no_scheme_invalid", + "string", + "example.com", + lambda f: check_url_format(F.col(f)), + lambda r: r is not None, + ), + _Case( + "url_ftp_scheme_invalid", + "string", + "ftp://example.com", + lambda f: check_url_format(F.col(f)), + lambda r: r is not None, + ), + _Case( + "url_null_passes", + "string", + None, + lambda f: check_url_format(F.col(f)), + lambda r: r is None, + ), + # --- check_url_length ----------------------------------------------------- + _Case( + "urllen_exceeds", + "string", + "https://example.com/" + "a" * 2064, # 2084 chars + lambda f: check_url_length(F.col(f)), + lambda r: r is not None, + ), + _Case( + "urllen_exactly_2083", + "string", + "https://example.com/" + "a" * 2063, # 2083 chars + lambda f: check_url_length(F.col(f)), + lambda r: r is None, + ), + _Case( + "urllen_null_passes", + "string", + None, + lambda f: check_url_length(F.col(f)), + lambda r: r is None, + ), + # --- check_email ---------------------------------------------------------- + _Case( + "email_valid", + "string", + "user@example.com", + lambda f: check_email(F.col(f)), + lambda r: r is None, + ), + _Case( + "email_no_at", + "string", + "userexample.com", + lambda f: check_email(F.col(f)), + lambda r: r is not None, + ), + _Case( + "email_no_domain", + "string", + "user@", + lambda f: check_email(F.col(f)), + lambda r: r is not None, + ), + _Case( + "email_spaces", + "string", + "user @example.com", + lambda f: check_email(F.col(f)), + lambda r: r is not None, + ), + _Case( + "email_null", + "string", + None, + lambda f: check_email(F.col(f)), + lambda r: r is None, + ), + _Case( + "email_trailing_period", + "string", + "user@example.com.", + lambda f: check_email(F.col(f)), + lambda r: r is not None, + ), + _Case( + "email_leading_period", + "string", + ".user@example.com", + lambda f: check_email(F.col(f)), + lambda r: r is not None, + ), + _Case( + "email_period_before_at", + "string", + "user.@example.com", + lambda f: check_email(F.col(f)), + lambda r: r is not None, + ), + _Case( + "email_period_after_at", + "string", + "user@.example.com", + lambda f: check_email(F.col(f)), + lambda r: r is not None, + ), + _Case( + "email_double_period_domain", + "string", + "user@example..com", + lambda f: check_email(F.col(f)), + lambda r: r is not None, + ), + _Case( + "email_dotted_local_valid", + "string", + "user.name@example.com", + lambda f: check_email(F.col(f)), + lambda r: r is None, + ), + _Case( + "email_subdomain_valid", + "string", + "user@mail.example.co.uk", + lambda f: check_email(F.col(f)), + lambda r: r is None, + ), + # --- phone via check_pattern ---------------------------------------------- + _Case( + "phone_valid_us", + "string", + "+1 555-555-5555", + lambda f: check_pattern(F.col(f), _PHONE_PATTERN, label=_PHONE_LABEL), + lambda r: r is None, + ), + _Case( + "phone_valid_international", + "string", + "+44 20 7946 0958", + lambda f: check_pattern(F.col(f), _PHONE_PATTERN, label=_PHONE_LABEL), + lambda r: r is None, + ), + _Case( + "phone_no_plus_invalid", + "string", + "555-555-5555", + lambda f: check_pattern(F.col(f), _PHONE_PATTERN, label=_PHONE_LABEL), + lambda r: r is not None and f"invalid {_PHONE_LABEL}" in r, + ), + _Case( + "phone_letters_invalid", + "string", + "+1 abc-defg", + lambda f: check_pattern(F.col(f), _PHONE_PATTERN, label=_PHONE_LABEL), + lambda r: r is not None, + ), + _Case( + "phone_null_passes", + "string", + None, + lambda f: check_pattern(F.col(f), _PHONE_PATTERN, label=_PHONE_LABEL), + lambda r: r is None, + ), + # --- wikidata id via check_pattern ---------------------------------------- + _Case( + "wd_valid", + "string", + "Q42", + lambda f: check_pattern(F.col(f), _WIKIDATA_PATTERN, label=_WIKIDATA_LABEL), + lambda r: r is None, + ), + _Case( + "wd_large_number", + "string", + "Q123456789", + lambda f: check_pattern(F.col(f), _WIKIDATA_PATTERN, label=_WIKIDATA_LABEL), + lambda r: r is None, + ), + _Case( + "wd_lowercase_q_invalid", + "string", + "q42", + lambda f: check_pattern(F.col(f), _WIKIDATA_PATTERN, label=_WIKIDATA_LABEL), + lambda r: r is not None and f"invalid {_WIKIDATA_LABEL}" in r, + ), + _Case( + "wd_no_digits_invalid", + "string", + "Q", + lambda f: check_pattern(F.col(f), _WIKIDATA_PATTERN, label=_WIKIDATA_LABEL), + lambda r: r is not None, + ), + _Case( + "wd_p_prefix_invalid", + "string", + "P42", + lambda f: check_pattern(F.col(f), _WIKIDATA_PATTERN, label=_WIKIDATA_LABEL), + lambda r: r is not None, + ), + _Case( + "wd_null_passes", + "string", + None, + lambda f: check_pattern(F.col(f), _WIKIDATA_PATTERN, label=_WIKIDATA_LABEL), + lambda r: r is None, + ), + # --- check_min_fields_set (multi-field -> struct column) ------------------ + # Count at threshold -> no error. + _Case( + "mfs_meets_threshold", + "struct", + {"a": 1, "b": 2, "c": None}, + lambda f: check_min_fields_set( + [F.col(f)["a"], F.col(f)["b"], F.col(f)["c"]], ["a", "b", "c"], 2 + ), + lambda r: r is None, + ), + # Count above threshold -> no error. + _Case( + "mfs_exceeds_threshold", + "struct", + {"a": 1, "b": 2, "c": 3}, + lambda f: check_min_fields_set( + [F.col(f)["a"], F.col(f)["b"], F.col(f)["c"]], ["a", "b", "c"], 2 + ), + lambda r: r is None, + ), + # Count below threshold -> error with field names and actual count. + _Case( + "mfs_below_threshold", + "struct", + {"a": 1, "b": None, "c": None}, + lambda f: check_min_fields_set( + [F.col(f)["a"], F.col(f)["b"], F.col(f)["c"]], ["a", "b", "c"], 2 + ), + lambda r: r is not None and "at least 2" in r and "a, b, c" in r and "1" in r, + ), + # All null -> error showing 0 non-null. + _Case( + "mfs_all_null", + "struct", + {"a": None, "b": None}, + lambda f: check_min_fields_set([F.col(f)["a"], F.col(f)["b"]], ["a", "b"], 1), + lambda r: r is not None and "0" in r, + ), + # Error message matches the expected format exactly. + _Case( + "mfs_message_format", + "struct", + {"x": None, "y": None}, + lambda f: check_min_fields_set([F.col(f)["x"], F.col(f)["y"]], ["x", "y"], 1), + lambda r: r == "at least 1 of x, y required, got 0 non-null", + ), + # --- check_bbox_completeness (bbox struct column) ------------------------- + _Case( + "bbox_valid", + "struct", + {"xmin": 0.0, "xmax": 1.0, "ymin": 0.0, "ymax": 1.0}, + lambda f: check_bbox_completeness(F.col(f)), + lambda r: r is None, + ), + _Case( + "bbox_null_passes", + "struct", + None, + lambda f: check_bbox_completeness(F.col(f)), + lambda r: r is None, + ), + _Case( + "bbox_null_subfield_fails", + "struct", + {"xmin": None, "xmax": 1.0, "ymin": 0.0, "ymax": 1.0}, + lambda f: check_bbox_completeness(F.col(f)), + lambda r: r is not None, + ), +] + + +@pytest.fixture(scope="module") +def results(spark: SparkSession) -> Any: + """Pack every case's input into one row, apply every check, collect once.""" + # A DDL schema string, not StructType.fromDDL -- fromDDL landed in PySpark + # 3.5, and createDataFrame parses the string itself on the >=3.4 floor. + schema = ", ".join(f"`{c.id}` {c.ddl}" for c in _CASES) + row = {c.id: c.value for c in _CASES} + # dict rows are read by field name against the explicit schema, a form the + # createDataFrame stubs don't model (they want tuple/Row for RowLike). + df = spark.createDataFrame([row], schema=schema, verifySchema=False) # type: ignore[call-overload] + return df.select(*[c.check(c.id).alias(c.id) for c in _CASES]).collect()[0] + + +@pytest.mark.parametrize("case", _CASES, ids=lambda c: c.id) +def test_constraint_expression(case: _Case, results: Any) -> None: + value = results[case.id] + assert case.expect(value), f"{case.id}: got {value!r}"