compat: Support dask query-planning, drop Python 3.9

.github/workflows/test.yaml

@@ -83,15 +83,15 @@ jobs:
         run: |
           MATRIX=$(jq -nsc '{
               "os": ["ubuntu-latest", "macos-latest", "windows-latest"],
-              "environment": ["test-39", "test-312"]
+              "environment": ["test-310", "test-312"]
           }')
           echo "MATRIX=$MATRIX" >> $GITHUB_ENV
       - name: Set test matrix with 'full' option
         if: env.MATRIX_OPTION == 'full'
         run: |
           MATRIX=$(jq -nsc '{
               "os": ["ubuntu-latest", "macos-latest", "windows-latest"],
-              "environment": ["test-39", "test-310", "test-311", "test-312"]
+              "environment": ["test-310", "test-311", "test-312"]
           }')
           echo "MATRIX=$MATRIX" >> $GITHUB_ENV
       - name: Set test matrix with 'downstream' option

pixi.toml

@@ -8,11 +8,9 @@ install = 'python -m pip install --no-deps --disable-pip-version-check -e .'
 
 [activation.env]
 PYTHONIOENCODING = "utf-8"
-DASK_DATAFRAME__QUERY_PLANNING = "False" # TODO: Support query planning
 USE_PYGEOS = '0'
 
 [environments]
-test-39 = ["py39", "test-core", "test", "test-task", "example", "test-example"]
 test-310 = ["py310", "test-core", "test", "test-task", "example", "test-example"]
 test-311 = ["py311", "test-core", "test", "test-task", "example", "test-example"]
 test-312 = ["py312", "test-core", "test", "test-task", "example", "test-example"]
@@ -22,17 +20,14 @@ lint = ["py311", "lint"]
 
 [dependencies]
 numba = "*"
-dask-core = "<2025.1"  # TODO: Does not work with new DataFrame interface
+dask-core = ">=2025.1"
 fsspec = "*"
 packaging = "*"
 pandas = "*"
 pip = "*"
 pyarrow = ">=10"
 retrying = "*"
 
-[feature.py39.dependencies]
-python = "3.9.*"
-
 [feature.py310.dependencies]
 python = "3.10.*"
 

pyproject.toml

@@ -27,7 +27,7 @@ classifiers = [
     "Topic :: Software Development :: Libraries",
 ]
 dependencies = [
-    'dask <2025.1',
+    'dask >=2025.1',
     'fsspec >=2022.8',
     'numba',
     'packaging',

spatialpandas/dask.py

@@ -13,25 +13,22 @@
 from dask import delayed
 from dask.dataframe.core import get_parallel_type
 from dask.dataframe.extensions import make_array_nonempty
-from dask.dataframe.partitionquantiles import partition_quantiles
+from dask.dataframe.utils import make_meta_obj, meta_nonempty
 from packaging.version import Version
 from retrying import retry
 
-try:
-    from dask.dataframe.backends import meta_nonempty
-    from dask.dataframe.dispatch import make_meta_dispatch
-except ImportError:
-    from dask.dataframe.utils import make_meta as make_meta_dispatch, meta_nonempty
-
 from .geodataframe import GeoDataFrame
 from .geometry.base import GeometryDtype, _BaseCoordinateIndexer
 from .geoseries import GeoSeries
 from .spatialindex import HilbertRtree
 
 
 class DaskGeoSeries(dd.Series):
-    def __init__(self, dsk, name, meta, divisions, *args, **kwargs):
-        super().__init__(dsk, name, meta, divisions)
+
+    _partition_type = GeoSeries
+
+    def __init__(self, expr, *args, **kwargs):
+        super().__init__(expr, *args, **kwargs)
 
         # Init backing properties
         self._partition_bounds = None
@@ -105,7 +102,7 @@ def persist(self, **kwargs):
         )
 
 
-@make_meta_dispatch.register(GeoSeries)
+@make_meta_obj.register(GeoSeries)
 def make_meta_series(s, index=None):
     result = s.head(0)
     if index is not None:
@@ -119,13 +116,20 @@ def meta_nonempty_series(s, index=None):
 
 
 @get_parallel_type.register(GeoSeries)
