diff --git a/src/vsparse/_anndata_class.py b/src/vsparse/_anndata_class.py index 69e479a..58e0ca2 100644 --- a/src/vsparse/_anndata_class.py +++ b/src/vsparse/_anndata_class.py @@ -245,10 +245,15 @@ def _write_group( g: Any, *, format: str = "vcsc", + convert_strings_to_categoricals: bool = True, dataset_kwargs: Mapping[str, Any] = MappingProxyType({}), ) -> None: if format not in _STORE_FORMATS: raise ValueError(f"format must be one of {_STORE_FORMATS}, got {format!r}") + if convert_strings_to_categoricals: + # Writing field-by-field skips what anndata's own writers do here, + # leaving low-cardinality columns as one string per row. + self.strings_to_categoricals() write_array = ad.io.write_elem if format == "vcsc" else _io.write_ivcs_elem if self._vcs_X is not None: write_array(g, "X", self._vcs_X, dataset_kwargs=dataset_kwargs) @@ -277,6 +282,7 @@ def write_h5ad( # ty: ignore[invalid-method-override] filename: str | PathLike[str], *, format: str = "vcsc", + convert_strings_to_categoricals: bool = True, dataset_kwargs: Mapping[str, Any] | None = None, **_kwargs: Any, ) -> None: @@ -291,6 +297,10 @@ def write_h5ad( # ty: ignore[invalid-method-override] the cost of extra work on write/read. Either way, ``X``/``raw_X`` come back from :meth:`read_h5ad` as ordinary VCSCArray/VCSRArray objects -- ``"ivcsc"`` is purely an on-disk storage format. + convert_strings_to_categoricals + Convert ``obs``/``var`` string columns to categorical in place + before writing, as ``anndata``'s own writers do. Only columns + with fewer categories than rows are converted. dataset_kwargs Passed to ``h5py.Group.create_dataset`` for every array written. Defaults to Blosc2+LZ4 compression; pass ``{}`` to store @@ -305,7 +315,12 @@ def write_h5ad( # ty: ignore[invalid-method-override] _compression.numeric_only_compression("h5"), h5py.File(filename, "w") as f, ): - self._write_group(f, format=format, dataset_kwargs=dataset_kwargs) + self._write_group( + f, + format=format, + convert_strings_to_categoricals=convert_strings_to_categoricals, + dataset_kwargs=dataset_kwargs, + ) @classmethod def read_h5ad(cls, filename: str | PathLike[str]) -> VCSCAnnData: @@ -320,14 +335,15 @@ def write_zarr( store: Any, *, format: str = "vcsc", + convert_strings_to_categoricals: bool = True, dataset_kwargs: Mapping[str, Any] | None = None, **_kwargs: Any, ) -> None: """Write to a zarr store. Read back with :meth:`read_zarr`. - See :meth:`write_h5ad` for ``format``/``dataset_kwargs`` (including the - numeric-only compression behavior); the default compression here is - Blosc+LZ4 via ``numcodecs``. + See :meth:`write_h5ad` for ``format``/``convert_strings_to_categoricals``/ + ``dataset_kwargs`` (including the numeric-only compression behavior); + the default compression here is Blosc+LZ4 via ``numcodecs``. """ import zarr @@ -335,7 +351,12 @@ def write_zarr( dataset_kwargs = _compression.zarr_dataset_kwargs() with _compression.numeric_only_compression("zarr"): f = zarr.open_group(store, mode="w") - self._write_group(f, format=format, dataset_kwargs=dataset_kwargs) + self._write_group( + f, + format=format, + convert_strings_to_categoricals=convert_strings_to_categoricals, + dataset_kwargs=dataset_kwargs, + ) @classmethod def read_zarr(cls, store: Any) -> VCSCAnnData: diff --git a/src/vsparse/_compression.py b/src/vsparse/_compression.py index d88a9d3..22c116d 100644 --- a/src/vsparse/_compression.py +++ b/src/vsparse/_compression.py @@ -11,11 +11,17 @@ ``_index``). At least some HDF5 filter-plugin builds (seen in this environment: h5py 3.16 / HDF5 2.0.0 / hdf5plugin's Blosc2) segfault (``SIGFPE``) when the Blosc2 filter is applied to a variable-length-string -dataset -- and there's no benefit to compressing already-tiny label arrays -anyway. :func:`numeric_only_compression` patches the relevant +dataset. :func:`numeric_only_compression` patches the relevant ``create_dataset``/``create_array`` calls for the duration of a write so string/object arrays always land uncompressed, regardless of what ``dataset_kwargs`` was passed in -- callers don't have to know about this. + +Strings are left uncompressed rather than given a different codec: gzip and +lzf are safe on variable-length strings here, but the size problem is the +per-row string itself, not its compression. Encoding low-cardinality +columns as categoricals turns the per-row data numeric, which the existing +Blosc2 path then compresses -- see +:meth:`vsparse.VCSCAnnData.write_h5ad`'s ``convert_strings_to_categoricals``. """ from __future__ import annotations diff --git a/tests/test_metadata_encoding.py b/tests/test_metadata_encoding.py new file mode 100644 index 0000000..94740c2 --- /dev/null +++ b/tests/test_metadata_encoding.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +import anndata as ad +import numpy as np +import pandas as pd +import pytest +import scipy.sparse as sp + +from vsparse import VCSCAnnData, VCSRArray + + +def _adata(n_cells: int = 400, n_genes: int = 20) -> ad.AnnData: + rng = np.random.default_rng(0) + obs = pd.DataFrame( + { + "cell_type": [f"type_{i}" for i in rng.integers(0, 5, n_cells)], + "sample_id": [f"SAMPLE_{i:03d}" for i in rng.integers(0, 12, n_cells)], + "total_counts": rng.normal(1000, 50, n_cells), + }, + index=[f"cell_{i:05d}" for i in range(n_cells)], + ) + var = pd.DataFrame( + { + "gene_symbol": [f"GENE{i:04d}" for i in range(n_genes)], # unique per gene + "chromosome": [f"chr{i}" for i in rng.integers(1, 5, n_genes)], # repeats + }, + index=[f"ENSG{i:05d}" for i in range(n_genes)], + ) + X = sp.random_array((n_cells, n_genes), density=0.2, format="csr", random_state=0) + X.data = np.round(X.data * 8 + 1) + return ad.AnnData(X=X, obs=obs, var=var) + + +@pytest.mark.parametrize("fmt", ["vcsc", "ivcsc"]) +def test_string_columns_are_written_as_categoricals(tmp_path, fmt): + va = VCSCAnnData.from_anndata(_adata(), format="csr") + assert not isinstance(va.obs["cell_type"].dtype, pd.CategoricalDtype) + + path = tmp_path / f"data.{fmt}.h5ad" + va.write_h5ad(path, format=fmt) + back = VCSCAnnData.read_h5ad(path) + + assert isinstance(back.obs["cell_type"].dtype, pd.CategoricalDtype) + assert isinstance(back.obs["sample_id"].dtype, pd.CategoricalDtype) + assert isinstance(back.var["chromosome"].dtype, pd.CategoricalDtype) + # gene_symbol is unique per gene: nothing to gain, so it's left alone. + assert not isinstance(back.var["gene_symbol"].dtype, pd.CategoricalDtype) + + +def test_values_survive_the_conversion(tmp_path): + original = _adata() + va = VCSCAnnData.from_anndata(original, format="csr") + path = tmp_path / "data.h5ad" + va.write_h5ad(path) + back = VCSCAnnData.read_h5ad(path) + + for col in ("cell_type", "sample_id"): + pd.testing.assert_series_equal( + back.obs[col].astype(str), original.obs[col].astype(str), check_names=False + ) + np.testing.assert_allclose(back.obs["total_counts"], original.obs["total_counts"]) + np.testing.assert_array_equal(back.obs_names, original.obs_names) + np.testing.assert_array_equal(back.var_names, original.var_names) + assert isinstance(back.X, VCSRArray) + np.testing.assert_allclose(back.X.toarray(), sp.csr_array(original.X).toarray()) + + +def test_conversion_can_be_turned_off(tmp_path): + va = VCSCAnnData.from_anndata(_adata(), format="csr") + path = tmp_path / "raw_strings.h5ad" + va.write_h5ad(path, convert_strings_to_categoricals=False) + back = VCSCAnnData.read_h5ad(path) + + assert not isinstance(back.obs["cell_type"].dtype, pd.CategoricalDtype) + assert not isinstance(va.obs["cell_type"].dtype, pd.CategoricalDtype) # not mutated + + +def test_categorical_encoding_shrinks_the_file(tmp_path): + """Categorical codes compress where per-row strings cannot.""" + va_plain = VCSCAnnData.from_anndata(_adata(n_cells=4000), format="csr") + va_cat = VCSCAnnData.from_anndata(_adata(n_cells=4000), format="csr") + + plain = tmp_path / "plain.h5ad" + cat = tmp_path / "cat.h5ad" + va_plain.write_h5ad(plain, convert_strings_to_categoricals=False) + va_cat.write_h5ad(cat, convert_strings_to_categoricals=True) + + assert cat.stat().st_size < plain.stat().st_size + + +def test_high_cardinality_columns_are_left_alone(tmp_path): + """A column with a distinct value per row is left alone.""" + adata = _adata(n_cells=100) + adata.obs["barcode"] = [f"barcode_{i}" for i in range(adata.n_obs)] + va = VCSCAnnData.from_anndata(adata, format="csr") + + path = tmp_path / "unique.h5ad" + va.write_h5ad(path) + back = VCSCAnnData.read_h5ad(path) + + assert not isinstance(back.obs["barcode"].dtype, pd.CategoricalDtype) + np.testing.assert_array_equal(back.obs["barcode"], adata.obs["barcode"]) + + +def test_zarr_write_converts_too(tmp_path): + va = VCSCAnnData.from_anndata(_adata(), format="csr") + store = tmp_path / "data.zarr" + va.write_zarr(store) + back = VCSCAnnData.read_zarr(store) + + assert isinstance(back.obs["cell_type"].dtype, pd.CategoricalDtype) + np.testing.assert_array_equal(back.obs["cell_type"].astype(str), va.obs["cell_type"].astype(str))