Repository: scverse/spatialdata

Issue:
`concatenate` does not account for custom region keys in the Table
When concatenating multiple tables from `spatialdata.aggregrate`, an error is thrown when using a custom region key.

Reproducible example:

```python
import pathlib
import spatialdata as sd
import natsort as ns

mibitof_zarr_path = pathlib.Path("~/Downloads/data.zarr").expanduser()


mibitof_sdata = sd.read_zarr(mibitof_zarr_path)
mibitof_sdata

# Delete the table
del mibitof_sdata.table

images = ns.natsorted(list(mibitof_sdata.images.keys()))
labels = ns.natsorted(list(mibitof_sdata.labels.keys()))
coords = ns.natsorted(list(mibitof_sdata.coordinate_systems))

agg_sd = []

for image, label, coord in zip(images, labels, coords):
    agg_sd.append(
        sd.aggregate(
            values=image,
            by=label,
            values_sdata=mibitof_sdata,
            by_sdata=mibitof_sdata,
            agg_func="sum",
            region_key="fov_id",
            instance_key="cell_id",
            target_coordinate_system=coord,
        )
    )

sd.concatenate(agg_sd, region_key="fov_id", instance_key="cell_id")
```


<details>
  <summary>Traceback </summary>
  
  ```shell
KeyError                                  Traceback (most recent call last)
File [~/.pyenv/versions/3.11.4/envs/ark-spatial/lib/python3.11/site-packages/pandas/core/indexes/base.py:3653](https://file+.vscode-resource.vscode-cdn.net/Users/srivarra/Dev/Python/ark/notebooks/~/.pyenv/versions/3.11.4/envs/ark-spatial/lib/python3.11/site-packages/pandas/core/indexes/base.py:3653), in Index.get_loc(self, key)
   3652 try:
-> 3653     return self._engine.get_loc(casted_key)
   3654 except KeyError as err:

File [~/.pyenv/versions/3.11.4/envs/ark-spatial/lib/python3.11/site-packages/pandas/_libs/index.pyx:147](https://file+.vscode-resource.vscode-cdn.net/Users/srivarra/Dev/Python/ark/notebooks/~/.pyenv/versions/3.11.4/envs/ark-spatial/lib/python3.11/site-packages/pandas/_libs/index.pyx:147), in pandas._libs.index.IndexEngine.get_loc()

File [~/.pyenv/versions/3.11.4/envs/ark-spatial/lib/python3.11/site-packages/pandas/_libs/index.pyx:176](https://file+.vscode-resource.vscode-cdn.net/Users/srivarra/Dev/Python/ark/notebooks/~/.pyenv/versions/3.11.4/envs/ark-spatial/lib/python3.11/site-packages/pandas/_libs/index.pyx:176), in pandas._libs.index.IndexEngine.get_loc()

File pandas/_libs/hashtable_class_helper.pxi:7080, in pandas._libs.hashtable.PyObjectHashTable.get_item()

File pandas/_libs/hashtable_class_helper.pxi:7088, in pandas._libs.hashtable.PyObjectHashTable.get_item()

KeyError: 'region'

The above exception was the direct cause of the following exception:

KeyError                                  Traceback (most recent call last)
Cell In[25], line 1
----> 1 sd.concatenate(agg_sd, region_key="fov_id", instance_key="cell_id")

File [~/.pyenv/versions/3.11.4/envs/ark-spatial/lib/python3.11/site-packages/spatialdata/_core/concatenate.py:115](https://file+.vscode-resource.vscode-cdn.net/Users/srivarra/Dev/Python/ark/notebooks/~/.pyenv/versions/3.11.4/envs/ark-spatial/lib/python3.11/site-packages/spatialdata/_core/concatenate.py:115), in concatenate(sdatas, region_key, instance_key, **kwargs)
    112 assert type(sdatas) == list, "sdatas must be a list"
    113 assert len(sdatas) > 0, "sdatas must be a non-empty list"
--> 115 merged_table = _concatenate_tables(
    116     [sdata.table for sdata in sdatas if sdata.table is not None], region_key, instance_key, **kwargs
    117 )
    119 return SpatialData(
    120     images=merged_images,
    121     labels=merged_labels,
   (...)
    124     table=merged_table,
    125 )

File [~/.pyenv/versions/3.11.4/envs/ark-spatial/lib/python3.11/site-packages/spatialdata/_core/concatenate.py:63](https://file+.vscode-resource.vscode-cdn.net/Users/srivarra/Dev/Python/ark/notebooks/~/.pyenv/versions/3.11.4/envs/ark-spatial/lib/python3.11/site-packages/spatialdata/_core/concatenate.py:63), in _concatenate_tables(tables, region_key, instance_key, **kwargs)
     59     tables_l.append(table)
     61 merged_table = ad.concat(tables_l, **kwargs)
     62 attrs = {
---> 63     TableModel.REGION_KEY: merged_table.obs[TableModel.REGION_KEY].unique().tolist(),
     64     TableModel.REGION_KEY_KEY: region_key,
     65     TableModel.INSTANCE_KEY: instance_key,
     66 }
     67 merged_table.uns[TableModel.ATTRS_KEY] = attrs
     69 return TableModel().validate(merged_table)

File [~/.pyenv/versions/3.11.4/envs/ark-spatial/lib/python3.11/site-packages/pandas/core/frame.py:3761](https://file+.vscode-resource.vscode-cdn.net/Users/srivarra/Dev/Python/ark/notebooks/~/.pyenv/versions/3.11.4/envs/ark-spatial/lib/python3.11/site-packages/pandas/core/frame.py:3761), in DataFrame.__getitem__(self, key)
   3759 if self.columns.nlevels > 1:
   3760     return self._getitem_multilevel(key)
-> 3761 indexer = self.columns.get_loc(key)
   3762 if is_integer(indexer):
   3763     indexer = [indexer]

File [~/.pyenv/versions/3.11.4/envs/ark-spatial/lib/python3.11/site-packages/pandas/core/indexes/base.py:3655](https://file+.vscode-resource.vscode-cdn.net/Users/srivarra/Dev/Python/ark/notebooks/~/.pyenv/versions/3.11.4/envs/ark-spatial/lib/python3.11/site-packages/pandas/core/indexes/base.py:3655), in Index.get_loc(self, key)
   3653     return self._engine.get_loc(casted_key)
   3654 except KeyError as err:
-> 3655     raise KeyError(key) from err
   3656 except TypeError:
   3657     # If we have a listlike key, _check_indexing_error will raise
   3658     #  InvalidIndexError. Otherwise we fall through and re-raise
   3659     #  the TypeError.
   3660     self._check_indexing_error(key)

KeyError: 'region'
  ```
