Points operations load the whole element into memory: polygon_query needs 17 GB / 15 s to return 25k of 42M transcripts

[!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

Measured on xenium_rep1_io (42,638,083 transcripts in 8 parquet partitions), each probe in a fresh process:

len(transcripts) (lazy)                      0.1 s   0.68 GB
get_extent(transcripts, exact=True)          2.1 s   7.58 GB
get_extent(transcripts, exact=False)         0.6 s   3.47 GB   (identical result, Scale-only transformation)
bounding_box_query(transcripts, 500x500)     0.7 s   8.79 GB   -> 24,699 rows
polygon_query(transcripts, same box)        14.6 s  17.31 GB   -> 24,699 rows
transform(transcripts, to 'global')          1.8 s   7.52 GB

bounding_box_query computes the full frame; polygon_query additionally converts all rows to shapely Point objects before sjoin although the polygon's bounding box could pre-filter; transform computes partition lengths, all non-axis columns and each axis; get_extent(exact=True) transforms first even when the transformation has no rotation/shear (the documentation states the approximate path is exact in that case). The memory limitation of bounding_box_query is documented in its Notes, but not the polygon_query gap.

Severity (agent's assessment): high for real data — a 16 GB laptop cannot run polygon_query on a standard Xenium dataset; bounding_box_query, transform and get_extent(exact=True) need 7.5–9 GB to return 0.06 % of the rows

Where: src/spatialdata/_core/query/spatial_query.py (DaskDataFrame overloads of bounding_box_query → points.compute(), and polygon_query → points_dask_dataframe_to_geopandas(points) + sjoin), operations/transform.py (DaskDataFrame overload), _core/data_extent.py

Expected behaviour

Memory roughly proportional to the query result (or to one partition), not to the whole element.

Reproduction

Requires the xenium_rep1_io dataset from https://github.com/giovp/spatialdata-sandbox; pass its data.zarr path as the first argument.

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",
# ]
# ///
"""Points queries / transform / get_extent load the whole element into memory.

Needs a real dataset: the Xenium `xenium_rep1_io` store from https://github.com/giovp/spatialdata-sandbox
(42.6 M transcripts). Usage: uv run repro.py /path/to/xenium_rep1_io/data.zarr
Each probe runs in a fresh subprocess and reports wall time and peak RSS.
"""
import json
import subprocess
import sys
import textwrap

path = sys.argv[1] if len(sys.argv) > 1 else "/Users/macbook/embl/projects/basel/spatialdata-sandbox/xenium_rep1_io/data.zarr"
PROBES = {
    "len(transcripts) (lazy)": "out = dict(n=len(sdata['transcripts']))",
    "get_extent(transcripts, exact=True)": "out = {k: [float(x) for x in v] for k, v in get_extent(sdata['transcripts']).items()}",
    "get_extent(transcripts, exact=False)": "out = {k: [float(x) for x in v] for k, v in get_extent(sdata['transcripts'], exact=False).items()}",
    "bounding_box_query(transcripts, 500x500 box)": "r = bounding_box_query(sdata['transcripts'], axes=('x', 'y'), min_coordinate=[1000, 1000], max_coordinate=[1500, 1500], target_coordinate_system='global'); out = dict(rows=len(r.compute()))",
    "polygon_query(transcripts, same 500x500 box)": "from shapely.geometry import box; r = polygon_query(sdata['transcripts'], polygon=box(1000, 1000, 1500, 1500), target_coordinate_system='global'); out = dict(rows=len(r.compute()))",
    "transform(transcripts, to 'global')": "r = transform(sdata['transcripts'], to_coordinate_system='global'); out = dict(npartitions=r.npartitions)",
}
TEMPLATE = textwrap.dedent("""
    import json, resource, time, warnings
    warnings.simplefilter("ignore")
    from spatialdata import read_zarr, get_extent, bounding_box_query, polygon_query, transform
    sdata = read_zarr({path!r})
    t0 = time.time()
    {body}
    print(json.dumps(dict(seconds=round(time.time() - t0, 1), peak_rss_gb=round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e9, 2), out=out)))
""")
results = {}
for name, body in PROBES.items():
    res = subprocess.run([sys.executable, "-c", TEMPLATE.format(path=path, body=body)], capture_output=True, text=True)
    lines = [l for l in res.stdout.splitlines() if l.startswith("{")]
    results[name] = json.loads(lines[-1]) if lines else {"error": res.stderr.strip().splitlines()[-1][:200] if res.stderr.strip() else "no output"}
    print(f"{name:48s} -> {results[name]}", flush=True)
bbox = results.get("bounding_box_query(transcripts, 500x500 box)", {})
poly = results.get("polygon_query(transcripts, same 500x500 box)", {})
bug = poly.get("peak_rss_gb", 0) > 8 or bbox.get("peak_rss_gb", 0) > 4
print("VERDICT:", "BUG REPRODUCED (several GB of RAM to return <0.1% of the rows)" if bug else "NOT REPRODUCED / dataset not available")

Observed output

len(transcripts) (lazy)                          -> {'seconds': 0.5, 'peak_rss_gb': 0.8, 'out': {'n': 42638083}}
get_extent(transcripts, exact=True)              -> {'seconds': 2.5, 'peak_rss_gb': 8.37, 'out': {'x': [-8.81606382482192, 35401.158088235294], 'y': [20.779946271110983, 25757.693014705885], 'z': [2.4367334842681885, 50.03297424316406]}}
get_extent(transcripts, exact=False)             -> {'seconds': 0.8, 'peak_rss_gb': 3.89, 'out': {'x': [-8.81606382482192, 35401.158088235294], 'y': [20.779946271110983, 25757.693014705885], 'z': [2.4367334842681885, 50.03297424316406]}}
bounding_box_query(transcripts, 500x500 box)     -> {'seconds': 0.9, 'peak_rss_gb': 8.84, 'out': {'rows': 24699}}
polygon_query(transcripts, same 500x500 box)     -> {'seconds': 15.6, 'peak_rss_gb': 14.75, 'out': {'rows': 24699}}
transform(transcripts, to 'global')              -> {'seconds': 2.0, 'peak_rss_gb': 8.56, 'out': {'npartitions': 8}}
VERDICT: BUG REPRODUCED (several GB of RAM to return <0.1% of the rows)

Possible fix direction (unverified)

1. polygon_query(points): apply the existing bounding-box mask in intrinsic coordinates first, then the exact test on the candidates. 2. bounding_box_query(points): filter with map_partitions instead of compute(). 3. transform(points): apply the affine per partition with map_partitions. 4. get_extent(points): use the exact=False path automatically when rotation and shear are identity (_decompose_transformation_full already exists). Related: #893, #210.

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).

Possibly related issues

#893, #210

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

Start by running repro.py against the xenium_rep1_io dataset, then inspect src/spatialdata/_core/query/spatial_query.py, operations/transform.py, and _core/data_extent.py. Compare the DaskDataFrame paths for bounding_box_query, polygon_query, transform, and get_extent; done means the reproduced queries return the same rows and extents without loading the whole element into memory.