Support duplicate columns and add test case

Signed-off-by: Michael Tiemann <[email protected]>
hgrecco · Oct 21, 2023 · 2808bb8 · 2808bb8
1 parent aa3c9f1
commit 2808bb8
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 7 deletions.
diff --git a/pint_pandas/pint_array.py b/pint_pandas/pint_array.py
@@ -1,7 +1,6 @@
 import copy
 import re
 import warnings
-from collections import OrderedDict
 from importlib.metadata import version
 
 import numpy as np
@@ -987,17 +986,29 @@ def formatter_func(dtype):
             for i, col in enumerate(df.columns)
         ]
 
-        data_for_df = OrderedDict()
+        data_for_df = []
         for i, col in enumerate(df.columns):
             if isinstance(df.dtypes.iloc[i], PintType):
-                data_for_df[tuple(df_columns.iloc[i])] = df.iloc[:, i].values.data
+                data_for_df.append(
+                    pd.Series(
+                        data=df.iloc[:, i].values.data,
+                        name=tuple(df_columns.iloc[i]),
+                        index=df.index,
+                        copy=False,
+                    )
+                )
             else:
-                data_for_df[tuple(df_columns.iloc[i])] = df.iloc[:, i].values
-
-        df_new = DataFrame(data_for_df, columns=data_for_df.keys())
+                data_for_df.append(
+                    pd.Series(
+                        data=df.iloc[:, i].values,
+                        name=tuple(df_columns.iloc[i]),
+                        index=df.index,
+                        copy=False,
+                    )
+                )
 
+        df_new = pd.concat(data_for_df, axis=1, copy=False)
         df_new.columns.names = df.columns.names + ["unit"]
-        df_new.index = df.index
 
         return df_new
 

diff --git a/pint_pandas/testsuite/test_issues.py b/pint_pandas/testsuite/test_issues.py
@@ -194,3 +194,38 @@ def test_issue_194(dtype):
     s2 = s1.astype(dtype)
 
     tm.assert_series_equal(s0, s2)
+
+
+class TestIssue202(BaseExtensionTests):
+    def test_dequantify(self):
+        df = pd.DataFrame()
+        df["test"] = pd.Series([1, 2, 3], dtype="pint[kN]")
+        df.insert(0, "test", df["test"], allow_duplicates=True)
+
+        expected = pd.DataFrame.from_dict(
+            data={
+                "index": [0, 1, 2],
+                "columns": [("test", "kilonewton")],
+                "data": [[1], [2], [3]],
+                "index_names": [None],
+                "column_names": [None, "unit"],
+            },
+            orient="tight",
+            dtype="Int64",
+        )
+        result = df.iloc[:, 1:].pint.dequantify()
+        tm.assert_frame_equal(expected, result)
+
+        expected = pd.DataFrame.from_dict(
+            data={
+                "index": [0, 1, 2],
+                "columns": [("test", "kilonewton"), ("test", "kilonewton")],
+                "data": [[1, 1], [2, 2], [3, 3]],
+                "index_names": [None],
+                "column_names": [None, "unit"],
+            },
+            orient="tight",
+            dtype="Int64",
+        )
+        result = df.pint.dequantify()
+        tm.assert_frame_equal(expected, result)