fix: Make unique work for nested types and improve performance (#333)

borchero · web-flow · commit b993da6d504f · 2026-04-22T10:02:40.000+02:00
diff --git a/dataframely/_base_schema.py b/dataframely/_base_schema.py
@@ -42,13 +42,6 @@ def _build_rules(
     if len(primary_key) > 0:
         rules["primary_key"] = Rule(pl.struct(primary_key).is_unique())
 
-    # Add unique column validation rules
-    unique_columns = _unique_columns(columns)
-    for col_name in unique_columns:
-        # wrap the column in a struct to make `is_unique` work with list/arrays
-        # https://github.com/pola-rs/polars/issues/27286
-        rules[f"{col_name}|unique"] = Rule(pl.struct(col_name).is_unique())
-
     # Add column-specific rules
     column_rules = {
         f"{col_name}|{rule_name}": Rule(expr)
@@ -79,10 +72,6 @@ def _primary_key(columns: dict[str, Column]) -> list[str]:
     return list(k for k, col in columns.items() if col.primary_key)
 
 
-def _unique_columns(columns: dict[str, Column]) -> list[str]:
-    return list(k for k, col in columns.items() if col.unique)
-
-
 # ------------------------------------------------------------------------------------ #
 #                                      SCHEMA META                                     #
 # ------------------------------------------------------------------------------------ #
@@ -315,11 +304,6 @@ def primary_key(cls) -> list[str]:
         """The primary key columns in this schema (possibly empty)."""
         return _primary_key(cls.columns())
 
-    @classmethod
-    def unique_columns(cls) -> list[str]:
-        """The columns with unique constraints in this schema (possibly empty)."""
-        return _unique_columns(cls.columns())
-
     @classmethod
     def _validation_rules(cls, *, with_cast: bool) -> dict[str, Rule]:
         return _build_rules(
diff --git a/dataframely/columns/_base.py b/dataframely/columns/_base.py
@@ -56,9 +56,9 @@ def __init__(
                 Explicitly set `nullable=True` if you want your column to be nullable.
             primary_key: Whether this column is part of the primary key of the schema.
                 If `True`, `nullable` is automatically set to `False`.
-            unique: Whether this column must contain unique values. Unlike ``primary_key``,
+            unique: Whether this column must contain unique values. Unlike `primary_key`,
                 this checks uniqueness for this column independently. Multiple columns
-                can each have ``unique=True`` without forming a composite constraint.
+                can each have `unique=True` without forming a composite constraint.
             check: A custom rule or multiple rules to run for this column. This can be:
                 - A single callable that returns a non-aggregated boolean expression.
                 The name of the rule is derived from the callable name, or defaults to
@@ -132,6 +132,9 @@ def validation_rules(self, expr: pl.Expr) -> dict[str, pl.Expr]:
         if not self.nullable:
             result["nullability"] = expr.is_not_null()
 
+        if self.unique:
+            result["unique"] = expr.is_unique()
+
         if self.check is not None:
             if isinstance(self.check, Mapping):
                 for rule_name, rule_callable in self.check.items():
diff --git a/dataframely/columns/array.py b/dataframely/columns/array.py
@@ -46,7 +46,9 @@ def __init__(
             shape: The shape of the array.
             nullable: Whether this column may contain null values.
             primary_key: Whether this column is part of the primary key of the schema.
-            unique: Whether this column must contain unique values.
+            unique: Whether this column must contain unique values. Unlike `primary_key`,
+                this checks uniqueness for this column independently. Multiple columns
+                can each have `unique=True` without forming a composite constraint.
             check: A custom rule or multiple rules to run for this column. This can be:
                 - A single callable that returns a non-aggregated boolean expression.
                 The name of the rule is derived from the callable name, or defaults to
@@ -90,6 +92,10 @@ def validation_rules(self, expr: pl.Expr) -> dict[str, pl.Expr]:
         array_rules: dict[str, pl.Expr] = {}
         if (rule := _list_primary_key_check(expr.arr, self.inner)) is not None:
             array_rules["primary_key"] = rule
+        if self.unique:
+            # Wrap the column in a struct to make `is_unique` work with arrays:
+            # https://github.com/pola-rs/polars/issues/27286
+            array_rules["unique"] = pl.struct(expr).is_unique()
 
         return {
             **super().validation_rules(expr),
diff --git a/dataframely/columns/categorical.py b/dataframely/columns/categorical.py
@@ -36,7 +36,9 @@ def __init__(
                 is not specified.
             primary_key: Whether this column is part of the primary key of the schema.
                 If `True`, `nullable` is automatically set to `False`.
-            unique: Whether this column must contain unique values.
+            unique: Whether this column must contain unique values. Unlike `primary_key`,
+                this checks uniqueness for this column independently. Multiple columns
+                can each have `unique=True` without forming a composite constraint.
             check: A custom rule or multiple rules to run for this column. This can be:
                 - A single callable that returns a non-aggregated boolean expression.
                 The name of the rule is derived from the callable name, or defaults to
diff --git a/dataframely/columns/datetime.py b/dataframely/columns/datetime.py
@@ -55,7 +55,9 @@ def __init__(
                 is not specified.
             primary_key: Whether this column is part of the primary key of the schema.
                 If `True`, `nullable` is automatically set to `False`.
-            unique: Whether this column must contain unique values.
+            unique: Whether this column must contain unique values. Unlike `primary_key`,
+                this checks uniqueness for this column independently. Multiple columns
+                can each have `unique=True` without forming a composite constraint.
             min: The minimum date for dates in this column (inclusive).
             min_exclusive: Like `min` but exclusive. May not be specified if `min`
                 is specified and vice versa.
@@ -191,7 +193,9 @@ def __init__(
                 is not specified.
             primary_key: Whether this column is part of the primary key of the schema.
                 If `True`, `nullable` is automatically set to `False`.
-            unique: Whether this column must contain unique values.
+            unique: Whether this column must contain unique values. Unlike `primary_key`,
+                this checks uniqueness for this column independently. Multiple columns
+                can each have `unique=True` without forming a composite constraint.
             min: The minimum time for times in this column (inclusive).
             min_exclusive: Like `min` but exclusive. May not be specified if `min`
                 is specified and vice versa.
@@ -335,7 +339,9 @@ def __init__(
                 is not specified.
             primary_key: Whether this column is part of the primary key of the schema.
                 If `True`, `nullable` is automatically set to `False`.
-            unique: Whether this column must contain unique values.
+            unique: Whether this column must contain unique values. Unlike `primary_key`,
+                this checks uniqueness for this column independently. Multiple columns
+                can each have `unique=True` without forming a composite constraint.
             min: The minimum datetime for datetimes in this column (inclusive).
             min_exclusive: Like `min` but exclusive. May not be specified if `min`
                 is specified and vice versa.
@@ -500,7 +506,9 @@ def __init__(
                 is not specified.
             primary_key: Whether this column is part of the primary key of the schema.
                 If `True`, `nullable` is automatically set to `False`.
-            unique: Whether this column must contain unique values.
+            unique: Whether this column must contain unique values. Unlike `primary_key`,
+                this checks uniqueness for this column independently. Multiple columns
+                can each have `unique=True` without forming a composite constraint.
             min: The minimum duration for durations in this column (inclusive).
             min_exclusive: Like `min` but exclusive. May not be specified if `min`
                 is specified and vice versa.
diff --git a/dataframely/columns/decimal.py b/dataframely/columns/decimal.py
@@ -48,7 +48,9 @@ def __init__(
                 is not specified.
             primary_key: Whether this column is part of the primary key of the schema.
                 If `True`, `nullable` is automatically set to `False`.
-            unique: Whether this column must contain unique values.
+            unique: Whether this column must contain unique values. Unlike `primary_key`,
+                this checks uniqueness for this column independently. Multiple columns
+                can each have `unique=True` without forming a composite constraint.
             min: The minimum value for decimals in this column (inclusive).
             min_exclusive: Like `min` but exclusive. May not be specified if `min`
                 is specified and vice versa.
diff --git a/dataframely/columns/enum.py b/dataframely/columns/enum.py
@@ -43,7 +43,9 @@ def __init__(
                 is not specified.
             primary_key: Whether this column is part of the primary key of the schema.
                 If `True`, `nullable` is automatically set to `False`.
-            unique: Whether this column must contain unique values.
+            unique: Whether this column must contain unique values. Unlike `primary_key`,
+                this checks uniqueness for this column independently. Multiple columns
+                can each have `unique=True` without forming a composite constraint.
             check: A custom rule or multiple rules to run for this column. This can be:
                 - A single callable that returns a non-aggregated boolean expression.
                 The name of the rule is derived from the callable name, or defaults to
diff --git a/dataframely/columns/float.py b/dataframely/columns/float.py
@@ -48,7 +48,9 @@ def __init__(
                 is not specified.
             primary_key: Whether this column is part of the primary key of the schema.
                 If `True`, `nullable` is automatically set to `False`.
-            unique: Whether this column must contain unique values.
+            unique: Whether this column must contain unique values. Unlike `primary_key`,
+                this checks uniqueness for this column independently. Multiple columns
+                can each have `unique=True` without forming a composite constraint.
             allow_inf: Whether this column may contain infinity values.
             allow_nan: Whether this column may contain NaN values.
             min: The minimum value for floats in this column (inclusive).
diff --git a/dataframely/columns/integer.py b/dataframely/columns/integer.py
@@ -44,7 +44,9 @@ def __init__(
                 is not specified.
             primary_key: Whether this column is part of the primary key of the schema.
                 If `True`, `nullable` is automatically set to `False`.
-            unique: Whether this column must contain unique values.
+            unique: Whether this column must contain unique values. Unlike `primary_key`,
+                this checks uniqueness for this column independently. Multiple columns
+                can each have `unique=True` without forming a composite constraint.
             min: The minimum value for integers in this column (inclusive).
             min_exclusive: Like `min` but exclusive. May not be specified if `min`
                 is specified and vice versa.
diff --git a/dataframely/columns/list.py b/dataframely/columns/list.py
@@ -54,7 +54,9 @@ def __init__(
                 In a future release, `nullable=False` will be the default if `nullable`
                 is not specified.
             primary_key: Whether this column is part of the primary key of the schema.
-            unique: Whether this column must contain unique values.
+            unique: Whether this column must contain unique values. Unlike `primary_key`,
+                this checks uniqueness for this column independently. Multiple columns
+                can each have `unique=True` without forming a composite constraint.
             check: A custom rule or multiple rules to run for this column. This can be:
                 - A single callable that returns a non-aggregated boolean expression.
                 The name of the rule is derived from the callable name, or defaults to
@@ -104,6 +106,10 @@ def validation_rules(self, expr: pl.Expr) -> dict[str, pl.Expr]:
         list_rules: dict[str, pl.Expr] = {}
         if (rule := _list_primary_key_check(expr.list, self.inner)) is not None:
             list_rules["primary_key"] = rule
+        if self.unique:
+            # Wrap the column in a struct to make `is_unique` work with lists:
+            # https://github.com/pola-rs/polars/issues/27286
+            list_rules["unique"] = pl.struct(expr).is_unique()
         if self.min_length is not None:
             list_rules["min_length"] = (
                 pl.when(expr.is_null())
diff --git a/dataframely/columns/string.py b/dataframely/columns/string.py
@@ -41,7 +41,9 @@ def __init__(
                 In a future release, `nullable=False` will be the default if `nullable`
                 is not specified.
             primary_key: Whether this column is part of the primary key of the schema.
-            unique: Whether this column must contain unique values.
+            unique: Whether this column must contain unique values. Unlike `primary_key`,
+                this checks uniqueness for this column independently. Multiple columns
+                can each have `unique=True` without forming a composite constraint.
             min_length: The minimum byte-length of string values in this column.
             max_length: The maximum byte-length of string values in this column.
             regex: A regex that the string values in this column must match. If the
diff --git a/dataframely/columns/struct.py b/dataframely/columns/struct.py
@@ -47,7 +47,9 @@ def __init__(
                 In a future release, `nullable=False` will be the default if `nullable`
                 is not specified.
             primary_key: Whether this column is part of the primary key of the schema.
-            unique: Whether this column must contain unique values.
+            unique: Whether this column must contain unique values. Unlike `primary_key`,
+                this checks uniqueness for this column independently. Multiple columns
+                can each have `unique=True` without forming a composite constraint.
             check: A custom rule or multiple rules to run for this column. This can be:
                 - A single callable that returns a non-aggregated boolean expression.
                 The name of the rule is derived from the callable name, or defaults to
diff --git a/tests/column_types/test_array.py b/tests/column_types/test_array.py
@@ -179,3 +179,23 @@ def test_outer_nullability() -> None:
     )
     df = pl.DataFrame({"nullable": [None, None]})
     schema.validate(df, cast=True)
+
+
+def test_unique() -> None:
+    schema = create_schema("test", {"a": dy.Array(dy.Integer(), shape=2, unique=True)})
+    df = pl.DataFrame(
+        {"a": [[1, 1], [1, 2], [1, 1], [1, 3]]}, schema={"a": pl.Array(pl.Int64, 2)}
+    )
+    _, failure = schema.filter(df)
+    assert failure.counts() == {"a|unique": 2}
+    assert validation_mask(df, failure).to_list() == [False, True, False, True]
+
+
+def test_inner_unique() -> None:
+    schema = create_schema("test", {"a": dy.Array(dy.Integer(unique=True), shape=2)})
+    df = pl.DataFrame(
+        {"a": [[1, 2], [1, 1], [2, 2], [1, 4]]}, schema={"a": pl.Array(pl.Int64, 2)}
+    )
+    _, failure = schema.filter(df)
+    assert failure.counts() == {"a|inner_unique": 2}
+    assert validation_mask(df, failure).to_list() == [True, False, False, True]
diff --git a/tests/column_types/test_list.py b/tests/column_types/test_list.py
@@ -192,3 +192,19 @@ def test_list_sampling_with_min_length(min_length: int) -> None:
     # Verify all lists have at least min_length elements
     min_list_len = cast(int, df["a"].list.len().min())
     assert min_list_len >= min_length
+
+
+def test_unique() -> None:
+    schema = create_schema("test", {"a": dy.List(dy.Integer(), unique=True)})
+    df = pl.DataFrame({"a": [[1, 1], [1, 2], [1, 1], [1, 1, 1]]})
+    _, failure = schema.filter(df)
+    assert failure.counts() == {"a|unique": 2}
+    assert validation_mask(df, failure).to_list() == [False, True, False, True]
+
+
+def test_inner_unique() -> None:
+    schema = create_schema("test", {"a": dy.List(dy.Integer(unique=True))})
+    df = pl.DataFrame({"a": [[1, 2, 3], [1, 1, 2], [1, 1], [1, 4]]})
+    _, failure = schema.filter(df)
+    assert failure.counts() == {"a|inner_unique": 2}
+    assert validation_mask(df, failure).to_list() == [True, False, False, True]
diff --git a/tests/schema/test_validate.py b/tests/schema/test_validate.py
@@ -289,10 +289,3 @@ def test_multiple_unique_columns_both_invalid(
     ):
         _validate_and_collect(MultiUniqueSchema, df, eager=eager)
     assert not MultiUniqueSchema.is_valid(df)
-
-
-def test_unique_columns_method() -> None:
-    assert UniqueSchema.unique_columns() == ["email"]
-    assert MultiUniqueSchema.unique_columns() == ["a", "b"]
-    assert NullableUniqueSchema.unique_columns() == ["b"]
-    assert MySchema.unique_columns() == []
