feat: Refine dy.Categorical to support pl.Categories

dataframely/_base_schema.py

@@ -202,7 +202,9 @@ def __new__(
         def __getattribute__(cls, name: str) -> Any:
             val = super().__getattribute__(name)
             # Dynamically set the name of the column if it is a `Column` instance.
+            # Also, we "register" the name of the schema that set the name.
             if isinstance(val, Column):
+                val._schema = f"{cls.__module__}:{cls.__name__}"
                 val._name = val.alias or name
             return val
 
@@ -305,7 +307,8 @@ def columns(cls) -> dict[str, Column]:
         """The column definitions of this schema."""
         columns: dict[str, Column] = getattr(cls, _COLUMN_ATTR)
         for name in columns.keys():
-            # Dynamically set the name of the columns.
+            # Dynamically set the name and source schema of the columns.
+            columns[name]._schema = f"{cls.__module__}:{cls.__name__}"
             columns[name]._name = name
         return columns
 

dataframely/columns/_base.py

@@ -92,6 +92,9 @@ def __init__(
         self.alias = alias
         self.metadata = metadata
         self.description = description
+
+        # The schema may be overridden by the schema on column access.
+        self._schema = ""
         # The name may be overridden by the schema on column access.
         self._name = ""
 

dataframely/columns/categorical.py

@@ -6,6 +6,7 @@
 from typing import Any
 
 import polars as pl
+from polars.datatypes import DataTypeClass
 
 from dataframely._compat import sa, sa_TypeEngine
 from dataframely.random import Generator
@@ -18,6 +19,7 @@ class Categorical(Column):
 
     def __init__(
         self,
+        categories: pl.Categories | pl.DataType | DataTypeClass | None = None,
         *,
         nullable: bool = False,
         primary_key: bool = False,
@@ -29,6 +31,14 @@ def __init__(
     ):
         """
         Args:
+            categories: An optional specification for how the categories for this
+                categorical are stored. If `None` is provided (default), the global
+                categories dictionary is used. When an instance of `pl.Categories` is
+                supplied, the categories are stored in the dictionary identified by
+                the name and namespace of the `pl.Categories` instance. When merely
+                a data type is provided, name and namespace are synthesized from the
+                enclosing schema and column name, automatically creating a column-
+                scoped categories dictionary.
             nullable: Whether this column may contain null values.
                 Explicitly set `nullable=True` if you want your column to be nullable.
                 In a future release, `nullable=False` will be the default if `nullable`
@@ -60,6 +70,14 @@ def __init__(
             metadata: A dictionary of metadata to attach to the column.
             description: A human-readable description of the column.
         """
+        if (
+            isinstance(categories, pl.DataType | DataTypeClass)
+            and categories != pl.UInt8
+            and categories != pl.UInt16
+            and categories != pl.UInt32
+        ):
+            raise ValueError("Category dtype must be one of [UInt8, UInt16, UInt32].")
+
         super().__init__(
             nullable=nullable,
             primary_key=primary_key,
@@ -69,10 +87,39 @@ def __init__(
             metadata=metadata,
             description=description,
         )
+        self.categories = categories
+
+    @property
+    def _categories(self) -> pl.Categories:
+        return self._resolve_categories(self.categories)
+
+    def _resolve_categories(
+        self, categories: pl.Categories | pl.DataType | DataTypeClass | None
+    ) -> pl.Categories:
+        if isinstance(categories, pl.Categories):
+            return categories
+        if isinstance(categories, pl.DataType | DataTypeClass):
+            return pl.Categories(
+                name=self._name,
+                namespace=self._schema,
+                physical=categories,
+            )
+        return pl.Categories()
 
     @property
     def dtype(self) -> pl.DataType:
-        return pl.Categorical()
+        return pl.Categorical(self._categories)
+
+    def _attributes_match(
+        self, lhs: Any, rhs: Any, name: str, column_expr: pl.Expr
+    ) -> bool:
+        if name == "categories":
+            # `categories` may be provided as `None`, a data type, or a
+            # `pl.Categories` instance. Compare the resolved categories so that
+            # equivalent specifications (e.g. `None` and the default global
+            # `pl.Categories`) are considered equal.
+            return self._resolve_categories(lhs) == self._resolve_categories(rhs)
+        return super()._attributes_match(lhs, rhs, name, column_expr)
 
     def sqlalchemy_dtype(self, dialect: sa.Dialect) -> sa_TypeEngine:
         return sa.String()

tests/column_types/test_categorical.py

@@ -0,0 +1,88 @@
+# Copyright (c) QuantCo 2025-2026
+# SPDX-License-Identifier: BSD-3-Clause
+
+from typing import cast
+
+import polars as pl
+import pytest
+
+import dataframely as dy
+from dataframely.testing.factory import create_schema
+
+
+def test_synthesized_categories_name() -> None:
+    class TestSchema(dy.Schema):
+        a = dy.Categorical(pl.UInt16)
+
+    assert cast(pl.Categorical, TestSchema.a.dtype).categories.name() == "a"
+    assert (
+        cast(pl.Categorical, TestSchema.a.dtype).categories.namespace()
+        == "column_types.test_categorical:TestSchema"
+    )
+
+
+@pytest.mark.parametrize(
+    ("column", "expected_dtype"),
+    [
+        (dy.Categorical(), pl.Categorical()),
+        (
+            dy.Categorical(pl.Categories("c", namespace="ns")),
+            pl.Categorical(pl.Categories("c", namespace="ns")),
+        ),
+    ],
+)
+def test_categories_dtype(column: dy.Categorical, expected_dtype: pl.DataType) -> None:
+    assert column.dtype == expected_dtype
+
+
+@pytest.mark.parametrize(
+    ("physical", "expected"),
+    [(pl.UInt8, pl.UInt8), (pl.UInt16, pl.UInt16), (pl.UInt32, pl.UInt32)],
+)
+def test_categories_physical(physical: pl.DataType, expected: pl.DataType) -> None:
+    schema = create_schema("test", {"a": dy.Categorical(physical)})
+    column = schema.columns()["a"]
+    assert isinstance(column, dy.Categorical)
+    assert column._categories.physical() == expected
+
+
+@pytest.mark.parametrize("physical", [pl.Int8, pl.Float64, pl.String])
+def test_categories_invalid_physical(physical: pl.DataType) -> None:
+    with pytest.raises(ValueError, match="Category dtype must be one of"):
+        dy.Categorical(physical)
+
+
+@pytest.mark.parametrize(
+    "column",
+    [
+        dy.Categorical(),
+        dy.Categorical(pl.Categories("c", namespace="ns", physical=pl.UInt16)),
+        dy.Categorical(pl.UInt16),
+    ],
+)
+@pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame])
+def test_valid(
+    df_type: type[pl.DataFrame] | type[pl.LazyFrame],
+    column: dy.Categorical,
+) -> None:
+    schema = create_schema("test", {"a": column})
+    df = df_type({"a": ["x", "y", "x"]}).cast(schema.columns()["a"].dtype)
+    assert schema.is_valid(df)
+
+
+def test_matches() -> None:
+    column = dy.Categorical(pl.Categories("c", physical=pl.UInt16))
+    expr = pl.col("a")
+    assert column.matches(dy.Categorical(pl.Categories("c", physical=pl.UInt16)), expr)
+    assert not column.matches(
+        dy.Categorical(pl.Categories("d", physical=pl.UInt16)), expr
+    )
+    assert not column.matches(
+        dy.Categorical(pl.Categories("c", physical=pl.UInt8)), expr
+    )
+    assert not column.matches(dy.Categorical(), expr)
+
+
+def test_matches_default() -> None:
+    expr = pl.col("a")
+    assert dy.Categorical().matches(dy.Categorical(), expr)