feat: Refine dy.Categorical to support pl.Categories

dataframely/__init__.py

@@ -26,6 +26,7 @@
     Binary,
     Bool,
     Categorical,
+    Categories,
     Column,
     Date,
     Datetime,
@@ -82,6 +83,7 @@
     "Binary",
     "Bool",
     "Categorical",
+    "Categories",
     "Column",
     "Date",
     "Datetime",

dataframely/columns/__init__.py

@@ -7,7 +7,7 @@
 from .array import Array
 from .binary import Binary
 from .bool import Bool
-from .categorical import Categorical
+from .categorical import Categorical, Categories
 from .datetime import Date, Datetime, Duration, Time
 from .decimal import Decimal
 from .enum import Enum
@@ -26,6 +26,7 @@
     "Binary",
     "Bool",
     "Categorical",
+    "Categories",
     "Date",
     "Datetime",
     "Decimal",

dataframely/columns/categorical.py

@@ -3,7 +3,9 @@
 
 from __future__ import annotations
 
-from typing import Any
+import dataclasses
+from dataclasses import dataclass
+from typing import Any, Literal
 
 import polars as pl
 
@@ -14,12 +16,30 @@
 from ._registry import register
 
 
+@dataclass(frozen=True)
+class Categories:
+    """The name, namespace, and physical type of a categorical's global categories.
+
+    Mirrors :class:`polars.Categories`, but is immutable and serializable.
+    """
+
+    name: str | None = None
+    namespace: str = ""
+    physical: Literal["u8", "u16", "u32"] = "u32"
+
+    def to_polars(self) -> pl.Categories:
+        """Convert this object into a :class:`polars.Categories`."""
+        physical = {"u8": pl.UInt8, "u16": pl.UInt16, "u32": pl.UInt32}[self.physical]
+        return pl.Categories(self.name, namespace=self.namespace, physical=physical)
+
+
 @register
 class Categorical(Column):
     """A column of categorical (string) values."""
 
     def __init__(
         self,
+        categories: Categories | None = None,
         *,
         nullable: bool = False,
         primary_key: bool = False,
@@ -31,6 +51,8 @@ def __init__(
     ):
         """
         Args:
+            categories: The global categories (name, namespace, and physical index type)
+                for this column. If omitted, the default global categories are used.
             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`
@@ -71,10 +93,24 @@ def __init__(
             metadata=metadata,
             description=description,
         )
+        self.categories = categories
 
     @property
     def dtype(self) -> pl.DataType:
-        return pl.Categorical()
+        return pl.Categorical(self.categories.to_polars() if self.categories else None)
+
+    def as_dict(self, expr: pl.Expr) -> dict[str, Any]:
+        result = super().as_dict(expr)
+        if self.categories is not None:
+            result["categories"] = dataclasses.asdict(self.categories)
+        return result
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> Categorical:
+        data = dict(data)
+        if data.get("categories") is not None:
+            data["categories"] = Categories(**data["categories"])
+        return super().from_dict(data)
 
     def sqlalchemy_dtype(self, dialect: sa.Dialect) -> sa_TypeEngine:
         return sa.String()

tests/column_types/test_categorical.py

@@ -0,0 +1,80 @@
+# Copyright (c) QuantCo 2025-2026
+# SPDX-License-Identifier: BSD-3-Clause
+
+from typing import Literal
+
+import polars as pl
+import pytest
+
+import dataframely as dy
+from dataframely.testing.factory import create_schema
+
+
+@pytest.mark.parametrize(
+    "column, expected_dtype",
+    [
+        (dy.Categorical(), pl.Categorical()),
+        (
+            dy.Categorical(dy.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
+
+
+def test_categories_equality() -> None:
+    assert dy.Categories("c") == dy.Categories("c")
+    assert dy.Categories("c") != dy.Categories("d")
+    assert dy.Categories("c", physical="u8") != dy.Categories("c", physical="u16")
+
+
+@pytest.mark.parametrize(
+    ("physical", "expected"),
+    [("u8", pl.UInt8), ("u16", pl.UInt16), ("u32", pl.UInt32)],
+)
+def test_categories_to_polars_physical(
+    physical: Literal["u8", "u16", "u32"], expected: pl.DataType
+) -> None:
+    assert dy.Categories("c", physical=physical).to_polars().physical() == expected
+
+
+@pytest.mark.parametrize(
+    "column",
+    [dy.Categorical(), dy.Categorical(dy.Categories("c", namespace="ns"))],
+)
+@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(column.dtype)
+    assert schema.is_valid(df)
+
+
+def test_matches() -> None:
+    column = dy.Categorical(dy.Categories("c", physical="u16"))
+    expr = pl.col("a")
+    assert column.matches(dy.Categorical(dy.Categories("c", physical="u16")), expr)
+    assert not column.matches(dy.Categorical(dy.Categories("d", physical="u16")), expr)
+    assert not column.matches(dy.Categorical(dy.Categories("c", physical="u8")), expr)
+    assert not column.matches(dy.Categorical(), expr)
+
+
+@pytest.mark.parametrize(
+    "column",
+    [dy.Categorical(), dy.Categorical(dy.Categories("c", namespace="ns"))],
+)
+def test_as_dict_from_dict(column: dy.Categorical) -> None:
+    restored = dy.Categorical.from_dict(column.as_dict(pl.element()))
+    assert restored.categories == column.categories
+
+
+def test_schema_serialization_roundtrip() -> None:
+    schema = create_schema(
+        "test", {"a": dy.Categorical(dy.Categories("c", namespace="ns"))}
+    )
+    decoded = dy.deserialize_schema(schema.serialize())
+    assert schema.matches(decoded)

tests/schema/test_serialization.py

@@ -49,6 +49,10 @@ def test_simple_serialization() -> None:
             },
         ),
         create_schema("test", {"a": dy.Enum(["a"])}),
+        create_schema("test", {"a": dy.Categorical()}),
+        create_schema(
+            "test", {"a": dy.Categorical(dy.Categories("c", namespace="ns"))}
+        ),
         create_schema("test", {"a": dy.Decimal(scale=2, min=Decimal("1.5"))}),
         create_schema("test", {"a": dy.Date(min=dt.date(2020, 1, 1))}),
         create_schema("test", {"a": dy.Datetime(min=dt.datetime(2020, 1, 1))}),