with_columns update on duplicates (#4179)

pola-rs · Jul 29, 2022 · 56cc086 · 56cc086
1 parent f5617af
commit 56cc086
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/polars/polars-lazy/src/physical_plan/executors/stack.rs b/polars/polars-lazy/src/physical_plan/executors/stack.rs
@@ -32,8 +32,14 @@ impl Executor for StackExec {
         state.clear_expr_cache();
 
         let schema = &*self.input_schema;
-        for s in res {
-            df.with_column_and_schema(s, schema)?;
+        for (i, s) in res.into_iter().enumerate() {
+            // we need to branch here
+            // because users can add multiple columns with the same name
+            if i == 0 || schema.get(s.name()).is_some() {
+                df.with_column_and_schema(s, schema)?;
+            } else {
+                df.with_column(s.clone())?;
+            }
         }
 
         Ok(df)

diff --git a/py-polars/tests/test_schema.py b/py-polars/tests/test_schema.py
@@ -31,3 +31,8 @@ def test_fill_null_minimal_upcast_4056() -> None:
     df = df.with_columns(pl.col("a").cast(pl.Int8))
     assert df.with_column(pl.col(pl.Int8).fill_null(-1)).dtypes[0] == pl.Int8
     assert df.with_column(pl.col(pl.Int8).fill_null(-1000)).dtypes[0] == pl.Int32
+
+
+def test_with_column_duplicates() -> None:
+    df = pl.DataFrame({"a": [0, None, 2, 3, None], "b": [None, 1, 2, 3, None]})
+    assert df.with_columns([pl.all().alias("same")]).columns == ["a", "b", "same"]