fix(rust, python): resolve duplicated column names in pivot (#5349)

pola-rs · Oct 27, 2022 · c2a029d · c2a029d
1 parent df9ca61
commit c2a029d
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 3 deletions.
diff --git a/polars/polars-ops/src/pivot/mod.rs b/polars/polars-ops/src/pivot/mod.rs
@@ -1,5 +1,7 @@
 mod positioning;
 
+use std::borrow::Cow;
+
 use polars_core::export::rayon::prelude::*;
 use polars_core::frame::groupby::expr::PhysicalAggExpr;
 use polars_core::prelude::*;
@@ -196,7 +198,11 @@ fn pivot_impl(
                 };
 
                 let headers = column_agg.unique_stable()?.cast(&DataType::Utf8)?;
-                let headers = headers.utf8().unwrap();
+                let mut headers = headers.utf8().unwrap().clone();
+                if values.len() > 1 {
+                    headers = headers.apply(|v| Cow::from(format!("{}_{}", value_col_name, v)))
+                }
+
                 let n_cols = headers.len();
                 let value_agg_phys = value_agg.to_physical_repr();
                 let logical_type = value_agg.dtype();
@@ -214,7 +220,7 @@ fn pivot_impl(
                                 &col_locations,
                                 $ca,
                                 logical_type,
-                                headers,
+                                &headers,
                             )
                         }};
                     }
@@ -227,7 +233,7 @@ fn pivot_impl(
                         &col_locations,
                         &value_agg_phys,
                         logical_type,
-                        headers,
+                        &headers,
                     )
                 };
 

diff --git a/py-polars/tests/unit/test_pivot.py b/py-polars/tests/unit/test_pivot.py
@@ -124,3 +124,21 @@ def test_pivot_categorical_index() -> None:
         "Car": [1, 2],
         "Ship": [1, None],
     }
+
+
+def test_pivot_multiple_values_column_names_5116() -> None:
+    df = pl.DataFrame(
+        {
+            "x1": [1, 2, 3, 4, 5, 6, 7, 8],
+            "x2": [8, 7, 6, 5, 4, 3, 2, 1],
+            "c1": ["A", "B"] * 4,
+            "c2": ["C", "C", "D", "D"] * 2,
+        }
+    )
+    assert df.pivot(values=["x1", "x2"], index="c1", columns="c2").to_dict(False) == {
+        "c1": ["A", "B"],
+        "x1_C": [1, 2],
+        "x1_D": [3, 4],
+        "x2_C": [8, 7],
+        "x2_D": [6, 5],
+    }