Copy spatial_partitions when copying GeoDataFrame (#56)

dask_geopandas/core.py

@@ -161,6 +161,16 @@ def to_crs(self, crs=None, epsg=None):
         token = f"{self._name}-to_crs"
         return self.map_partitions(M.to_crs, crs=crs, epsg=epsg, token=token)
 
+    def copy(self):
+        """Make a copy of the dataframe
+
+        Creates shallow copies of the computational graph and spatial partitions.
+        Does not affect the underlying data.
+        """
+        self_copy = super().copy()
+        self_copy.spatial_partitions = self.spatial_partitions.copy()
+        return self_copy
+
     @property
     @derived_from(geopandas.base.GeoPandasBase)
     def total_bounds(self):

tests/test_core.py

@@ -354,6 +354,15 @@ def test_to_crs_geoseries(geoseries_points_crs):
     assert all(new.compute() == s.to_crs(new_crs))
 
 
+def test_copy_spatial_partitions(geodf_points):
+    dask_obj = dask_geopandas.from_geopandas(geodf_points, npartitions=2)
+    dask_obj.calculate_spatial_partitions()
+    dask_obj_copy = dask_obj.copy()
+    pd.testing.assert_series_equal(
+        dask_obj.spatial_partitions, dask_obj_copy.spatial_partitions
+    )
+
+
 @pytest.mark.skipif(
     LooseVersion(geopandas.__version__) <= LooseVersion("0.8.1"),
     reason="geopandas 0.8 has bug in apply",