From cda0aec10808b2e2677909a58f0c7f02812d9bcb Mon Sep 17 00:00:00 2001 From: Oliver Borchert Date: Tue, 15 Sep 2026 17:02:01 -0700 Subject: [PATCH] fix: synthesize category names for nested columns --- dataframely/_base_schema.py | 6 ++-- dataframely/columns/_base.py | 5 ++++ dataframely/columns/array.py | 6 ++++ dataframely/columns/categorical.py | 3 ++ dataframely/columns/list.py | 6 ++++ dataframely/columns/struct.py | 7 +++++ tests/column_types/test_categorical.py | 38 ++++++++++++++++++++------ 7 files changed, 59 insertions(+), 12 deletions(-) diff --git a/dataframely/_base_schema.py b/dataframely/_base_schema.py index 204c67d..766a74b 100644 --- a/dataframely/_base_schema.py +++ b/dataframely/_base_schema.py @@ -190,8 +190,7 @@ def __getattribute__(cls, name: str) -> Any: # 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 + val._bind(f"{cls.__module__}:{cls.__name__}", val.alias or name) return val @staticmethod @@ -294,8 +293,7 @@ def columns(cls) -> dict[str, Column]: columns: dict[str, Column] = getattr(cls, _COLUMN_ATTR) for name in columns.keys(): # Dynamically set the name and source schema of the columns. - columns[name]._schema = f"{cls.__module__}:{cls.__name__}" - columns[name]._name = name + columns[name]._bind(f"{cls.__module__}:{cls.__name__}", name) return columns @classmethod diff --git a/dataframely/columns/_base.py b/dataframely/columns/_base.py index cc1e403..abf4d26 100644 --- a/dataframely/columns/_base.py +++ b/dataframely/columns/_base.py @@ -276,6 +276,11 @@ def _pydantic_field_kwargs(self) -> dict[str, Any]: # ------------------------------------ HELPER ------------------------------------ # + def _bind(self, schema: str, name: str) -> None: + """Set the schema and column path used by schema-scoped data types.""" + self._schema = schema + self._name = name + @property def name(self) -> str: """Get the name of the column in a schema.""" diff --git a/dataframely/columns/array.py b/dataframely/columns/array.py index 60480c7..00a7de0 100644 --- a/dataframely/columns/array.py +++ b/dataframely/columns/array.py @@ -5,6 +5,7 @@ import math import warnings +from copy import copy from typing import Any, cast import polars as pl @@ -75,6 +76,11 @@ def __init__( self.inner = inner self.shape = shape if isinstance(shape, tuple) else (shape,) + def _bind(self, schema: str, name: str) -> None: + super()._bind(schema, name) + self.inner = copy(self.inner) + self.inner._bind(schema, f"{name}.inner") + @property def dtype(self) -> pl.DataType: return pl.Array(self.inner.dtype, self.shape) diff --git a/dataframely/columns/categorical.py b/dataframely/columns/categorical.py index 0d89ce0..aa4f5fc 100644 --- a/dataframely/columns/categorical.py +++ b/dataframely/columns/categorical.py @@ -39,6 +39,9 @@ def __init__( a data type is provided, name and namespace are synthesized from the enclosing schema and column name, automatically creating a column- scoped categories dictionary. + List and array elements append `.inner` to the containing column's + name; struct fields append their field names with dots (for example, + `items.inner.kind` for a field in a list of structs). 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` diff --git a/dataframely/columns/list.py b/dataframely/columns/list.py index 73b7552..ab2d116 100644 --- a/dataframely/columns/list.py +++ b/dataframely/columns/list.py @@ -3,6 +3,7 @@ from __future__ import annotations +from copy import copy from itertools import chain from typing import Any, cast @@ -85,6 +86,11 @@ def __init__( self.min_length = min_length self.max_length = max_length + def _bind(self, schema: str, name: str) -> None: + super()._bind(schema, name) + self.inner = copy(self.inner) + self.inner._bind(schema, f"{name}.inner") + @property def dtype(self) -> pl.DataType: return pl.List(self.inner.dtype) diff --git a/dataframely/columns/struct.py b/dataframely/columns/struct.py index eca3b26..e0d323c 100644 --- a/dataframely/columns/struct.py +++ b/dataframely/columns/struct.py @@ -3,6 +3,7 @@ from __future__ import annotations +from copy import copy from typing import Any, cast import polars as pl @@ -76,6 +77,12 @@ def __init__( ) self.inner = inner + def _bind(self, schema: str, name: str) -> None: + super()._bind(schema, name) + self.inner = {field: copy(col) for field, col in self.inner.items()} + for field, col in self.inner.items(): + col._bind(schema, f"{name}.{field}") + @property def dtype(self) -> pl.DataType: return pl.Struct({name: col.dtype for name, col in self.inner.items()}) diff --git a/tests/column_types/test_categorical.py b/tests/column_types/test_categorical.py index 1d65f42..0945331 100644 --- a/tests/column_types/test_categorical.py +++ b/tests/column_types/test_categorical.py @@ -7,18 +7,40 @@ import pytest import dataframely as dy +from dataframely._polars import PolarsDataType from dataframely.testing.factory import create_schema -def test_synthesized_categories_name() -> None: +@pytest.mark.parametrize( + ("column", "expected_name"), + [ + (dy.Categorical(pl.UInt16), "a"), + (dy.List(dy.Categorical(pl.UInt16)), "a.inner"), + (dy.Array(dy.Categorical(pl.UInt16), 2), "a.inner"), + (dy.Struct({"x": dy.Categorical(pl.UInt16)}), "a.x"), + ( + dy.List(dy.Array(dy.List(dy.Categorical(pl.UInt16)), 2)), + "a.inner.inner.inner", + ), + ( + dy.List(dy.Struct({"x": dy.List(dy.Categorical(pl.UInt16))})), + "a.inner.x.inner", + ), + ], +) +def test_synthesized_categories_name(column: dy.Column, expected_name: str) -> 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" - ) + a = column + + for bound_column in (TestSchema.a, TestSchema.columns()["a"]): + dtype: PolarsDataType = bound_column.dtype + while isinstance(dtype, pl.List | pl.Array | pl.Struct): + dtype = ( + dtype.fields[0].dtype if isinstance(dtype, pl.Struct) else dtype.inner + ) + categories = cast(pl.Categorical, dtype).categories + assert categories.name() == expected_name + assert categories.namespace() == f"{__name__}:TestSchema" @pytest.mark.parametrize(