ENH: use total_bounds of spatial_partitions if available for Hilbert/Morton distance (#161)

Author: jorisvandenbossche

dask_geopandas/core.py

@@ -348,8 +348,10 @@ def hilbert_distance(self, total_bounds=None, level=16):
         total_bounds : 4-element array, optional
             The spatial extent in which the curve is constructed (used to
             rescale the geometry midpoints). By default, the total bounds
-            of the full dask GeoDataFrame will be computed. If known, you
-            can pass the total bounds to avoid this extra computation.
+            of the full dask GeoDataFrame will be computed (from the spatial
+            partitions, if available, otherwise computed from the full
+            dataframe). If known, you can pass the total bounds to avoid this
+            extra computation.
         level : int (1 - 16), default 16
             Determines the precision of the curve (points on the curve will
             have coordinates in the range [0, 2^level - 1]).
@@ -362,7 +364,10 @@ def hilbert_distance(self, total_bounds=None, level=16):
         """
         # Compute total bounds of all partitions rather than each partition
         if total_bounds is None:
-            total_bounds = self.total_bounds
+            if self.spatial_partitions is not None:
+                total_bounds = self.spatial_partitions.total_bounds
+            else:
+                total_bounds = self.total_bounds
 
         # Calculate hilbert distances for each partition
         distances = self.map_partitions(
@@ -396,20 +401,25 @@ def morton_distance(self, total_bounds=None, level=16):
         total_bounds : 4-element array, optional
             The spatial extent in which the curve is constructed (used to
             rescale the geometry midpoints). By default, the total bounds
-            of the full dask GeoDataFrame will be computed. If known, you
-            can pass the total bounds to avoid this extra computation.
+            of the full dask GeoDataFrame will be computed (from the spatial
+            partitions, if available, otherwise computed from the full
+            dataframe). If known, you can pass the total bounds to avoid this
+            extra computation.
         level : int (1 - 16), default 16
             Determines the precision of the Morton curve.
 
         Returns
         -------
         dask.Series
             Series containing distances along the Morton curve
-        """
 
+        """
         # Compute total bounds of all partitions rather than each partition
         if total_bounds is None:
-            total_bounds = self.total_bounds
+            if self.spatial_partitions is not None:
+                total_bounds = self.spatial_partitions.total_bounds
+            else:
+                total_bounds = self.total_bounds
 
         # Calculate Morton distances for each partition
         distances = self.map_partitions(

tests/test_hilbert_distance.py

@@ -108,6 +108,15 @@ def test_specified_total_bounds(geoseries_polygons):
     assert_series_equal(result.compute(), expected.compute())
 
 
+def test_total_bounds_from_partitions(geoseries_polygons):
+    ddf = from_geopandas(geoseries_polygons, npartitions=2)
+    expected = ddf.hilbert_distance().compute()
+
+    ddf.calculate_spatial_partitions()
+    result = ddf.hilbert_distance().compute()
+    assert_series_equal(result, expected)
+
+
 def test_world():
     # world without Fiji
     hilbert_distance_dask(

tests/test_morton_distance.py

@@ -77,6 +77,15 @@ def test_specified_total_bounds(geoseries_polygons):
     assert_series_equal(result.compute(), expected.compute())
 
 
+def test_total_bounds_from_partitions(geoseries_polygons):
+    ddf = from_geopandas(geoseries_polygons, npartitions=2)
+    expected = ddf.morton_distance().compute()
+
+    ddf.calculate_spatial_partitions()
+    result = ddf.morton_distance().compute()
+    assert_series_equal(result, expected)
+
+
 def test_world():
     # world without Fiji
     morton_distance_dask(