fix(rust, python): streaming hstack allow duplicates (#5538)

polars/polars-core/src/frame/mod.rs

@@ -1172,6 +1172,19 @@ impl DataFrame {
         Ok(())
     }
 
+    pub fn _add_columns(&mut self, columns: Vec<Series>, schema: &Schema) -> PolarsResult<()> {
+        for (i, s) in columns.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() {
+                self.with_column_and_schema(s, schema)?;
+            } else {
+                self.with_column(s.clone())?;
+            }
+        }
+        Ok(())
+    }
+
     /// Add a new column to this `DataFrame` or replace an existing one.
     /// Uses an existing schema to amortize lookups.
     /// If the schema is incorrect, we will fallback to linear search.

polars/polars-lazy/polars-pipe/src/executors/operators/projection.rs

@@ -2,6 +2,7 @@ use std::sync::Arc;
 
 use polars_core::error::PolarsResult;
 use polars_core::frame::DataFrame;
+use polars_core::schema::SchemaRef;
 
 use crate::expressions::PhysicalPipedExpr;
 use crate::operators::{DataChunk, Operator, OperatorResult, PExecutionContext};
@@ -53,6 +54,7 @@ impl Operator for ProjectionOperator {
 #[derive(Clone)]
 pub(crate) struct HstackOperator {
     pub(crate) exprs: Vec<Arc<dyn PhysicalPipedExpr>>,
+    pub(crate) input_schema: SchemaRef,
 }
 
 impl Operator for HstackOperator {
@@ -67,7 +69,9 @@ impl Operator for HstackOperator {
             .map(|e| e.evaluate(chunk, context.execution_state.as_ref()))
             .collect::<PolarsResult<Vec<_>>>()?;
 
-        let df = chunk.data.hstack(&projected)?;
+        let mut df = chunk.data.clone();
+        let schema = &*self.input_schema;
+        df._add_columns(projected, schema)?;
 
         let chunk = chunk.with_data(df);
         Ok(OperatorResult::Finished(chunk))

polars/polars-lazy/polars-pipe/src/pipeline/convert.rs

@@ -221,9 +221,10 @@ where
             };
             Box::new(op) as Box<dyn Operator>
         }
-        HStack { exprs, .. } => {
+        HStack { exprs, schema, .. } => {
             let op = operators::HstackOperator {
                 exprs: exprs_to_physical(exprs, expr_arena, &to_physical)?,
+                input_schema: schema.clone(),
             };
             Box::new(op) as Box<dyn Operator>
         }

polars/polars-lazy/src/physical_plan/executors/stack.rs

@@ -30,15 +30,7 @@ impl StackExec {
         state.clear_expr_cache();
 
         let schema = &*self.input_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())?;
-            }
-        }
+        df._add_columns(res, schema)?;
 
         Ok(df)
     }

py-polars/tests/unit/test_projections.py

@@ -94,3 +94,13 @@ def test_unnest_columns_available() -> None:
         "genre3": ["Fantasy", "Thriller", "Drama"],
         "genre4": ["Sci-Fi", None, "Romance"],
     }
+
+
+def test_streaming_duplicate_cols_5537() -> None:
+    assert pl.DataFrame({"a": [1, 2, 3], "b": [1, 2, 3]}).lazy().with_columns(
+        [(pl.col("a") * 2).alias("foo"), (pl.col("a") * 3)]
+    ).collect(allow_streaming=True).to_dict(False) == {
+        "a": [3, 6, 9],
+        "b": [1, 2, 3],
+        "foo": [2, 4, 6],
+    }