Merge branch 'exploding-frames' of https://github.com/MarcoGorelli/pandas into exploding-frames

Author: MarcoGorelli

doc/source/whatsnew/v1.0.0.rst

@@ -143,7 +143,7 @@ Indexing
 ^^^^^^^^
 
 - Bug in assignment using a reverse slicer (:issue:`26939`)
--
+- Bug in :meth:`DataFrame.explode` would duplicate frame in the presence of duplicates in the index (:issue:`28010`)
 
 Missing
 ^^^^^^^

pandas/core/frame.py

@@ -6259,12 +6259,14 @@ def explode(self, column: Union[str, Tuple]) -> "DataFrame":
         if not self.columns.is_unique:
             raise ValueError("columns must be unique")
 
-        result = self[column].explode()
-        return (
-            self.drop([column], axis=1)
-            .join(result)
-            .reindex(columns=self.columns, copy=False)
-        )
+        result = self.copy()
+        exploded_col = result.pop(column).reset_index(drop=True).explode()
+        result = result.reset_index().join(exploded_col)
+        result.set_index(result.columns[: self.index.nlevels].tolist(), inplace=True)
+        result.index.names = self.index.names
+        result = result.reindex(columns=self.columns, copy=False)
+
+        return result
 
     def unstack(self, level=-1, fill_value=None):
         """

pandas/tests/frame/test_explode.py

@@ -118,3 +118,84 @@ def test_usecase():
         index=[0, 0, 1, 1],
     )
     tm.assert_frame_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    "df, expected",
+    [
+        (
+            pd.DataFrame({"col": [[1, 2], [3, 4]]}, index=[0, 0]),
+            pd.DataFrame({"col": [1, 2, 3, 4]}, index=[0, 0, 0, 0], dtype=object),
+        ),
+        (
+            pd.DataFrame(
+                {"col": [[1, 2], [3, 4]], "other_col": ["a", "b"]}, index=[0, 0]
+            ),
+            pd.DataFrame(
+                {"col": [1, 2, 3, 4], "other_col": ["a", "a", "b", "b"]},
+                index=[0, 0, 0, 0],
+                dtype=object,
+            ),
+        ),
+        (
+            pd.DataFrame(
+                {"col": [[1, 2], [3, 4]], "other_col": ["a", "b"], "my_index": [0, 0]}
+            ).set_index("my_index"),
+            pd.DataFrame(
+                {
+                    "col": [1, 2, 3, 4],
+                    "other_col": ["a", "a", "b", "b"],
+                    "my_index": [0, 0, 0, 0],
+                },
+                dtype=object,
+            ).set_index("my_index"),
+        ),
+        (
+            pd.DataFrame(
+                {
+                    "col": [[1, 2], [3, 4]],
+                    "other_col": ["a", "b"],
+                    "my_first_index": [0, 0],
+                    "my_second_index": [1, 1],
+                }
+            ).set_index(["my_first_index", "my_second_index"]),
+            pd.DataFrame(
+                {
+                    "col": [1, 2, 3, 4],
+                    "other_col": ["a", "a", "b", "b"],
+                    "my_first_index": [0, 0, 0, 0],
+                    "my_second_index": [1, 1, 1, 1],
+                },
+                dtype=object,
+            ).set_index(["my_first_index", "my_second_index"]),
+        ),
+        (
+            pd.DataFrame(
+                {"col": [[1, 2], [3, 4]], "other_col": ["a", "b"]},
+                pd.MultiIndex.from_tuples([(0, 1), (0, 1)]),
+            ),
+            pd.DataFrame(
+                {"col": [1, 2, 3, 4], "other_col": ["a", "a", "b", "b"]},
+                pd.MultiIndex.from_tuples([(0, 1), (0, 1), (0, 1), (0, 1)]),
+                dtype=object,
+            ),
+        ),
+        (
+            pd.DataFrame(
+                {"col": [[1, 2], [3, 4]], "other_col": ["a", "b"]},
+                pd.MultiIndex.from_arrays([[0, 0], [1, 1]], names=["foo", None]),
+            ),
+            pd.DataFrame(
+                {"col": [1, 2, 3, 4], "other_col": ["a", "a", "b", "b"]},
+                pd.MultiIndex.from_arrays(
+                    [[0, 0, 0, 0], [1, 1, 1, 1]], names=["foo", None]
+                ),
+                dtype=object,
+            ),
+        ),
+    ],
+)
+def test_duplicate_index(df, expected):
+    # GH 28005
+    result = df.explode("col")
+    tm.assert_frame_equal(result, expected)

pandas/tests/series/test_explode.py

@@ -111,3 +111,11 @@ def test_nested_EA():
         pd.date_range("20170101", periods=6, tz="UTC"), index=[0, 0, 0, 1, 1, 1]
     )
     tm.assert_series_equal(result, expected)
+
+
+def test_duplicate_index():
+    # GH 28005
+    s = pd.Series([[1, 2], [3, 4]], index=[0, 0])
+    result = s.explode()
+    expected = pd.Series([1, 2, 3, 4], index=[0, 0, 0, 0], dtype=object)
+    tm.assert_series_equal(result, expected)