From 5759f0c4ec981bcd99f59f3da3f4fdae45992aa9 Mon Sep 17 00:00:00 2001
From: Brendan Ward <bcward@astutespruce.com>
Date: Thu, 28 Jan 2021 00:44:56 -0800
Subject: [PATCH] ENH: integrate pygeos.get_parts for faster explode() (#1693)

---
 benchmarks/geom_methods.py | 12 ++++++++++--
 geopandas/_compat.py       |  5 +++++
 geopandas/geoseries.py     | 37 +++++++++++++++++++++++++++++++++++--
 3 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/benchmarks/geom_methods.py b/benchmarks/geom_methods.py
index ab6ac61..0da8b4c 100644
--- a/benchmarks/geom_methods.py
+++ b/benchmarks/geom_methods.py
@@ -2,7 +2,7 @@ import random
 
 import numpy as np
 from geopandas import GeoSeries
-from shapely.geometry import Point, LineString, Polygon
+from shapely.geometry import Point, Polygon, MultiPolygon
 
 
 def with_attributes(**attrs):
@@ -25,10 +25,15 @@ class Bench:
         triangles3 = GeoSeries([Polygon([(random.random(), random.random())
                                          for _ in range(3)])
                                 for _ in range(10000)])
+        triangles4 = GeoSeries([
+            MultiPolygon([
+                Polygon([(random.random(), random.random()) for _ in range(3)])
+            ]) for _ in range(10000)])
         triangle = Polygon([(random.random(), random.random())
                             for _ in range(3)])
         self.triangles, self.triangles2 = triangles, triangles2
         self.triangles_big = triangles3
+        self.multi_triangles = triangles4
         self.triangle = triangle
 
     @with_attributes(param_names=['op'],
@@ -98,7 +103,10 @@ class Bench:
     def time_buffer(self, *args):
         self.points.buffer(2)
 
+    def time_explode(self, *args):
+        self.multi_triangles.explode()
+
 
 # TODO
-# project, interpolate, affine_transform, translate, rotate, scale, skew, explode
+# project, interpolate, affine_transform, translate, rotate, scale, skew
 # cx indexer
diff --git a/geopandas/_compat.py b/geopandas/_compat.py
index 413a424..36501f5 100644
--- a/geopandas/_compat.py
+++ b/geopandas/_compat.py
@@ -26,14 +26,19 @@ SHAPELY_GE_17 = str(shapely.__version__) >= LooseVersion("1.7.0")
 SHAPELY_GE_18 = str(shapely.__version__) >= LooseVersion("1.8")
 SHAPELY_GE_20 = str(shapely.__version__) >= LooseVersion("2.0")
 
+
 HAS_PYGEOS = None
 USE_PYGEOS = None
 PYGEOS_SHAPELY_COMPAT = None
 
+PYGEOS_GE_09 = None
+
 try:
     import pygeos  # noqa
 
     HAS_PYGEOS = True
+    PYGEOS_GE_09 = str(pygeos.__version__) >= LooseVersion("0.9")
+
 except ImportError:
     HAS_PYGEOS = False
 
diff --git a/geopandas/geoseries.py b/geopandas/geoseries.py
index 59b6d92..e9381a9 100644
--- a/geopandas/geoseries.py
+++ b/geopandas/geoseries.py
@@ -15,7 +15,7 @@ from geopandas.plotting import plot_series
 from .array import GeometryArray, GeometryDtype, from_shapely
 from .base import is_geometry_type
 from . import _vectorized as vectorized
-from ._compat import ignore_shapely2_warnings
+from . import _compat as compat
 
 
 _SERIES_WARNING_MSG = """\
@@ -197,7 +197,7 @@ class GeoSeries(GeoPandasBase, Series):
             # https://github.com/pandas-dev/pandas/issues/26469
             kwargs.pop("dtype", None)
             # Use Series constructor to handle input data
-            with ignore_shapely2_warnings():
+            with compat.ignore_shapely2_warnings():
                 s = pd.Series(data, index=index, name=name, **kwargs)
             # prevent trying to convert non-geometry objects
             if s.dtype != object:
@@ -669,6 +669,39 @@ class GeoSeries(GeoPandasBase, Series):
         GeoDataFrame.explode
 
         """
+
+        if compat.USE_PYGEOS and compat.PYGEOS_GE_09:
+            import pygeos  # noqa
+
+            geometries, outer_idx = pygeos.get_parts(
+                self.values.data, return_index=True
+            )
+
+            if len(outer_idx):
+                # Generate inner index as a range per value of outer_idx
+                # 1. identify the start of each run of values in outer_idx
+                # 2. count number of values per run
+                # 3. use cumulative sums to create an incremental range
+                #    starting at 0 in each run
+                run_start = np.r_[True, outer_idx[:-1] != outer_idx[1:]]
+                counts = np.diff(np.r_[np.nonzero(run_start)[0], len(outer_idx)])
+                inner_index = (~run_start).cumsum()
+                inner_index -= np.repeat(inner_index[run_start], counts)
+
+            else:
+                inner_index = []
+
+            # extract original index values based on integer index
+            outer_index = self.index.take(outer_idx)
+
+            index = MultiIndex.from_arrays(
+                [outer_index, inner_index], names=self.index.names + [None]
+            )
+
+            return GeoSeries(geometries, index=index, crs=self.crs).__finalize__(self)
+
+        # else PyGEOS is not available or version <= 0.8
+
         index = []
         geometries = []
         for idx, s in self.geometry.iteritems():