Functions relying on `transform` break with partitioned, dask-backed elements (dask 2025)

[ppm1337]

Description

After upgrading to spatialdata 0.7.0 and dask 2025, plotting of points and other elements with dask dataframe as the underlying data structure appears to be broken.

The issue originates from the calls to transform that render_* functions perform:

transform(sdata["transcripts"], to_coordinate_system="global")

Results in ValueError: The number of items in 'lengths' does not match the number of partitions. if sdata["transcripts"].npartitions is > 1.

MRE

Since I am working with Xenium data, I used the xenium_rep1_io data for my MRE listed in your docs.

import spatialdata as sd
import spatialdata_plot

sdata = sd.read_zarr("temp/xenium/example_data/xenium_v1_human_lung/data.zarr")
feat = sdata["table"].var_names.tolist()[0]
(
 sdata
 .pl.render_images("he_image")
 .pl.render_points(
 "transcripts",
 table_name="table",
 color="feature_name",
 groups=feat,
 palette="red",
 )
 .pl.show()
)

Yields ValueError: The number of items in 'lengths' does not match the number of partitions. 2 != 4

Expand for stacktrace

---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
Cell In[41], line 11
 1 feat = sdata["table"].var_names.tolist()[0]
 2 (
 3 sdata.pl.render_images("he_image")
 4 .pl.render_points(
 5 "transcripts",
 6 table_name="table",
 7 color="feature_name",
 8 groups=feat,
 9 palette="red",
 10 )
---> 11 .pl.show()
 12 )
File ~/projects/ddr_st/ddr_st_xenium/lib/python3.12/site-packages/spatialdata_plot/pl/basic.py:1163, in PlotAccessor.show(self, coordinate_systems, legend_fontsize, legend_fontweight, legend_loc, legend_fontoutline, na_in_legend, colorbar, colorbar_params, wspace, hspace, ncols, frameon, figsize, dpi, fig, title, share_extent, pad_extent, ax, return_ax, save)

1160 ax.set_title(t)

1161 ax.set_aspect("equal")

-> 1163 extent = get_extent(

1164 sdata,

1165 coordinate_system=cs,

1166 has_images=has_images and wants_images,

1167 has_labels=has_labels and wants_labels,

1168 has_points=has_points and wants_points,

1169 has_shapes=has_shapes and wants_shapes,

1170 elements=wanted_elements,

1171 )

1172 cs_x_min, cs_x_max = extent["x"]

1173 cs_y_min, cs_y_max = extent["y"]
File ~/projects/ddr_st/ddr_st_xenium/lib/python3.12/functools.py:912, in singledispatch..wrapper(*args, **kw)

908 if not args:

909 raise TypeError(f'{funcname} requires at least '

910 '1 positional argument')

--> 912 return dispatch(args[0].class)(*args, **kw)
File ~/projects/ddr_st/ddr_st_xenium/lib/python3.12/site-packages/spatialdata/_core/data_extent.py:213, in _(e, coordinate_system, exact, has_images, has_labels, has_points, has_shapes, elements)

211 if coordinate_system in coordinate_systems:

212 if isinstance(element_obj, DaskDataFrame | GeoDataFrame):

--> 213 extent = get_extent(element_obj, coordinate_system=coordinate_system, exact=exact)

214 else:

215 extent = get_extent(element_obj, coordinate_system=coordinate_system)
File ~/projects/ddr_st/ddr_st_xenium/lib/python3.12/functools.py:912, in singledispatch..wrapper(*args, **kw)

908 if not args:

909 raise TypeError(f'{funcname} requires at least '

910 '1 positional argument')

--> 912 return dispatch(args[0].class)(*args, **kw)
File ~/projects/ddr_st/ddr_st_xenium/lib/python3.12/site-packages/spatialdata/_core/data_extent.py:291, in _(e, coordinate_system, exact)

285 extent = _get_extent_of_points(e)

286 return _compute_extent_in_coordinate_system(

287 element=e,

288 coordinate_system=coordinate_system,

289 extent=extent,

290 )

--> 291 transformed = transform(e, to_coordinate_system=coordinate_system)

292 return _get_extent_of_points(transformed)
File ~/projects/ddr_st/ddr_st_xenium/lib/python3.12/functools.py:912, in singledispatch..wrapper(*args, **kw)

908 if not args:

909 raise TypeError(f'{funcname} requires at least '

910 '1 positional argument')

--> 912 return dispatch(args[0].class)(*args, **kw)
File ~/projects/ddr_st/ddr_st_xenium/lib/python3.12/site-packages/spatialdata/_core/operations/transform.py:443, in _(data, transformation, maintain_positioning, to_coordinate_system)

441 arrays = []

442 for ax in axes:

--> 443 arrays.append(data[ax].to_dask_array(lengths=True).reshape(-1, 1))

444 xdata = DataArray(da.concatenate(arrays, axis=1), coords={"points": range(len(data)), "dim": list(axes)})

