inner/right joins fail with KeyError when the points element has more than one dask partition

[!NOTE]
This whole message is AI-generated. The issue was automatically discovered and reported by an AI agent (Claude) during an autonomous bug hunt on the spatialdata code base. It has not been verified or triaged by a human yet; the needs: triage label is set so that a maintainer can confirm it. The reproduction script below was executed by the agent in an isolated environment (see Environment) and its output is pasted verbatim.

Summary

mask_values is the global list of instance ids to keep and df.loc[list] is applied to every partition; pandas raises KeyError for labels absent from that partition, which is always the case with ≥ 2 partitions. join_spatialelement_table(how="inner"|"right"), match_element_to_table, match_sdata_to_table and filter_by_table_query (default how="right") all fail; the single-partition case works.

Severity (agent's assessment): medium — real points elements written by spatialdata-io are multi-partition (e.g. 8 partitions for Xenium transcripts)

Where: src/spatialdata/_core/query/relational_query.py::_get_masked_element (element.map_partitions(lambda df: df.loc[mask_values], meta=element))

Expected behaviour

Same result as for a single partition.

Reproduction

Save as repro.py and run uv run repro.py (the PEP 723 header pins spatialdata to the commit the bug was found on; replace the URL fragment with @main to test the current main branch).

# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "spatialdata @ git+https://github.com/scverse/spatialdata.git@ccf1ea048d054b6624214bf618008a9f9ae223e0",
# ]
# ///
"""inner/right joins fail with KeyError when the points element has more than one dask partition."""
import warnings
import numpy as np
import pandas as pd
import dask.dataframe as dd
from anndata import AnnData
from spatialdata import SpatialData, join_spatialelement_table
from spatialdata.models import PointsModel, TableModel

warnings.simplefilter("ignore")
df = pd.DataFrame({"x": np.arange(10.0), "y": np.arange(10.0)}, index=np.arange(10))
obs = pd.DataFrame({"region": pd.Categorical(["pts"] * 4), "instance_id": [1, 2, 7, 8]})
table = TableModel.parse(AnnData(X=np.zeros((4, 2)), obs=obs), region="pts", region_key="region", instance_key="instance_id")
bug = False
for npartitions in [1, 2]:
    points = PointsModel.parse(dd.from_pandas(df, npartitions=npartitions))
    sdata = SpatialData(points={"pts": points}, tables={"t": table})
    for how in ["inner", "right"]:
        try:
            elements, joined = join_spatialelement_table(sdata=sdata, spatial_element_names="pts", table_name="t", how=how)
            print(f"npartitions={npartitions} how={how:5s}: OK, element index {elements['pts'].compute().index.tolist()}")
        except Exception as e:  # noqa: BLE001
            print(f"npartitions={npartitions} how={how:5s}: {type(e).__name__}: {e}")
            bug |= npartitions > 1
print("VERDICT:", "BUG REPRODUCED" if bug else "NOT REPRODUCED")

Observed output

npartitions=1 how=inner: OK, element index [1, 2, 7, 8]
npartitions=1 how=right: OK, element index [1, 2, 7, 8]
npartitions=2 how=inner: KeyError: '[1, 2] not in index'
npartitions=2 how=right: KeyError: '[7, 8] not in index'
VERDICT: BUG REPRODUCED

Possible fix direction (unverified)

Filter by membership per partition (df[df.index.isin(mask_values)]) and reorder afterwards when match_rows="right" is requested.

Environment

uv run repro.py with the PEP 723 metadata in the script (fresh, isolated environment; spatialdata built from main @ ccf1ea0 (2026-08-28); Python 3.13, latest releases of the dependencies at run time: pandas 3.0, anndata 0.13, zarr 3.3, dask 2026.8, numpy 2.5, geopandas 1.1, shapely 2.1). macOS (arm64). Also reproduced in a second environment with pandas 2.3.3 / anndata 0.12.11 / numpy 2.4.4 / zarr 3.2.1.

Automatically generated; discovered by an AI agent (Claude) and not yet reviewed by a human.

Start with _get_masked_element in src/spatialdata/_core/query/relational_query.py and run the provided repro.py with one and two Dask partitions. Check the partition-level filtering used by inner and right joins, then verify that join_spatialelement_table and related query paths produce the same indices without KeyError for multi-partition points.