|
1 | 1 | import numpy as np
|
2 | 2 |
|
3 | 3 | from dask.base import tokenize
|
| 4 | +from dask.dataframe import from_graph |
4 | 5 | from dask.highlevelgraph import HighLevelGraph
|
5 | 6 | from dask.utils import derived_from
|
6 | 7 |
|
7 | 8 | import geopandas
|
8 | 9 |
|
9 |
| -from . import backends |
10 |
| - |
11 | 10 |
|
12 | 11 | @derived_from(geopandas.tools)
|
13 | 12 | def clip(gdf, mask, keep_geom_type=False):
|
14 |
| - |
15 |
| - if backends.QUERY_PLANNING_ON: |
16 |
| - from .expr import GeoDataFrame, GeoSeries |
17 |
| - else: |
18 |
| - from .core import GeoDataFrame, GeoSeries |
| 13 | + from dask_geopandas import GeoDataFrame, GeoSeries |
19 | 14 |
|
20 | 15 | if isinstance(mask, (GeoDataFrame, GeoSeries)):
|
21 | 16 | raise NotImplementedError("Mask cannot be a Dask GeoDataFrame or GeoSeries.")
|
@@ -45,17 +40,8 @@ def clip(gdf, mask, keep_geom_type=False):
|
45 | 40 | }
|
46 | 41 | divisions = [None] * (len(dsk) + 1)
|
47 | 42 | graph = HighLevelGraph.from_collections(name, dsk, dependencies=[gdf])
|
48 |
| - if backends.QUERY_PLANNING_ON: |
49 |
| - from dask_expr import from_graph |
50 |
| - |
51 |
| - result = from_graph(graph, gdf._meta, tuple(divisions), dsk.keys(), "clip") |
52 |
| - else: |
53 |
| - from .core import GeoDataFrame, GeoSeries |
54 | 43 |
|
55 |
| - if isinstance(gdf, GeoDataFrame): |
56 |
| - result = GeoDataFrame(graph, name, gdf._meta, tuple(divisions)) |
57 |
| - elif isinstance(gdf, GeoSeries): |
58 |
| - result = GeoSeries(graph, name, gdf._meta, tuple(divisions)) |
| 44 | + result = from_graph(graph, gdf._meta, tuple(divisions), dsk.keys(), "clip") |
59 | 45 |
|
60 | 46 | result.spatial_partitions = new_spatial_partitions
|
61 | 47 | return result
|
0 commit comments