445 xtransformed = transformation._transform_coordinates(xdata)
File ~/projects/ddr_st/ddr_st_xenium/lib/python3.12/site-packages/dask/dataframe/dask_expr/_collection.py:1422, in FrameBase.to_dask_array(self, lengths, meta, optimize, **optimize_kwargs)

1418 lengths = tuple(self.map_partitions(len).compute())

1420 arr = self.values

-> 1422 chunks = self._validate_chunks(arr, lengths)

1423 arr._chunks = chunks

1425 if meta is not None:
File ~/projects/ddr_st/ddr_st_xenium/lib/python3.12/site-packages/dask/dataframe/dask_expr/_collection.py:2471, in FrameBase._validate_chunks(self, arr, lengths)

2468 lengths = tuple(lengths)

2470 if len(lengths) != self.npartitions:

-> 2471 raise ValueError(

2472 "The number of items in 'lengths' does not match the number of "

2473 f"partitions. {len(lengths)} != {self.npartitions}"

2474 )

2476 if self.ndim == 1:

2477 chunks = normalize_chunks((lengths,))
ValueError: The number of items in 'lengths' does not match the number of partitions. 2 != 4

Workaround

print(sdata["transcripts"].npartitions)
# 4

sdata["transcripts"] = sd.models.PointsModel.parse(
 sdata["transcripts"].repartition(npartitions=1).reset_index(drop=True),
 transformations={"global": sd.transformations.get_transformation(sdata["transcripts"])},
)

print(sdata["transcripts"].npartitions)
# 1

# works without error
(
 sdata
 .pl.render_images("he_image")
 .pl.render_points(
 "transcripts",
 table_name="table",
 color="feature_name",
 groups=feat,
 palette="red",
 )
 .pl.show()
)

Additional context

# works
sd.get_extent(sdata["transcripts"], coordinate_system="global", exact=False)
# raises ValueError
sd.get_extent(sdata["transcripts"], coordinate_system="global", exact=True)
# originating from
sd.transform(sdata["transcripts"], to_coordinate_system="global")

The call to reset_index(drop=True) is necessary in the MRE to avoid ValueError: cannot reindex on an axis with duplicate labels.

This issue cannot be reproduced by partitioning the blobs dataset:

from spatialdata.datasets import blobs

sdata = blobs()

print(sdata["blobs_points"].npartitions)
# 1

sdata["blobs_points"] = sd.models.PointsModel.parse(
 sdata["blobs_points"].repartition(npartitions=4),
 transformations={"global": sd.transformations.get_transformation(sdata["blobs_points"])},
)

print(sdata["blobs_points"].npartitions)
# 4

# works
sd.transform(sdata["blobs_points"], to_coordinate_system="global")

This error occurs with

• spatialdata 0.7.0 + dask 2025.12, dask 2025.9 or dask 2025.2.0

This error does not occur with

• spatialdata 0.6.1 + dask 2024.11.2

Session Info

Click to view session information

-----
dask 2025.12.0
session_info v1.0.1
spatialdata 0.7.0
spatialdata_plot 0.2.14
-----

Click to view modules imported as dependencies

PIL 12.1.0
anndata 0.12.9
annsel 0.1.2
anyio NA
argcomplete NA
arrow 1.4.0
asttokens NA
attr 25.4.0
attrs 25.4.0
babel 2.18.0
certifi 2026.01.04
cffi 2.0.0
charset_normalizer 3.4.4
cloudpickle 3.1.2
comm 0.2.3
cuda 12.9.4
cycler 0.12.1
cython_runtime NA
dask_image NA
datashader 0.18.2
dateutil 2.9.0.post0
debugpy 1.8.20
decorator 5.2.1
defusedxml 0.7.1
donfig 0.8.1.post1
executing 2.2.1
fast_array_utils NA
fastjsonschema NA
fqdn NA
fsspec 2026.2.0
geopandas 1.1.2
google_crc32c NA
h5py 3.15.1
idna 3.11
ipykernel 7.1.0
ipywidgets 8.1.8
isoduration NA
jaraco NA
jedi 0.19.2
jinja2 3.1.6
joblib 1.5.3
json5 0.13.0
jsonpointer 3.0.0
jsonschema 4.26.0
jsonschema_specifications NA
jupyter_events 0.12.0
jupyter_server 2.17.0
jupyterlab_server 2.28.0
kiwisolver 1.4.9
lark 1.3.1
lazy_loader 0.4
legacy_api_wrap NA
llvmlite 0.46.0
locket NA
markupsafe 3.0.3
matplotlib 3.10.8
matplotlib_inline 0.2.1
matplotlib_scalebar 0.9.0
more_itertools 10.8.0
mpl_toolkits NA
msgpack 1.1.2
multipledispatch 0.6.0
multiscale_spatial_image 2.0.3
narwhals 2.16.0
natsort 8.4.0
nbformat 5.10.4
networkx 3.6.1
numba 0.63.1
numcodecs 0.16.5
numexpr 2.14.1
numpy 2.3.5
ome_zarr NA
packaging 26.0
pandas 2.3.3
parso 0.8.5
partd 1.4.2
pathlib_abc NA
pkg_resources NA
platformdirs 4.5.1
pooch v1.9.0
prometheus_client NA
prompt_toolkit 3.0.52
psutil 7.2.2
pure_eval 0.2.3
pyarrow 22.0.0
pycparser 3.00
pyct 0.6.0
pydev_ipython NA
pydevconsole NA
pydevd 3.2.3
pydevd_file_utils NA
pydevd_plugins NA
pydevd_tracing NA
pygments 2.19.2
pyparsing 3.3.2
pyproj 3.7.2
pythonjsonlogger NA
pytz 2025.2
referencing NA
requests 2.32.5
rfc3339_validator 0.1.4
rfc3986_validator 0.1.1
rfc3987_syntax NA
rich NA
rpds NA
scanpy 1.12
scipy 1.16.0
send2trash NA
shapely 2.1.2
six 1.17.0
skimage 0.26.0
sklearn 1.8.0
spatial_image 1.2.3
stack_data 0.6.3
testing NA
threadpoolctl 3.6.0
tlz 1.1.0
toolz 1.1.0
tornado 6.5.4
tqdm 4.67.3
traitlets 5.14.3
typing_extensions NA
upath 0.3.9
uri_template NA
urllib3 2.6.3
vscode NA
wcwidth 0.5.3
webcolors NA
websocket 1.9.0
xarray 2026.1.0
xarray_dataclass 3.0.0
xarray_schema 0.0.3
xrspatial 0.5.2
yaml 6.0.3
zarr 3.1.5
zmq 27.1.0
zoneinfo NA

