From 212cb9c2b5112fdee4ce403bf0a51a395d528628 Mon Sep 17 00:00:00 2001 From: Martin Fleischmann Date: Sun, 23 Jan 2022 19:04:39 +0000 Subject: [PATCH] BUG: explode incorrectly expects unsorted index (#2292) --- geopandas/geodataframe.py | 43 ++++------- geopandas/tests/test_geom_methods.py | 103 +++++++++++++++++++++++++++ 2 files changed, 116 insertions(+), 30 deletions(-) diff --git a/geopandas/geodataframe.py b/geopandas/geodataframe.py index b6872fa..788e627 100644 --- a/geopandas/geodataframe.py +++ b/geopandas/geodataframe.py @@ -1684,43 +1684,26 @@ individually so that features may have different properties ) index_parts = True - df_copy = self.copy() + exploded_geom = self.geometry.reset_index(drop=True).explode(index_parts=True) - level_str = f"level_{df_copy.index.nlevels}" - - if level_str in df_copy.columns: # GH1393 - df_copy = df_copy.rename(columns={level_str: f"__{level_str}"}) - - if index_parts: - exploded_geom = df_copy.geometry.explode(index_parts=True) - exploded_index = exploded_geom.index - exploded_geom = exploded_geom.reset_index(level=-1, drop=True) - else: - exploded_geom = df_copy.geometry.explode(index_parts=True).reset_index( - level=-1, drop=True - ) - exploded_index = exploded_geom.index - - df = ( - df_copy.drop(df_copy._geometry_column_name, axis=1) - .join(exploded_geom) - .set_geometry(self._geometry_column_name) - .__finalize__(self) - ) + df = GeoDataFrame( + self.drop(self._geometry_column_name, axis=1).take( + exploded_geom.index.droplevel(-1) + ), + geometry=exploded_geom.values, + ).__finalize__(self) if ignore_index: df.reset_index(inplace=True, drop=True) elif index_parts: # reset to MultiIndex, otherwise df index is only first level of # exploded GeoSeries index. - df.set_index(exploded_index, inplace=True) - df.index.names = list(self.index.names) + [None] - else: - df.set_index(exploded_index, inplace=True) - df.index.names = self.index.names - - if f"__{level_str}" in df.columns: - df = df.rename(columns={f"__{level_str}": level_str}) + df = df.set_index( + exploded_geom.index.droplevel( + list(range(exploded_geom.index.nlevels - 1)) + ), + append=True, + ) return df diff --git a/geopandas/tests/test_geom_methods.py b/geopandas/tests/test_geom_methods.py index 221eca3..941b70f 100644 --- a/geopandas/tests/test_geom_methods.py +++ b/geopandas/tests/test_geom_methods.py @@ -1090,6 +1090,109 @@ class TestGeomMethods: test_df = df.explode(ignore_index=True, index_parts=True) assert_frame_equal(test_df, expected_df) + def test_explode_order(self): + df = GeoDataFrame( + {"vals": [1, 2, 3]}, + geometry=[MultiPoint([(x, x), (x, 0)]) for x in range(3)], + index=[2, 9, 7], + ) + test_df = df.explode(index_parts=True) + + expected_index = MultiIndex.from_arrays( + [[2, 2, 9, 9, 7, 7], [0, 1, 0, 1, 0, 1]], + ) + expected_geometry = GeoSeries( + [ + Point(0, 0), + Point(0, 0), + Point(1, 1), + Point(1, 0), + Point(2, 2), + Point(2, 0), + ], + index=expected_index, + ) + expected_df = GeoDataFrame( + {"vals": [1, 1, 2, 2, 3, 3]}, + geometry=expected_geometry, + index=expected_index, + ) + assert_geodataframe_equal(test_df, expected_df) + + def test_explode_order_no_multi(self): + df = GeoDataFrame( + {"vals": [1, 2, 3]}, + geometry=[Point(0, x) for x in range(3)], + index=[2, 9, 7], + ) + test_df = df.explode(index_parts=True) + + expected_index = MultiIndex.from_arrays( + [[2, 9, 7], [0, 0, 0]], + ) + expected_df = GeoDataFrame( + {"vals": [1, 2, 3]}, + geometry=[Point(0, x) for x in range(3)], + index=expected_index, + ) + assert_geodataframe_equal(test_df, expected_df) + + def test_explode_order_mixed(self): + df = GeoDataFrame( + {"vals": [1, 2, 3]}, + geometry=[MultiPoint([(x, x), (x, 0)]) for x in range(2)] + [Point(0, 10)], + index=[2, 9, 7], + ) + test_df = df.explode(index_parts=True) + + expected_index = MultiIndex.from_arrays( + [[2, 2, 9, 9, 7], [0, 1, 0, 1, 0]], + ) + expected_geometry = GeoSeries( + [ + Point(0, 0), + Point(0, 0), + Point(1, 1), + Point(1, 0), + Point(0, 10), + ], + index=expected_index, + ) + expected_df = GeoDataFrame( + {"vals": [1, 1, 2, 2, 3]}, + geometry=expected_geometry, + index=expected_index, + ) + assert_geodataframe_equal(test_df, expected_df) + + def test_explode_duplicated_index(self): + df = GeoDataFrame( + {"vals": [1, 2, 3]}, + geometry=[MultiPoint([(x, x), (x, 0)]) for x in range(3)], + index=[1, 1, 2], + ) + test_df = df.explode(index_parts=True) + expected_index = MultiIndex.from_arrays( + [[1, 1, 1, 1, 2, 2], [0, 1, 0, 1, 0, 1]], + ) + expected_geometry = GeoSeries( + [ + Point(0, 0), + Point(0, 0), + Point(1, 1), + Point(1, 0), + Point(2, 2), + Point(2, 0), + ], + index=expected_index, + ) + expected_df = GeoDataFrame( + {"vals": [1, 1, 2, 2, 3, 3]}, + geometry=expected_geometry, + index=expected_index, + ) + assert_geodataframe_equal(test_df, expected_df) + # # Test '&', '|', '^', and '-' #