</details>

I believe the cause of the issue is here:

https://github.com/scverse/spatialdata/blob/054a7bf3b4cbc70f451a9a124a715b79d30be9e0/src/spatialdata/_core/concatenate.py#L62-L66

on line 63, where the `obs` is indexed with `TableModel.REGION_KEY` which is `"region"` instead of the function parameter `region_key`, which would be `"fov_id"` in this case. I'd like to contribute with a fix if this is not the expected behavior.

Discussion:
@LucaMarconato: @srivarra thank you for reporting and sorry for the late answer, the thesis writing got in between and I missed some GitHub notifications. I will look into this and report back.@srivarra: @LucaMarconato No worries, hope your thesis is going well! Let me know if I can contribute with a PR if it's an actual issue!@LucaMarconato: Hi @srivarra, this issue fell outside my radar. I made a quick fix in the linked PR.

Solution patch:
diff --git a/src/spatialdata/_core/concatenate.py b/src/spatialdata/_core/concatenate.py
index 6b91e22..58e278f 100644
--- a/src/spatialdata/_core/concatenate.py
+++ b/src/spatialdata/_core/concatenate.py
@@ -65,7 +65,7 @@ def _concatenate_tables(
 
     merged_table = ad.concat(tables_l, **kwargs)
     attrs = {
-        TableModel.REGION_KEY: merged_table.obs[TableModel.REGION_KEY].unique().tolist(),
+        TableModel.REGION_KEY: merged_table.obs[region_key].unique().tolist(),
         TableModel.REGION_KEY_KEY: region_key,
         TableModel.INSTANCE_KEY: instance_key,
     }


Test patch:
diff --git a/tests/core/operations/test_spatialdata_operations.py b/tests/core/operations/test_spatialdata_operations.py
index d4644d4..9b669fc 100644
--- a/tests/core/operations/test_spatialdata_operations.py
+++ b/tests/core/operations/test_spatialdata_operations.py
@@ -20,8 +20,14 @@ from spatialdata.models import (
     TableModel,
     get_table_keys,
 )
-from spatialdata.testing import assert_elements_dict_are_identical, assert_spatial_data_objects_are_identical
-from spatialdata.transformations.operations import get_transformation, set_transformation
+from spatialdata.testing import (
+    assert_elements_dict_are_identical,
+    assert_spatial_data_objects_are_identical,
+)
+from spatialdata.transformations.operations import (
+    get_transformation,
+    set_transformation,
+)
 from spatialdata.transformations.transformations import (
     Affine,
     BaseTransformation,
@@ -30,7 +36,7 @@ from spatialdata.transformations.transformations import (
     Sequence,
     Translation,
 )
-from tests.conftest import _get_table
+from tests.conftest import _get_shapes, _get_table
 
 
 def test_element_names_unique() -> None:
@@ -183,7 +189,10 @@ def test_filter_by_coordinate_system_also_table(full_sdata: SpatialData) -> None
     del adata.uns[TableModel.ATTRS_KEY]
     del full_sdata.tables["table"]
     full_sdata.table = TableModel.parse(
-        adata, region=["circles", "poly"], region_key="annotated_shapes", instance_key="instance_id"
+        adata,
+        region=["circles", "poly"],
+        region_key="annotated_shapes",
+        instance_key="instance_id",
     )
 
     scale = Scale([2.0], axes=("x",))
@@ -201,11 +210,19 @@ def test_filter_by_coordinate_system_also_table(full_sdata: SpatialData) -> None
 def test_rename_coordinate_systems(full_sdata: SpatialData) -> None:
     # all the elements point to global, add new coordinate systems
     set_transformation(
-        element=full_sdata.shapes["circles"], transformation=Identity(), to_coordinate_system="my_space0"
+        element=full_sdata.shapes["circles"],
+        transformation=Identity(),
+        to_coordinate_system="my_space0",
+    )
+    set_transformation(
+        element=full_sdata.shapes["poly"],
+        transformation=Identity(),
+        to_coordinate_system="my_space1",
     )
-    set_transformation(element=full_sdata.shapes["poly"], transformation=Identity(), to_coordinate_system="my_space1")
     set_transformation(
-        element=full_sdata.shapes["multipoly"], transformation=Identity(), to_coordinate_system="my_space2"
+        element=full_sdata.shapes["multipoly"],
+        transformation=Identity(),
+        to_coordinate_system="my_space2",
     )
 
     elements_in_global_before = {
@@ -234,7 +251,11 @@ def test_rename_coordinate_systems(full_sdata: SpatialData) -> None:
     # renaming, as it doesn't exist at the time of the function call)
     with pytest.raises(ValueError):
         full_sdata.rename_coordinate_systems(
-            {"my_space00": "my_space3", "my_space11": "my_space3", "my_space3": "my_space4"}
+            {
+                "my_space00": "my_space3",
+                "my_space11": "my_space3",
+                "my_space3": "my_space4",
+            }
         )
 
     # valid renaming with collisions
@@ -276,33 +297,79 @@ def test_concatenate_tables() -> None:
         "instance_key": "instance_id",
     }
 
-    table3 = _get_table(region="shapes/circles", region_key="annotated_shapes_other", instance_key="instance_id")
+    table3 = _get_table(
+        region="shapes/circles",
+        region_key="annotated_shapes_other",
+        instance_key="instance_id",
+    )
     with pytest.raises(ValueError):
         _concatenate_tables([table0, table3], region_key="region")
 
     table4 = _get_table(
-        region=["shapes/circles1", "shapes/poly1"], region_key="annotated_shape0", instance_key="instance_id"
+        region=["shapes/circles1", "shapes/poly1"],
+        region_key="annotated_shape0",
+        instance_key="instance_id",
     )
     table5 = _get_table(
-        region=["shapes/circles2", "shapes/poly2"], region_key="annotated_shape0", instance_key="instance_id"
+        region=["shapes/circles2", "shapes/poly2"],
+        region_key="annotated_shape0",
+        instance_key="instance_id",
     )
     table6 = _get_table(
-        region=["shapes/circles3", "shapes/poly3"], region_key="annotated_shape1", instance_key="instance_id"
+        region=["shapes/circles3", "shapes/poly3"],
+        region_key="annotated_shape1",
+        instance_key="instance_id",
     )
-    with pytest.raises(ValueError, match="`region_key` must be specified if tables have different region keys"):
+    with pytest.raises(
+        ValueError,
+        match="`region_key` must be specified if tables have different region keys",
+    ):
         _concatenate_tables([table4, table5, table6])
     assert len(_concatenate_tables([table4, table5, table6], region_key="region")) == len(table4) + len(table5) + len(
         table6
     )
 
 
+def test_concatenate_custom_table_metadata() -> None:
+    # test for https://github.com/scverse/spatialdata/issues/349
+    shapes0 = _get_shapes()
+    shapes1 = _get_shapes()
+    n = len(shapes0["poly"])
+    table0 = TableModel.parse(
+        AnnData(obs={"my_region": ["poly0"] * n, "my_instance_id": list(range(n))}),
+        region="poly0",
+        region_key="my_region",
+        instance_key="my_instance_id",
+    )
+    table1 = TableModel.parse(
+        AnnData(obs={"my_region": ["poly1"] * n, "my_instance_id": list(range(n))}),
+        region="poly1",
+        region_key="my_region",
+        instance_key="my_instance_id",
+    )
+    sdata0 = SpatialData.init_from_elements({"poly0": shapes0["poly"], "table": table0})
+    sdata1 = SpatialData.init_from_elements({"poly1": shapes1["poly"], "table": table1})
+    sdata = concatenate([sdata0, sdata1], concatenate_tables=True)
+    assert len(sdata["table"]) == 2 * n
+
+
 def test_concatenate_sdatas(full_sdata: SpatialData) -> None:
     with pytest.raises(KeyError):
         concatenate([full_sdata, SpatialData(images={"image2d": full_sdata.images["image2d"]})])
     with pytest.raises(KeyError):
-        concatenate([full_sdata, SpatialData(labels={"labels2d": full_sdata.labels["labels2d"]})])
+        concatenate(
+            [
+                full_sdata,
+                SpatialData(labels={"labels2d": full_sdata.labels["labels2d"]}),
+            ]
+        )
     with pytest.raises(KeyError):
-        concatenate([full_sdata, SpatialData(points={"points_0": full_sdata.points["points_0"]})])
+        concatenate(
+            [
+                full_sdata,
+                SpatialData(points={"points_0": full_sdata.points["points_0"]}),
+            ]
+        )
     with pytest.raises(KeyError):
         concatenate([full_sdata, SpatialData(shapes={"circles": full_sdata.shapes["circles"]})])
 
@@ -335,9 +402,15 @@ def test_concatenate_sdatas_from_iterable(concatenate_tables: bool, obs_names_ma
     sdatas = {"sample0": sdata0, "sample1": sdata1}
     with pytest.raises(KeyError, match="Images must have unique names across the SpatialData objects"):
         _ = concatenate(
-            sdatas.values(), concatenate_tables=concatenate_tables, obs_names_make_unique=obs_names_make_unique
+            sdatas.values(),
+            concatenate_tables=concatenate_tables,
+            obs_names_make_unique=obs_names_make_unique,
         )
-    merged = concatenate(sdatas, obs_names_make_unique=obs_names_make_unique, concatenate_tables=concatenate_tables)
+    merged = concatenate(
+        sdatas,
+        obs_names_make_unique=obs_names_make_unique,
+        concatenate_tables=concatenate_tables,
+    )
 
     if concatenate_tables:
         assert len(merged.tables) == 1
@@ -449,10 +522,18 @@ def test_subset(full_sdata: SpatialData) -> None:
 
     adata = AnnData(
         shape=(10, 0),
-        obs={"region": ["circles"] * 5 + ["poly"] * 5, "instance_id": [0, 1, 2, 3, 4, 0, 1, 2, 3, 4]},
+        obs={
+            "region": ["circles"] * 5 + ["poly"] * 5,
+            "instance_id": [0, 1, 2, 3, 4, 0, 1, 2, 3, 4],
+        },
     )
     del full_sdata.tables["table"]
-    sdata_table = TableModel.parse(adata, region=["circles", "poly"], region_key="region", instance_key="instance_id")
+    sdata_table = TableModel.parse(
+        adata,
+        region=["circles", "poly"],
+        region_key="region",
+        instance_key="instance_id",
+    )
     full_sdata["table"] = sdata_table
     full_sdata.tables["second_table"] = sdata_table
     subset1 = full_sdata.subset(["poly", "second_table"])
@@ -494,7 +575,12 @@ def test_transform_to_data_extent(full_sdata: SpatialData, maintain_positioning:
     points["z"] = points["x"]
     points = PointsModel.parse(points)
     full_sdata["points_0_3d"] = points
-    sdata = transform_to_data_extent(full_sdata, "global", target_width=1000, maintain_positioning=maintain_positioning)
+    sdata = transform_to_data_extent(
+        full_sdata,
+        "global",
+        target_width=1000,
+        maintain_positioning=maintain_positioning,
+    )
 
     first_a: ArrayLike | None = None
     for _, name, el in sdata.gen_spatial_elements():
@@ -505,7 +591,13 @@ def test_transform_to_data_extent(full_sdata: SpatialData, maintain_positioning:
             first_a = a
         else:
             # we are not pixel perfect because of this bug: https://github.com/scverse/spatialdata/issues/165
-            if maintain_positioning and name in ["points_0_3d", "points_0", "poly", "circles", "multipoly"]:
+            if maintain_positioning and name in [
+                "points_0_3d",
+                "points_0",
+                "poly",
+                "circles",
+                "multipoly",
+            ]:
                 # Again, due to the "pixel perfect" bug, the 0.5 translation forth and back in the z axis that is added
                 # by rasterize() (like the one in the example belows), amplifies the error also for x and y beyond the
                 # rtol threshold below. So, let's skip that check and to an absolute check up to 0.5 (due to the