Refactor relational queries (#1131)

* fix test_filter_by_coordinate_system_also_table from @giovp

* remove unused _filter_table_in_coordinate_systems(); replace assert

* Refactor _filter_table_by_elements to delegate to join_spatialelement_table

- Replace the bespoke numpy implementation in _filter_table_by_elements with
  a call to join_spatialelement_table(how="left"), removing the unused
  match_rows parameter and consolidating to a single code path.
- Fix join functions (_left_join, _inner_join, _right_exclusive_join) to
  update spatialdata_attrs region metadata after filtering.
- Fix _right_exclusive_join: restore groupby pattern (consistent with other
  joins) and use reset_index so integer positions are used as the group index,
  avoiding IndexError when obs names are duplicated and fixing a latent bug
  where pd.concat of per-group masks produced a partial-length boolean mask.
- Simplify _get_filtered_or_unfiltered_tables to use _filter_table_by_elements.
- Replace SpatialData.init_from_elements with get_model() lookup when
  sdata=None in join_spatialelement_table, removing an expensive importlib call.
- Add test_join_updates_spatialdata_attrs covering metadata update for all join
  types; update test_filter_by_table_query_edge_cases to match new behaviour.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* Remove _filter_table_by_element_names and replace with _filter_table_by_elements

The old function filtered only by region name, ignoring instance IDs.
Replace its only caller (filter_tables by coordinate system) with
_filter_table_by_elements, which correctly filters by both region and instance.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: LucaMarconato

src/spatialdata/_core/concatenate.py

@@ -219,14 +219,6 @@ def concatenate(
     return sdata
 
 
-def _filter_table_in_coordinate_systems(table: AnnData, coordinate_systems: list[str]) -> AnnData:
-    table_mapping_metadata = table.uns[TableModel.ATTRS_KEY]
-    region_key = table_mapping_metadata[TableModel.REGION_KEY_KEY]
-    new_table = table[table.obs[region_key].isin(coordinate_systems)].copy()
-    new_table.uns[TableModel.ATTRS_KEY][TableModel.REGION_KEY] = new_table.obs[region_key].unique().tolist()
-    return new_table
-
-
 def _fix_ensure_unique_element_names(
     sdatas: dict[str, SpatialData],
     rename_tables: bool,

src/spatialdata/_core/query/_utils.py

@@ -226,7 +226,6 @@ def _get_filtered_or_unfiltered_tables(
         return {
             name: filtered_table
             for name, table in sdata.tables.items()
-            if (filtered_table := _filter_table_by_elements(table, elements)) and len(filtered_table) != 0
+            if (filtered_table := _filter_table_by_elements(table, elements)) is not None
         }
-
     return sdata.tables

src/spatialdata/_core/query/relational_query.py

@@ -59,31 +59,6 @@ def get_element_annotators(sdata: SpatialData, element_name: str) -> set[str]:
     return table_names
 
 
-def _filter_table_by_element_names(table: AnnData | None, element_names: str | list[str]) -> AnnData | None:
-    """
-    Filter an AnnData table to keep only the rows that are in the coordinate system.
-
-    Parameters
-    ----------
-    table
-        The table to filter; if None, returns None
-    element_names
-        The element_names to keep in the tables obs.region column
-
-    Returns
-    -------
-    The filtered table, or None if the input table was None
-    """
-    if table is None or not table.uns.get(TableModel.ATTRS_KEY):
-        return None
-    table_mapping_metadata = table.uns[TableModel.ATTRS_KEY]
-    region_key = table_mapping_metadata[TableModel.REGION_KEY_KEY]
-    table.obs = pd.DataFrame(table.obs)
-    table = table[table.obs[region_key].isin(element_names)].copy()
-    table.uns[TableModel.ATTRS_KEY][TableModel.REGION_KEY] = table.obs[region_key].unique().tolist()
-    return table
-
-
 @singledispatch
 def get_element_instances(
     element: SpatialElement,
@@ -113,7 +88,8 @@ def _(
     return_background: bool = False,
 ) -> pd.Index:
     model = get_model(element)
-    assert model in [Labels2DModel, Labels3DModel], "Expected a `Labels` element. Found an `Image` instead."
+    if model not in [Labels2DModel, Labels3DModel]:
+        raise ValueError("Expected a `Labels` element. Found an `Image` instead.")
     if isinstance(element, DataArray):
         # get unique labels value (including 0 if present)
         instances = da.unique(element.data).compute()
@@ -144,88 +120,43 @@ def _(
     return element.index
 
 
-# TODO: replace function use throughout repo by `join_sdata_spatialelement_table`
-# TODO: benchmark against join operations before removing
-def _filter_table_by_elements(
-    table: AnnData | None, elements_dict: dict[str, dict[str, Any]], match_rows: bool = False
-) -> AnnData | None:
+def _filter_table_by_elements(table: AnnData | None, elements_dict: dict[str, dict[str, Any]]) -> AnnData | None:
     """
-    Filter an AnnData table to keep only the rows that are in the elements.
+    Filter an AnnData table to keep only the rows annotating elements in elements_dict.
 
     Parameters
     ----------
     table
         The table to filter; if None, returns None
     elements_dict
-        The elements to use to filter the table
-    match_rows
-        If True, reorder the table rows to match the order of the elements
+        The elements to use to filter the table, structured as ``{element_type: {name: element}}``.
+        Image elements are ignored since tables cannot annotate images.
 
     Returns
     -------
-    The filtered table (eventually with reordered rows), or None if the input table was None.
+    The filtered table, or None if the input table is None or no rows match.
     """
-    assert set(elements_dict.keys()).issubset({"images", "labels", "shapes", "points"})
-    assert len(elements_dict) > 0, "elements_dict must not be empty"
-    assert any(len(elements) > 0 for elements in elements_dict.values()), (
-        "elements_dict must contain at least one dict which contains at least one element"
-    )
     if table is None:
         return None
-    to_keep = np.zeros(len(table), dtype=bool)
-    region_key = table.uns[TableModel.ATTRS_KEY][TableModel.REGION_KEY_KEY]
-    instance_key = table.uns[TableModel.ATTRS_KEY][TableModel.INSTANCE_KEY]
-    instances = None
-    for _, elements in elements_dict.items():
-        for name, element in elements.items():
-            if get_model(element) == Labels2DModel or get_model(element) == Labels3DModel:
-                if isinstance(element, DataArray):
-                    # get unique labels value (including 0 if present)
-                    instances = da.unique(element.data).compute()
-                else:
-                    assert isinstance(element, DataTree)
-                    v = element["scale0"].values()
-                    assert len(v) == 1
-                    xdata = next(iter(v))
-                    # can be slow
-                    instances = da.unique(xdata.data).compute()
-                instances = np.sort(instances)
-            elif get_model(element) == ShapesModel:
-                instances = element.index.to_numpy()
-            elif get_model(element) == PointsModel:
-                instances = element.compute().index.to_numpy()
-            else:
-                continue
-            indices = ((table.obs[region_key] == name) & (table.obs[instance_key].isin(instances))).to_numpy()
-            to_keep = to_keep | indices
-    original_table = table
-    table.obs = pd.DataFrame(table.obs)
-    table = table[to_keep, :]
-    if match_rows:
-        assert instances is not None
-        assert isinstance(instances, np.ndarray)
-        assert np.sum(to_keep) != 0, "No row matches in the table annotates the element"
-        if np.sum(to_keep) != len(instances):
-            if len(elements_dict) > 1 or len(elements_dict) == 1 and len(next(iter(elements_dict.values()))) > 1:
-                raise NotImplementedError("Sorting is not supported when filtering by multiple elements")
-            # case in which the instances in the table and the instances in the element don't correspond
-            assert "element" in locals()
-            assert "name" in locals()
-            n0 = np.setdiff1d(instances, table.obs[instance_key].to_numpy())
-            n1 = np.setdiff1d(table.obs[instance_key].to_numpy(), instances)
-            assert len(n1) == 0, f"The table contains {len(n1)} instances that are not in the element: {n1}"
-            # some instances have not a corresponding row in the table
-            instances = np.setdiff1d(instances, n0)
-        assert np.sum(to_keep) == len(instances)
-        assert sorted(set(instances.tolist())) == sorted(set(table.obs[instance_key].tolist()))  # type: ignore[type-var]
-        table_df = pd.DataFrame({instance_key: table.obs[instance_key], "position": np.arange(len(instances))})
-        merged = pd.merge(table_df, pd.DataFrame(index=instances), left_on=instance_key, right_index=True, how="right")
-        matched_positions = merged["position"].to_numpy()
-        table = table[matched_positions, :]
-    table = table.copy()
-    _inplace_fix_subset_categorical_obs(subset_adata=table, original_adata=original_table)
-    table.uns[TableModel.ATTRS_KEY][TableModel.REGION_KEY] = table.obs[region_key].unique().tolist()
-    return table
+    elements_by_name = {
+        name: element
+        for element_type, name_to_element in elements_dict.items()
+        if element_type != "images"
+        for name, element in name_to_element.items()
+    }
+    if not elements_by_name:
+        return None
+    # Suppress "element not annotated by table" warnings: the table may annotate
+    # only a subset of the elements passed in, which is expected here.
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", UserWarning)
+        _, filtered = join_spatialelement_table(
+            spatial_element_names=list(elements_by_name.keys()),
+            spatial_elements=list(elements_by_name.values()),
+            table=table,
+            how="left",
+        )
+    return filtered if filtered is not None and len(filtered) > 0 else None
 
 
 def _get_joined_table_indices(
@@ -253,7 +184,7 @@ def _get_joined_table_indices(
     -------
         The indices that of the table that match the SpatialElement indices.
     """
-    mask = table_instance_key_column.isin(element_indices)
+    mask = np.isin(table_instance_key_column.values, element_indices)
     if joined_indices is None:
         if match_rows == "left":
             _, joined_indices = _match_rows(table_instance_key_column, mask, element_indices, match_rows)
@@ -294,7 +225,7 @@ def _get_masked_element(
     -------
     The masked spatial element based on the provided indices and match rows.
     """
-    mask = table_instance_key_column.isin(element_indices)
+    mask = np.isin(table_instance_key_column.values, element_indices)
     masked_table_instance_key_column = table_instance_key_column[mask]
     mask_values = mask_values if len(mask_values := masked_table_instance_key_column.values) != 0 else None
     if match_rows in ["left", "right"]:
@@ -320,8 +251,11 @@ def _right_exclusive_join_spatialelement_table(
     regions, region_column_name, instance_key = get_table_keys(table)
     if isinstance(regions, str):
         regions = [regions]
-    groups_df = table.obs.groupby(by=region_column_name, observed=False)
-    mask = []
+    # reset_index so group_df.index gives integer positions — safe with duplicate obs names
+    obs = table.obs.reset_index()
+    groups_df = obs.groupby(by=region_column_name, observed=False)
+    keep = np.zeros(len(table), dtype=bool)
+    has_match = False
     for element_type, name_element in element_dict.items():
         for name, element in name_element.items():
             if name in regions:
@@ -331,24 +265,23 @@ def _right_exclusive_join_spatialelement_table(
                     element_indices = element.index
                 else:
                     element_indices = get_element_instances(element)
-
-                element_dict[element_type][name] = None
                 submask = ~table_instance_key_column.isin(element_indices)
-                mask.append(submask)
+                keep[group_df.index[submask.values]] = True
+                has_match = True
+                element_dict[element_type][name] = None
             else:
                 warnings.warn(
                     f"The element `{name}` is not annotated by the table. Skipping", UserWarning, stacklevel=2
                 )
                 element_dict[element_type][name] = None
                 continue
 
-    if len(mask) != 0:
-        mask = pd.concat(mask)
-        exclusive_table = table[mask, :].copy() if mask.sum() != 0 else None  # type: ignore[attr-defined]
-    else:
-        exclusive_table = None
-
+    exclusive_table = table[keep, :] if has_match and keep.any() else None
     _inplace_fix_subset_categorical_obs(subset_adata=exclusive_table, original_adata=table)
+    if exclusive_table is not None:
+        exclusive_table.uns[TableModel.ATTRS_KEY][TableModel.REGION_KEY] = (
+            exclusive_table.obs[region_column_name].unique().tolist()
+        )
     return element_dict, exclusive_table
 
 
@@ -449,6 +382,10 @@ def _inner_join_spatialelement_table(
     joined_table = table[joined_indices.tolist(), :].copy() if joined_indices is not None else None
 
     _inplace_fix_subset_categorical_obs(subset_adata=joined_table, original_adata=table)
+    if joined_table is not None:
+        joined_table.uns[TableModel.ATTRS_KEY][TableModel.REGION_KEY] = (
+            joined_table.obs[region_column_name].unique().tolist()
+        )
     return element_dict, joined_table
 
 
@@ -528,7 +465,10 @@ def _left_join_spatialelement_table(
             joined_indices = joined_indices.astype(int)
     joined_table = table[joined_indices.tolist(), :].copy() if joined_indices is not None else None
     _inplace_fix_subset_categorical_obs(subset_adata=joined_table, original_adata=table)
-
+    if joined_table is not None:
+        joined_table.uns[TableModel.ATTRS_KEY][TableModel.REGION_KEY] = (
+            joined_table.obs[region_column_name].unique().tolist()
+        )
     return element_dict, joined_table
 
 
@@ -723,11 +663,16 @@ def join_spatialelement_table(
     if sdata is not None:
         elements_dict = _create_sdata_elements_dict_for_join(sdata, spatial_element_names)
     else:
-        derived_sdata = SpatialData.init_from_elements(dict(zip(spatial_element_names, spatial_elements, strict=True)))
-        element_types = ["labels", "shapes", "points"]
+        _model_to_type = {
+            Labels2DModel: "labels",
+            Labels3DModel: "labels",
+            ShapesModel: "shapes",
+            PointsModel: "points",
+        }
         elements_dict = defaultdict(lambda: defaultdict(dict))
-        for element_type in element_types:
-            for name, element in getattr(derived_sdata, element_type).items():
+        for name, element in zip(spatial_element_names, spatial_elements, strict=True):
+            element_type = _model_to_type.get(get_model(element))
+            if element_type is not None:
                 elements_dict[element_type][name] = element
 
     elements_dict_joined, table = _call_join(elements_dict, table, how, match_rows, filter_label_pixels)

src/spatialdata/_core/spatialdata.py

@@ -691,12 +691,17 @@ def _filter_tables(
                     continue
                 # each mode here requires paths or elements, using assert here to avoid mypy errors.
                 if by == "cs":
-                    from spatialdata._core.query.relational_query import (
-                        _filter_table_by_element_names,
-                    )
+                    from spatialdata._core.query.relational_query import _filter_table_by_elements
 
                     assert element_names is not None
-                    table = _filter_table_by_element_names(table, element_names)
+                    elements_dict = {}
+                    for element_type in ["images", "labels", "shapes", "points"]:
+                        elements = getattr(self, element_type)
+                        if elements:  # Check if the dictionary is not empty
+                            elements_dict[element_type] = {
+                                name: elements[name] for name in element_names if name in elements
+                            }
+                    table = _filter_table_by_elements(table, elements_dict=elements_dict)
                     if table is not None and len(table) != 0:
                         tables[table_name] = table
                 elif by == "elements":

tests/core/operations/test_spatialdata_operations.py

@@ -156,11 +156,14 @@ def test_filter_by_coordinate_system(full_sdata: SpatialData) -> None:
 def test_filter_by_coordinate_system_also_table(full_sdata: SpatialData) -> None:
     from spatialdata.models import TableModel
 
-    rng = np.random.default_rng(seed=0)
-    full_sdata["table"].obs["annotated_shapes"] = pd.Categorical(
-        rng.choice(["circles", "poly"], size=full_sdata["table"].shape[0])
-    )
-    adata = full_sdata["table"]
+    adata = full_sdata["table"].copy()
+
+    circles_instances = full_sdata["circles"].index.values
+    poly_instances = full_sdata["poly"].index.values
+
+    adata = adata[: len(circles_instances) + len(poly_instances), :].copy()
+    adata.obs["annotated_shapes"] = ["circles"] * len(circles_instances) + ["poly"] * len(poly_instances)
+    adata.obs["instance_id"] = np.concatenate([circles_instances, poly_instances])
     del adata.uns[TableModel.ATTRS_KEY]
     full_sdata["table"] = TableModel.parse(
         adata,

tests/core/query/test_relational_query.py

@@ -246,6 +246,41 @@ def test_left_inner_right_exclusive_join(sdata_query_aggregation):
     assert element_dict["by_polygons"] is None
 
 
+def test_join_updates_spatialdata_attrs(sdata_query_aggregation):
+    sdata = sdata_query_aggregation
+    # table annotates ["values_circles", "values_polygons"]
+    original_regions = sdata["table"].uns["spatialdata_attrs"]["region"]
+    assert set(original_regions) == {"values_circles", "values_polygons"}
+
+    # left join on a single element: region list must shrink to just that element
+    _, table = join_spatialelement_table(
+        sdata=sdata, spatial_element_names="values_circles", table_name="table", how="left"
+    )
+    assert table.uns["spatialdata_attrs"]["region"] == ["values_circles"]
+
+    # inner join on a single element
+    _, table = join_spatialelement_table(
+        sdata=sdata, spatial_element_names="values_circles", table_name="table", how="inner"
+    )
+    assert table.uns["spatialdata_attrs"]["region"] == ["values_circles"]
+
+    # right_exclusive join: pass a truncated circles element so some table rows have no match.
+    # values_circles has 9 instances (0-8); keep only 5 → 4 table rows are exclusive.
+    # Use sdata=None mode so we can pass a truncated element under the original region name.
+    _, table = join_spatialelement_table(
+        spatial_element_names=["values_circles"],
+        spatial_elements=[sdata["values_circles"].iloc[:5]],
+        table=sdata["table"],
+        how="right_exclusive",
+    )
+    assert table is not None
+    assert table.n_obs == 4
+    assert table.uns["spatialdata_attrs"]["region"] == ["values_circles"]
+
+    # original table metadata must be unchanged
+    assert set(sdata["table"].uns["spatialdata_attrs"]["region"]) == {"values_circles", "values_polygons"}
+
+
 def test_join_spatialelement_table_fail(full_sdata):
     with pytest.raises(ValueError, match=" not supported for join operation."):
         join_spatialelement_table(
@@ -1214,7 +1249,7 @@ def test_filter_by_table_query_edge_cases(complex_sdata):
         assert var_name.startswith("feature_") and int(var_name.split("_")[1]) < 5
 
     # Test 6: Invalid element_names (element doesn't exist)
-    with pytest.raises(AssertionError, match="elements_dict must not be empty"):
+    with pytest.raises(KeyError, match="shapes_table"):
         filter_by_table_query(
             sdata=sdata,
             table_name="shapes_table",