diff --git a/geopandas/tests/test_overlay.py b/geopandas/tests/test_overlay.py index a07bf62..ec9b682 100644 --- a/geopandas/tests/test_overlay.py +++ b/geopandas/tests/test_overlay.py @@ -575,6 +575,11 @@ def test_keep_geom_type_geometry_collection(): df1 = read_file(os.path.join(DATA, "geom_type", "df1.geojson")) df2 = read_file(os.path.join(DATA, "geom_type", "df2.geojson")) + with pytest.warns(UserWarning, match="`keep_geom_type=True` in overlay"): + intersection = overlay(df1, df2, keep_geom_type=None) + assert len(intersection) == 1 + assert (intersection.geom_type == "Polygon").all() + intersection = overlay(df1, df2, keep_geom_type=True) assert len(intersection) == 1 assert (intersection.geom_type == "Polygon").all() @@ -584,6 +589,44 @@ def test_keep_geom_type_geometry_collection(): assert (intersection.geom_type == "GeometryCollection").all() +def test_keep_geom_type_geometry_collection2(): + polys1 = [ + box(0, 0, 1, 1), + box(1, 1, 3, 3).union(box(1, 3, 5, 5)), + ] + + polys2 = [ + box(0, 0, 1, 1), + box(3, 1, 4, 2).union(box(4, 1, 5, 4)), + ] + df1 = GeoDataFrame({"left": [0, 1], "geometry": polys1}) + df2 = GeoDataFrame({"right": [0, 1], "geometry": polys2}) + + result1 = overlay(df1, df2, keep_geom_type=True) + expected1 = GeoDataFrame( + { + "left": [0, 1], + "right": [0, 1], + "geometry": [box(0, 0, 1, 1), box(4, 3, 5, 4)], + } + ) + assert_geodataframe_equal(result1, expected1) + + result1 = overlay(df1, df2, keep_geom_type=False) + expected1 = GeoDataFrame( + { + "left": [0, 1, 1], + "right": [0, 0, 1], + "geometry": [ + box(0, 0, 1, 1), + Point(1, 1), + GeometryCollection([box(4, 3, 5, 4), LineString([(3, 1), (3, 2)])]), + ], + } + ) + assert_geodataframe_equal(result1, expected1) + + @pytest.mark.parametrize("make_valid", [True, False]) def test_overlap_make_valid(make_valid): bowtie = Polygon([(1, 1), (9, 9), (9, 1), (1, 9), (1, 1)]) diff --git a/geopandas/tools/overlay.py b/geopandas/tools/overlay.py index ef47fa1..ca516f8 100644 --- a/geopandas/tools/overlay.py +++ b/geopandas/tools/overlay.py @@ -185,8 +185,8 @@ def overlay(df1, df2, how="intersection", keep_geom_type=None, make_valid=True): 1 2.0 1.0 POLYGON ((2.00000 2.00000, 2.00000 3.00000, 3.... 2 2.0 2.0 POLYGON ((4.00000 4.00000, 4.00000 3.00000, 3.... 3 1.0 NaN POLYGON ((2.00000 0.00000, 0.00000 0.00000, 0.... - 4 2.0 NaN MULTIPOLYGON (((4.00000 3.00000, 4.00000 2.000... - 5 NaN 1.0 MULTIPOLYGON (((3.00000 2.00000, 3.00000 1.000... + 4 2.0 NaN MULTIPOLYGON (((3.00000 3.00000, 4.00000 3.000... + 5 NaN 1.0 MULTIPOLYGON (((2.00000 2.00000, 3.00000 2.000... 6 NaN 2.0 POLYGON ((3.00000 5.00000, 5.00000 5.00000, 5.... >>> geopandas.overlay(df1, df2, how='intersection') @@ -198,8 +198,8 @@ def overlay(df1, df2, how="intersection", keep_geom_type=None, make_valid=True): >>> geopandas.overlay(df1, df2, how='symmetric_difference') df1_data df2_data geometry 0 1.0 NaN POLYGON ((2.00000 0.00000, 0.00000 0.00000, 0.... - 1 2.0 NaN MULTIPOLYGON (((4.00000 3.00000, 4.00000 2.000... - 2 NaN 1.0 MULTIPOLYGON (((3.00000 2.00000, 3.00000 1.000... + 1 2.0 NaN MULTIPOLYGON (((3.00000 3.00000, 4.00000 3.000... + 2 NaN 1.0 MULTIPOLYGON (((2.00000 2.00000, 3.00000 2.000... 3 NaN 2.0 POLYGON ((3.00000 5.00000, 5.00000 5.00000, 5.... >>> geopandas.overlay(df1, df2, how='difference') @@ -213,7 +213,7 @@ def overlay(df1, df2, how="intersection", keep_geom_type=None, make_valid=True): 1 2.0 1.0 POLYGON ((2.00000 2.00000, 2.00000 3.00000, 3.... 2 2.0 2.0 POLYGON ((4.00000 4.00000, 4.00000 3.00000, 3.... 3 1.0 NaN POLYGON ((2.00000 0.00000, 0.00000 0.00000, 0.... - 4 2.0 NaN MULTIPOLYGON (((4.00000 3.00000, 4.00000 2.000... + 4 2.0 NaN MULTIPOLYGON (((3.00000 3.00000, 4.00000 3.000... See also -------- @@ -317,32 +317,64 @@ def overlay(df1, df2, how="intersection", keep_geom_type=None, make_valid=True): result = dfunion[dfunion["__idx1"].notnull()].copy() if keep_geom_type: - key_order = result.keys() - exploded = result.reset_index(drop=True).explode() - exploded = exploded.reset_index(level=0) - - orig_num_geoms = result.shape[0] geom_type = df1.geom_type.iloc[0] + + # First we filter the geometry types inside GeometryCollections objects + # (e.g. GeometryCollection([polygon, point]) -> polygon) + # we do this separately on only the relevant rows, as this is an expensive + # operation (an expensive no-op for geometry types other than collections) + is_collection = result.geom_type == "GeometryCollection" + if is_collection.any(): + geom_col = result._geometry_column_name + collections = result[[geom_col]][is_collection] + + exploded = collections.reset_index(drop=True).explode() + exploded = exploded.reset_index(level=0) + + orig_num_geoms_exploded = exploded.shape[0] + if geom_type in polys: + exploded = exploded.loc[exploded.geom_type.isin(polys)] + elif geom_type in lines: + exploded = exploded.loc[exploded.geom_type.isin(lines)] + elif geom_type in points: + exploded = exploded.loc[exploded.geom_type.isin(points)] + else: + raise TypeError( + "`keep_geom_type` does not support {}.".format(geom_type) + ) + num_dropped_collection = orig_num_geoms_exploded - exploded.shape[0] + + # level_0 created with above reset_index operation + # and represents the original geometry collections + # TODO avoiding dissolve to call unary_union in this case could further + # improve performance (we only need to collect geometries in their + # respective Multi version) + dissolved = exploded.dissolve(by="level_0") + result.loc[is_collection, geom_col] = dissolved[geom_col].values + else: + num_dropped_collection = 0 + + # Now we filter all geometries (in theory we don't need to do this + # again for the rows handled above for GeometryCollections, but filtering + # them out is probably more expensive as simply including them when this + # is typically about only a few rows) + orig_num_geoms = result.shape[0] if geom_type in polys: - exploded = exploded.loc[exploded.geom_type.isin(polys)] + result = result.loc[result.geom_type.isin(polys)] elif geom_type in lines: - exploded = exploded.loc[exploded.geom_type.isin(lines)] + result = result.loc[result.geom_type.isin(lines)] elif geom_type in points: - exploded = exploded.loc[exploded.geom_type.isin(points)] + result = result.loc[result.geom_type.isin(points)] else: raise TypeError("`keep_geom_type` does not support {}.".format(geom_type)) + num_dropped = orig_num_geoms - result.shape[0] - # level_0 created with above reset_index operation - # and represents the original geometry collections - result = exploded.dissolve(by="level_0")[key_order] - - if (result.shape[0] != orig_num_geoms) and keep_geom_type_warning: - num_dropped = orig_num_geoms - result.shape[0] + if (num_dropped > 0 or num_dropped_collection > 0) and keep_geom_type_warning: warnings.warn( "`keep_geom_type=True` in overlay resulted in {} dropped " "geometries of different geometry types than df1 has. " "Set `keep_geom_type=False` to retain all " - "geometries".format(num_dropped), + "geometries".format(num_dropped + num_dropped_collection), UserWarning, stacklevel=2, )