-----
IPython 9.10.0
jupyter_client 8.8.0
jupyter_core 5.9.1
jupyterlab 4.5.3
notebook 7.5.3
-----
Python 3.12.12 | packaged by conda-forge | (main, Jan 26 2026, 23:51:32) [GCC 14.3.0]
Linux-6.8.0-90-generic-x86_64-with-glibc2.39
-----
Session information updated at 2026-02-06 13:16

[LucaMarconato]

Thanks for the comprehensive bug report! CC @melonora.

I'll make also @timtreis aware, but this should get fixed on the spatialdata side, not spatialdata-plot.

[melonora]

Thanks for reporting, going to investigate and report back

[melonora]

From first glance I can already tell that the issue is due to the new dask expressions and reading data back in from parquet. The new graph optimization that comes with dask expressions basically does not respect the partitions and collapses them thereby causing a difference in length. We had an issue like this before elsewhere and there was at some a fix you could apply from dask so will look into that.

[melonora]

@ppm1337 Following up, the blobs dataset can reproduce the same error. The problem originates from read_parquet in dask. The problem you see occurs when you write a partitioned dask dataframe to parquet, then read it back in and perform certain operations. With the new dask expressions it tries to optimize the dask graph, due to indices overlapping, this collapses the partitions hence the error occurs. There is a fix for this that I will outline below just so you can run code without errors already, but first the MRE:

MRE:

from spatialdata.datasets import blobs
import spatialdata as sd

sdata = blobs()

print(sdata["blobs_points"].npartitions)
# 1

sdata["blobs_points"] = sd.models.PointsModel.parse(
 sdata["blobs_points"].repartition(npartitions=4),
 transformations={"global": sd.transformations.get_transformation(sdata["blobs_points"])},
)

# This is necessary for reproducing the error
sdata.write("temp.zarr")

sdata = sd.read_zarr("temp.zarr")
sd.transform(sdata["blobs_points"], to_coordinate_system="global") # error now occurs here

With fix:

from spatialdata.datasets import blobs
import spatialdata as sd
import dask

sdata = blobs()

print(sdata["blobs_points"].npartitions)
# 1

sdata["blobs_points"] = sd.models.PointsModel.parse(
 sdata["blobs_points"].repartition(npartitions=4),
 transformations={"global": sd.transformations.get_transformation(sdata["blobs_points"])},
)

sdata.write("temp.zarr")

sdata = sd.read_zarr("temp.zarr")
# We prevent dask from optimizing the dask graph with this and thereby prevent the collaps of the partitions.
dask.config.set({"optimization.tune.active": False})
sd.transform(sdata["blobs_points"], to_coordinate_system="global")
dask.config.set({"optimization.tune.active": True})

For now this will allow you to run the code. I will implement this fix directly in the transform code where partitions are being accessed, but it can be that users perform operations themselves as well that are outside the SpatialData library, so I will also document this workaround for now. I will investigate whether in dask I can prevent this collaps without having to adjust this config setting.

[melonora]

For the transform code I also have another potential fix, will open the PR in a bit and cc. Main thing is that when out of a series you create a dask dataframe, instead of setting argument lengths to True, calculating partition lengths automatically, you can explicitly pass the current lengths of the partition so it will not do the automatic determination.

edit:
while the fix initially works, it causes a missing dependency problem in the dask graph. I will use the config setting in the transform code to have the workaround for now implemented.

[ppm1337]

Thanks for investigating this. Can confirm that your MRE and workaround work.

[melonora]

No problem, PR opened