improved memory footprint of streamings

Author: antonkulaga

src/cell2sentence4longevity/preprocessing/obs_stream.py

@@ -71,6 +71,39 @@ def _read_dataset_slice(
     return data
 
 
+git def _decode_string_array(values: np.ndarray) -> np.ndarray:
+    if values.dtype.kind in {"S", "O"}:
+        decoded = np.empty(len(values), dtype=object)
+        for idx, value in enumerate(values):
+            decoded[idx] = _decode_scalar(value)
+        return decoded
+    if values.dtype.kind == "U":
+        return values.astype(object)
+    return values
+
+
+def _read_nullable_group_slice(
+    field_node: h5py.Group,
+    field_name: str,
+    start_idx: int,
+    end_idx: int
+) -> np.ndarray:
+    if "values" not in field_node or "mask" not in field_node:
+        msg = f"Nullable obs field '{field_name}' is missing required datasets"
+        raise ValueError(msg)
+    values = field_node["values"][start_idx:end_idx]
+    mask = field_node["mask"][start_idx:end_idx]
+    values_array = _ensure_numpy_array(values)
+    if values_array.dtype.kind in {"S", "O", "U"}:
+        values_array = _decode_string_array(values_array)
+    mask_array = np.asarray(mask, dtype=bool)
+    if mask_array.size == 0 or not mask_array.any():
+        return values_array
+    result = values_array.astype(object, copy=True)
+    result[mask_array] = None
+    return result
+
+
 def _read_obs_field_slice(
     field_node: h5py.Dataset | h5py.Group,
     field_name: str,
@@ -87,6 +120,11 @@ def _read_obs_field_slice(
             categorical_cache[field_name] = _load_categorical_values(field_node)
         categories = categorical_cache[field_name]
         return _decode_categorical_codes(field_node, categories, start_idx, end_idx)
+    if encoding_type in {"nullable-boolean", "nullable-integer", "nullable-string-array"}:
+        if not isinstance(field_node, h5py.Group):
+            msg = f"Unexpected nullable node type for field '{field_name}'"
+            raise ValueError(msg)
+        return _read_nullable_group_slice(field_node, field_name, start_idx, end_idx)
     if not isinstance(field_node, h5py.Dataset):
         msg = f"Unsupported obs field storage for '{field_name}'"
         raise ValueError(msg)
@@ -244,6 +282,8 @@ def preload_complex_obs_fields(
         is_dataset = isinstance(node, h5py.Dataset)
         if is_dataset or encoding_type == "categorical":
             continue
+        if encoding_type in {"nullable-boolean", "nullable-integer", "nullable-string-array"}:
+            continue
         try:
             values = read_elem(node)
         except Exception as exc:
@@ -270,7 +310,7 @@ def infer_obs_schema(obs_group: h5py.Group) -> dict[str, pl.datatypes.DataType]:
         if name == "index":
             continue
         encoding_type = node.attrs.get("encoding-type", "array")
-        if encoding_type in {"categorical", "string-array"}:
+        if encoding_type in {"categorical", "string-array", "nullable-string-array"}:
             schema[name] = pl.String
             continue
         if encoding_type == "nullable-boolean":
@@ -279,9 +319,6 @@ def infer_obs_schema(obs_group: h5py.Group) -> dict[str, pl.datatypes.DataType]:
         if encoding_type == "nullable-integer":
             schema[name] = pl.Int64
             continue
-        if encoding_type == "nullable-string-array":
-            schema[name] = pl.String
-            continue
         if not isinstance(node, h5py.Dataset):
             continue
         kind = node.dtype.kind

tests/test_obs_stream_nullable.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import anndata as ad
+import h5py
+import numpy as np
+import pandas as pd
+import polars as pl
+import pytest
+
+from cell2sentence4longevity.preprocessing.obs_stream import (
+    build_obs_chunk_dataframe,
+    infer_obs_schema,
+    list_obs_columns_from_group,
+    preload_complex_obs_fields,
+)
+
+
+def test_nullable_obs_fields_stream_without_preload(tmp_path: Path) -> None:
+    """Ensure nullable obs columns stream chunk-by-chunk without eager preloading."""
+    original_setting = ad.settings.allow_write_nullable_strings
+    ad.settings.allow_write_nullable_strings = True
+    try:
+        n_cells = 6
+        obs = pd.DataFrame(
+            {
+                "nullable_int": pd.Series([1, None, 3, None, 5, 6], dtype="Int64"),
+                "nullable_bool": pd.Series([True, None, False, True, None, False], dtype="boolean"),
+            }
+        )
+        var = pd.DataFrame(index=[f"gene_{idx}" for idx in range(3)])
+        adata = ad.AnnData(X=np.zeros((n_cells, var.shape[0])), obs=obs, var=var)
+        h5ad_path = tmp_path / "nullable.h5ad"
+        adata.write_h5ad(h5ad_path)
+    finally:
+        ad.settings.allow_write_nullable_strings = original_setting
+
+    with h5py.File(h5ad_path, "r") as handle:
+        obs_group = handle["obs"]
+        fields = list_obs_columns_from_group(obs_group)
+        schema = infer_obs_schema(obs_group)
+        string_fields = {col for col, dtype in schema.items() if dtype == pl.String}
+        preloaded = preload_complex_obs_fields(obs_group, fields, total_rows=n_cells)
+        assert preloaded == {}
+
+        chunk_df = build_obs_chunk_dataframe(
+            obs_group=obs_group,
+            fields=fields,
+            start_idx=0,
+            end_idx=n_cells,
+            categorical_cache={},
+            string_fields=string_fields,
+            string_fill_value=None,
+            preloaded_fields=preloaded,
+        )
+    assert chunk_df["nullable_int"].to_list() == [1, None, 3, None, 5, 6]
+    assert chunk_df["nullable_bool"].to_list() == [True, None, False, True, None, False]
+