Removed dask pinning (#570)

* removed dask pinning

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fixed top level error tests

* wip

* testing against 3.12

* fix _get_backing_files; wip fix .attrs

* wip

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* wip

* wip

* all tests should be fixed (dask-expr installed, but backend disabled)

* attempt fix python3.12 tests

* improved search for dask backing files

* fix test backing files

* add comment

* fix ruff

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* wip pickled dask graphs for tests

* removing pickled tests; requiring minimum version of dask instead

* forgot to re-disable dask-expr

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

.github/workflows/build.yaml

@@ -11,10 +11,10 @@ jobs:
         runs-on: ubuntu-latest
         steps:
             - uses: actions/checkout@v2
-            - name: Set up Python 3.10
+            - name: Set up Python 3.12
               uses: actions/setup-python@v4
               with:
-                  python-version: "3.10"
+                  python-version: "3.12"
                   cache: pip
             - name: Install build dependencies
               run: python -m pip install --upgrade pip wheel twine build

.github/workflows/release.yaml

@@ -14,10 +14,10 @@ jobs:
             - name: Checkout code
               uses: actions/checkout@v3
 
-            - name: Set up Python 3.10
+            - name: Set up Python 3.12
               uses: actions/setup-python@v4
               with:
-                  python-version: "3.10"
+                  python-version: "3.12"
 
             - name: Install hatch
               run: pip install hatch

.github/workflows/test.yaml

@@ -18,15 +18,15 @@ jobs:
         strategy:
             fail-fast: false
             matrix:
-                python: ["3.9", "3.10"]
+                python: ["3.9", "3.12"]
                 os: [ubuntu-latest]
                 include:
                     - os: macos-latest
                       python: "3.9"
                     - os: macos-latest
-                      python: "3.10"
+                      python: "3.12"
                       pip-flags: "--pre"
-                      name: "Python 3.10 (pre-release)"
+                      name: "Python 3.12 (pre-release)"
 
         env:
             OS: ${{ matrix.os }}

pyproject.toml

@@ -25,7 +25,7 @@ dependencies = [
     "anndata>=0.9.1",
     "click",
     "dask-image",
-    "dask<=2024.2.1",
+    "dask>=2024.2.1",
     "fsspec<=2023.6",
     "geopandas>=0.14",
     "multiscale_spatial_image>=1.0.0",
@@ -36,6 +36,7 @@ dependencies = [
     "pooch",
     "pyarrow",
     "rich",
+    "setuptools",
     "shapely>=2.0.1",
     "spatial_image>=1.1.0",
     "scikit-image",

src/spatialdata/__init__.py

@@ -1,5 +1,20 @@
 from __future__ import annotations
 
+import dask
+
+dask.config.set({"dataframe.query-planning": False})
+from dask.dataframe import DASK_EXPR_ENABLED
+
+# Setting `dataframe.query-planning` to False is effective only if run before `dask.dataframe` is initialized. In
+# the case in which the user had initilized `dask.dataframe` before, we would have DASK_EXPER_ENABLED set to `True`.
+# Here we check that this does not happen.
+if DASK_EXPR_ENABLED:
+    raise RuntimeError(
+        "Unsupported backend: dask-expr has been detected as the backend of dask.dataframe. Please "
+        "use:\nimport dask\ndask.config.set({'dataframe.query-planning': False})\nbefore importing "
+        "dask.dataframe to disable dask-expr. The support is being worked on, for more information please see"
+        "https://github.com/scverse/spatialdata/pull/570"
+    )
 from importlib.metadata import version
 
 __version__ = version("spatialdata")

src/spatialdata/_core/_deepcopy.py

@@ -6,7 +6,7 @@
 from anndata import AnnData
 from dask.array.core import Array as DaskArray
 from dask.array.core import from_array
-from dask.dataframe.core import DataFrame as DaskDataFrame
+from dask.dataframe import DataFrame as DaskDataFrame
 from datatree import DataTree
 from geopandas import GeoDataFrame
 from xarray import DataArray
@@ -64,9 +64,9 @@ def _(element: DataArray) -> DataArray:
 
 @deepcopy.register(DataTree)
 def _(element: DataTree) -> DataTree:
-    # the complexity here is due to the fact that the parsers don't accept MultiscaleSpatialImage types and that we need
-    # to convert the DataTree to a MultiscaleSpatialImage. This will be simplified once we support
-    # multiscale_spatial_image 1.0.0
+    # TODO: now that multiscale_spatial_image 1.0.0 is supported, this code can probably be simplified. Check
+    # https://github.com/scverse/spatialdata/pull/587/files#diff-c74ebf49cb8cbddcfaec213defae041010f2043cfddbded24175025b6764ef79
+    # to understand the original motivation.
     model = get_model(element)
     for key in element:
         ds = element[key].ds

src/spatialdata/_core/_elements.py

@@ -8,7 +8,7 @@
 from warnings import warn
 
 from anndata import AnnData
-from dask.dataframe.core import DataFrame as DaskDataFrame
+from dask.dataframe import DataFrame as DaskDataFrame
 from geopandas import GeoDataFrame
 
 from spatialdata._types import Raster_T

src/spatialdata/_core/centroids.py

@@ -6,7 +6,7 @@
 import dask.array as da
 import pandas as pd
 import xarray as xr
-from dask.dataframe.core import DataFrame as DaskDataFrame
+from dask.dataframe import DataFrame as DaskDataFrame
 from datatree import DataTree
 from geopandas import GeoDataFrame
 from shapely import MultiPolygon, Point, Polygon

src/spatialdata/_core/data_extent.py

@@ -7,7 +7,7 @@
 
 import numpy as np
 import pandas as pd
-from dask.dataframe.core import DataFrame as DaskDataFrame
+from dask.dataframe import DataFrame as DaskDataFrame
 from datatree import DataTree
 from geopandas import GeoDataFrame
 from shapely import MultiPolygon, Point, Polygon

src/spatialdata/_core/operations/_utils.py

@@ -62,7 +62,7 @@ def transform_to_data_extent(
     Notes
     -----
         - The data extent is the smallest rectangle that contains all the images and geometries.
-        - MultiscaleSpatialImage objects will be converted to SpatialImage objects.
+        - DataTree objects (multiscale images) will be converted to DataArray (single-scale images) objects.
         - This helper function will be deprecated when https://github.com/scverse/spatialdata/issues/308 is closed,
           as this function will be easily recovered by `transform_to_coordinate_system()`
     """

src/spatialdata/_core/operations/aggregate.py

@@ -9,7 +9,7 @@
 import geopandas as gpd
 import numpy as np
 import pandas as pd
-from dask.dataframe.core import DataFrame as DaskDataFrame
+from dask.dataframe import DataFrame as DaskDataFrame
 from datatree import DataTree
 from geopandas import GeoDataFrame
 from scipy import sparse

src/spatialdata/_core/operations/rasterize.py

@@ -3,8 +3,8 @@
 import dask_image.ndinterp
 import datashader as ds
 import numpy as np
-from dask.array.core import Array as DaskArray
-from dask.dataframe.core import DataFrame as DaskDataFrame
+from dask.array import Array as DaskArray
+from dask.dataframe import DataFrame as DaskDataFrame
 from datatree import DataTree
 from geopandas import GeoDataFrame
 from shapely import Point

src/spatialdata/_core/operations/rasterize_bins.py

@@ -10,7 +10,7 @@
 from numpy.random import default_rng
 from scipy.sparse import csc_matrix
 from skimage.transform import estimate_transform
-from spatial_image import SpatialImage
+from xarray import DataArray
 
 from spatialdata._core.query.relational_query import get_values
 from spatialdata.models import Image2DModel, get_table_keys
@@ -30,7 +30,7 @@ def rasterize_bins(
     col_key: str,
     row_key: str,
     value_key: str | list[str] | None = None,
-) -> SpatialImage:
+) -> DataArray:
     """
     Rasterizes grid-like binned shapes/points annotated by a table (e.g. Visium HD data).
 

src/spatialdata/_core/operations/transform.py

@@ -9,11 +9,10 @@
 import dask_image.ndinterp
 import numpy as np
 from dask.array.core import Array as DaskArray
-from dask.dataframe.core import DataFrame as DaskDataFrame
+from dask.dataframe import DataFrame as DaskDataFrame
 from datatree import DataTree
 from geopandas import GeoDataFrame
 from shapely import Point
-from spatial_image import SpatialImage
 from xarray import DataArray
 
 from spatialdata._core.spatialdata import SpatialData
@@ -177,6 +176,9 @@ def _set_transformation_for_transformed_elements(
 
     d = get_transformation(element, get_all=True)
     assert isinstance(d, dict)
+    if DEFAULT_COORDINATE_SYSTEM not in d:
+        raise RuntimeError(f"Coordinate system {DEFAULT_COORDINATE_SYSTEM} not found in element")
+    assert isinstance(d, dict)
     assert len(d) == 1
     assert isinstance(d[DEFAULT_COORDINATE_SYSTEM], Identity)
     remove_transformation(element, remove_all=True)
@@ -389,9 +391,7 @@ def _(
             raster_translation = raster_translation_single_scale
         # we set a dummy empty dict for the transformation that will be replaced with the correct transformation for
         # each scale later in this function, when calling set_transformation()
-        transformed_dict[k] = SpatialImage(
-            transformed_dask, dims=xdata.dims, name=xdata.name, attrs={TRANSFORM_KEY: {}}
-        )
+        transformed_dict[k] = DataArray(transformed_dask, dims=xdata.dims, name=xdata.name, attrs={TRANSFORM_KEY: {}})
 
     # mypy thinks that schema could be ShapesModel, PointsModel, ...
     transformed_data = DataTree.from_dict(transformed_dict)
@@ -435,6 +435,9 @@ def _(
     transformed = data.drop(columns=list(axes)).copy()
     # dummy transformation that will be replaced by _adjust_transformation()
     transformed.attrs[TRANSFORM_KEY] = {DEFAULT_COORDINATE_SYSTEM: Identity()}
+    # TODO: the following line, used in place of the line before, leads to an incorrect aggregation result. Look into
+    #  this! Reported here: ...
+    # transformed.attrs = {TRANSFORM_KEY: {DEFAULT_COORDINATE_SYSTEM: Identity()}}
     assert isinstance(transformed, DaskDataFrame)
     for ax in axes:
         indices = xtransformed["dim"] == ax

src/spatialdata/_core/query/relational_query.py

@@ -12,7 +12,7 @@
 import numpy as np
 import pandas as pd
 from anndata import AnnData
-from dask.dataframe.core import DataFrame as DaskDataFrame
+from dask.dataframe import DataFrame as DaskDataFrame
 from datatree import DataTree
 from geopandas import GeoDataFrame
 from xarray import DataArray

src/spatialdata/_core/query/spatial_query.py

@@ -9,7 +9,7 @@
 import dask.array as da
 import dask.dataframe as dd
 import numpy as np
-from dask.dataframe.core import DataFrame as DaskDataFrame
+from dask.dataframe import DataFrame as DaskDataFrame
 from datatree import DataTree
 from geopandas import GeoDataFrame
 from shapely.geometry import MultiPolygon, Polygon

src/spatialdata/_core/spatialdata.py

@@ -11,16 +11,14 @@
 import pandas as pd
 import zarr
 from anndata import AnnData
+from dask.dataframe import DataFrame as DaskDataFrame
 from dask.dataframe import read_parquet
-from dask.dataframe.core import DataFrame as DaskDataFrame
 from dask.delayed import Delayed
 from datatree import DataTree
 from geopandas import GeoDataFrame
-from multiscale_spatial_image.multiscale_spatial_image import MultiscaleSpatialImage
 from ome_zarr.io import parse_url
 from ome_zarr.types import JSONDict
 from shapely import MultiPolygon, Polygon
-from spatial_image import SpatialImage
 from xarray import DataArray
 
 from spatialdata._core._elements import Images, Labels, Points, Shapes, Tables
@@ -97,9 +95,7 @@ class SpatialData:
     -----
     The SpatialElements are stored with standard types:
 
-        - images and labels are stored as :class:`spatial_image.SpatialImage` or
-            :class:`multiscale_spatial_image.MultiscaleSpatialImage` objects, which are respectively equivalent to
-            :class:`xarray.DataArray` and to a :class:`datatree.DataTree` of :class:`xarray.DataArray` objects.
+        - images and labels are stored as :class:`xarray.DataArray` or :class:`datatree.DataTree` objects.
         - points are stored as :class:`dask.dataframe.DataFrame` objects.
         - shapes are stored as :class:`geopandas.GeoDataFrame`.
         - the table are stored as :class:`anndata.AnnData` objects,  with the spatial coordinates stored in the obsm
@@ -856,8 +852,8 @@ def transform_element_to_coordinate_system(
         else:
             # When maintaining positioning is true, and if the element has a transformation to target_coordinate_system
             # (this may not be the case because it could be that the element is not directly mapped to that coordinate
-            # system), then the transformation to the target coordinate system is not needed # because the data is now
-            # already transformed; here we remove such transformation.
+            # system), then the transformation to the target coordinate system is not needed
+            # because the data is now already transformed; here we remove such transformation.
             d = get_transformation(transformed, get_all=True)
             assert isinstance(d, dict)
             if target_coordinate_system in d:
@@ -1595,7 +1591,7 @@ def add_image(
     def add_labels(
         self,
         name: str,
-        labels: SpatialImage | MultiscaleSpatialImage,
+        labels: DataArray | DataTree,
         storage_options: JSONDict | list[JSONDict] | None = None,
         overwrite: bool = False,
     ) -> None:
@@ -1743,8 +1739,8 @@ def h(s: str) -> str:
                     descr += f"{h(attr + 'level1.1')}{k!r}: {descr_class} " f"shape: {v.shape} (2D shapes)"
                 elif attr == "points":
                     length: int | None = None
-                    if len(v.dask.layers) == 1:
-                        name, layer = v.dask.layers.items().__iter__().__next__()
+                    if len(v.dask) == 1:
+                        name, layer = v.dask.items().__iter__().__next__()
                         if "read-parquet" in name:
                             t = layer.creation_info["args"]
                             assert isinstance(t, tuple)