-def get_parallel_type_dataframe(df):
+def get_parallel_type_series(df):
+    return DaskGeoSeries
+
+
+@dd.get_collection_type.register(GeoSeries)
+def get_collection_type_series(df):
     return DaskGeoSeries
 
 
 class DaskGeoDataFrame(dd.DataFrame):
-    def __init__(self, dsk, name, meta, divisions):
-        super().__init__(dsk, name, meta, divisions)
+    _partition_type = GeoDataFrame
+
+    def __init__(self, expr, *args, **kwargs):
+        super().__init__(expr, *args, **kwargs)
         self._partition_sindex = {}
         self._partition_bounds = {}
 
@@ -312,8 +316,9 @@ def move_retry(p1, p2):
         ddf = self._with_hilbert_distance_column(p)
 
         # Compute output hilbert_distance divisions
-        quantiles = partition_quantiles(
-            ddf.hilbert_distance, npartitions
+        from dask.dataframe.dask_expr import RepartitionQuantiles, new_collection
+        quantiles = new_collection(
+            RepartitionQuantiles(ddf.hilbert_distance, npartitions)
         ).compute().values
 
         # Add _partition column containing output partition number of each row
@@ -328,7 +333,7 @@ def move_retry(p1, p2):
             for out_partition in out_partitions
         ]
         part_output_paths = [
-            os.path.join(path, f"part.{int(out_partition)}.parquet")
+            os.path.join(path, f"part.{out_partition}.parquet")
             for out_partition in out_partitions
         ]
 
@@ -338,7 +343,7 @@ def move_retry(p1, p2):
             rm_retry(path)
 
         for out_partition in out_partitions:
-            part_dir = os.path.join(path, f"part.{int(out_partition)}.parquet" )
+            part_dir = os.path.join(path, f"part.{out_partition}.parquet" )
             mkdirs_retry(part_dir)
             tmp_part_dir = tempdir_format.format(partition=out_partition, uuid=dataset_uuid)
             mkdirs_retry(tmp_part_dir)
@@ -360,7 +365,7 @@ def process_partition(df, i):
             for out_partition, df_part in df.groupby('_partition'):
                 part_path = os.path.join(
                     tempdir_format.format(partition=out_partition, uuid=dataset_uuid),
-                    f'part{int(i)}.parquet',
+                    f'part{i}.parquet',
                 )
                 df_part = (
                     df_part
@@ -584,7 +589,7 @@ def __getitem__(self, key):
         return result
 
 
-@make_meta_dispatch.register(GeoDataFrame)
+@make_meta_obj.register(GeoDataFrame)
 def make_meta_dataframe(df, index=None):
     result = df.head(0)
     if index is not None:
@@ -598,10 +603,14 @@ def meta_nonempty_dataframe(df, index=None):
 
 
 @get_parallel_type.register(GeoDataFrame)
-def get_parallel_type_series(s):
+def get_parallel_type_dataframe(s):
     return DaskGeoDataFrame
 
 
+@dd.get_collection_type.register(GeoDataFrame)
+def get_collection_type_dataframe(df):
+    return DaskGeoDataFrame
+
 class _DaskCoordinateIndexer(_BaseCoordinateIndexer):
     def __init__(self, obj, sindex):
         super().__init__(sindex)

spatialpandas/io/parquet.py

@@ -1,4 +1,5 @@
 import json
+import warnings
 from collections.abc import Iterable
 from functools import reduce
 from glob import has_magic
@@ -180,17 +181,24 @@ def to_parquet_dask(
     spatial_metadata = {'partition_bounds': partition_bounds}
     b_spatial_metadata = json.dumps(spatial_metadata).encode('utf')
 
-    dd_to_parquet(
-        ddf,
-        path,
-        engine="pyarrow",
-        compression=compression,
-        storage_options=storage_options,
-        custom_metadata={b'spatialpandas': b_spatial_metadata},
-        write_metadata_file=True,
-        **engine_kwargs,
-        **kwargs,
-    )
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore",
+            category=UserWarning,
+            message="Dask annotations .* detected. Annotations will be ignored when using query-planning.",
+        )
+        dd_to_parquet(
+            ddf,
+            path,
+            engine="pyarrow",
+            compression=compression,
+            storage_options=storage_options,
+            custom_metadata={b'spatialpandas': b_spatial_metadata},
+            write_metadata_file=True,
+            **engine_kwargs,
+            **kwargs,
+        )
 
 
 def read_parquet_dask(