diff --git a/dataframely/columns/categorical.py b/dataframely/columns/categorical.py index aa4f5fc..dad676d 100644 --- a/dataframely/columns/categorical.py +++ b/dataframely/columns/categorical.py @@ -132,7 +132,10 @@ def _python_type(self) -> Any: return str def _sample_unchecked(self, generator: Generator, n: int) -> pl.Series: - # We simply sample low-cardinality strings here + # Two-letter strings allow 702 categories, exceeding UInt8's capacity. + max_length = 1 if self._categories.physical() == pl.UInt8 else 2 return generator.sample_string( - n, regex=r"[a-z]{1,2}", null_probability=self._null_probability + n, + regex=rf"[a-z]{{1,{max_length}}}", + null_probability=self._null_probability, ).cast(self.dtype) diff --git a/tests/columns/test_sample.py b/tests/columns/test_sample.py index 666516d..7b84cb4 100644 --- a/tests/columns/test_sample.py +++ b/tests/columns/test_sample.py @@ -175,6 +175,31 @@ def test_sample_enum(generator: Generator) -> None: assert set(samples) == {"a", "b", "c"} +@pytest.mark.parametrize("physical", [pl.UInt8, pl.UInt16, pl.UInt32]) +@pytest.mark.parametrize("explicit_categories", [False, True]) +@pytest.mark.parametrize("nullable", [False, True]) +@pytest.mark.parametrize("n", [0, 10_000]) +def test_sample_categorical( + physical: type[pl.DataType], + explicit_categories: bool, + nullable: bool, + n: int, + generator: Generator, +) -> None: + categories = ( + pl.Categories("sample", physical=physical) if explicit_categories else physical + ) + schema = create_schema("test", {"a": dy.Categorical(categories, nullable=nullable)}) + column = schema.columns()["a"] + samples = sample_and_validate(column, generator, n=n) + assert len(samples) == n + assert samples.dtype == column.dtype + assert samples.to_physical().dtype == physical + assert samples.drop_nulls().n_unique() <= (256 if physical == pl.UInt8 else 702) + if n: + assert samples.is_null().any() == nullable + + def test_sample_list(generator: Generator) -> None: column = dy.List( dy.String(regex="[abc]"), nullable=True, min_length=5, max_length=10