fix: lazy index.name not destroyed (#2358)

Author: ilan-gold

docs/release-notes/2358.fix.md

@@ -0,0 +1 @@
+Ensure the index name along `obs` and `var` are not lost when reading/writing to disk after reading with {func}`~anndata.experimental.read_lazy` {user}`ilan-gold`

src/anndata/_core/merge.py

@@ -1313,7 +1313,8 @@ def make_xarray_extension_dtypes_dask(
         )
 
 
-DS_CONCAT_DUMMY_INDEX_NAME = "concat_index"
+DS_CONCAT_DUMMY_INDEX_NAME = "_anndata_concat_index"
+DS_MERGE_DUMMY_INDEX_NAME = "_anndata_merge_index"
 
 
 def concat_dataset2d_on_annot_axis(
@@ -1725,7 +1726,7 @@ def concat(  # noqa: PLR0912, PLR0913, PLR0915
             if a.true_index_dim != a.index_dim:
                 a.index = a.true_index
         annotations_with_only_dask = [
-            a.ds.rename({a.true_index_dim: "merge_index"})
+            a.ds.rename({a.true_index_dim: DS_MERGE_DUMMY_INDEX_NAME})
             for a in annotations_with_only_dask
         ]
         alt_annot = Dataset2D(

src/anndata/_core/xarray.py

@@ -243,6 +243,11 @@ def to_memory(self, *, copy: bool = False) -> pd.DataFrame:
         -------
             :class:`pandas.DataFrame` with index set accordingly.
         """
+        from anndata._core.merge import (
+            DS_CONCAT_DUMMY_INDEX_NAME,
+            DS_MERGE_DUMMY_INDEX_NAME,
+        )
+
         index_key = self.ds.attrs.get("indexing_key", None)
         all_columns = {*self.columns, *([] if index_key is None else [index_key])}
         # https://github.com/pydata/xarray/issues/10419
@@ -258,7 +263,12 @@ def to_memory(self, *, copy: bool = False) -> pd.DataFrame:
             )
         if df.index.name != index_key and index_key is not None:
             df = df.set_index(index_key)
-        df.index.name = None  # matches old AnnData object
+        if df.index.name in {
+            "_index",
+            DS_CONCAT_DUMMY_INDEX_NAME,
+            DS_MERGE_DUMMY_INDEX_NAME,
+        }:
+            df.index.name = None  # matches old AnnData object
         return df
 
     @property

src/anndata/tests/helpers.py

@@ -1193,12 +1193,7 @@ class AccessTrackingStore(LocalStore):
     _accessed_keys: defaultdict[str, list[str]]
 
     def __init__(self, *args, **kwargs):
-        import traceback
-
-        traceback.print_stack()
-        print(kwargs)
         super().__init__(*args, **kwargs)
-        print(self._read_only)
         self._access_count = Counter()
         self._accessed = defaultdict(set)
         self._accessed_keys = defaultdict(list)

tests/lazy/conftest.py

@@ -124,7 +124,8 @@ def adata_remote_with_store_tall_skinny_path(
     orig_path = tmp_path_factory.mktemp(f"orig_{worker_id}.zarr")
     M = 1000
     N = 5
-    obs_names = pd.Index(f"cell{i}" for i in range(M))
+    # One named, one unnamed
+    obs_names = pd.Index((f"cell{i}" for i in range(M)), name="obs_names")
     var_names = pd.Index(f"gene{i}" for i in range(N))
     obs = gen_typed_df(M, obs_names)
     var = gen_typed_df(N, var_names)

tests/lazy/test_read.py

@@ -94,16 +94,16 @@ def test_access_count_index(
 ) -> None:
     adata_orig = read_zarr(adata_remote_with_store_tall_skinny_path)
 
-    remote_store_tall_skinny.initialize_key_trackers(["obs/_index"])
+    remote_store_tall_skinny.initialize_key_trackers(["obs/obs_names"])
     read_lazy(remote_store_tall_skinny, load_annotation_index=False)
-    remote_store_tall_skinny.assert_access_count("obs/_index", 0)
+    remote_store_tall_skinny.assert_access_count("obs/obs_names", 0)
 
     read_lazy(remote_store_tall_skinny)
     n_chunks = 4
     count_expected = (  # *2 when mask exists
         n_chunks * 2 if adata_orig.obs.index.dtype == "string" else n_chunks
     )
-    remote_store_tall_skinny.assert_access_count("obs/_index", count_expected)
+    remote_store_tall_skinny.assert_access_count("obs/obs_names", count_expected)
 
 
 def test_access_count_dtype(

tests/test_xarray.py

@@ -86,7 +86,6 @@ def test_true_index_dim_column_subset(dataset2d, df):
     df_expected = dataset2d[cols].to_memory()
     # account for the fact that we manually set `true_index_dim`
     df.index = df[col]
-    df.index.name = None
     pd.testing.assert_frame_equal(df_expected, df[cols])