diff --git a/dataframely/_base_schema.py b/dataframely/_base_schema.py index e6ed611..e7ab585 100644 --- a/dataframely/_base_schema.py +++ b/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 diff --git a/dataframely/columns/_base.py b/dataframely/columns/_base.py index 79fb10b..cc1e403 100644 --- a/dataframely/columns/_base.py +++ b/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 = "" diff --git a/dataframely/columns/categorical.py b/dataframely/columns/categorical.py index db24599..0d89ce0 100644 --- a/dataframely/columns/categorical.py +++ b/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() diff --git a/tests/column_types/test_categorical.py b/tests/column_types/test_categorical.py new file mode 100644 index 0000000..1d65f42 --- /dev/null +++ b/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)