Repository: pola-rs/polars (language: python)

Issue:
Alternative ``concat`` way several times faster than ``with_columns()`` when adding many columns to the df
### Checks

- [X] I have checked that this issue has not already been reported.
- [X] I have confirmed this bug exists on the [latest version](https://pypi.org/project/polars/) of Polars.

### Reproducible example

```python
import polars as pl
import numpy as np

# Set the number of rows and columns
num_rows = 10000
num_columns = 10000

np.random.seed(42)

data1 = {f"col_{i}": np.random.rand(num_rows) for i in range(num_columns)}
df1 = pl.DataFrame(data1)

data2 = {f"feature_{i}": np.random.rand(num_rows) for i in range(num_columns)}
df2=pl.DataFrame(data2)

%timeit df1.with_columns(df2)

%timeit pl.concat([df1,df2], how="horizontal")
```
### Log output

```shell
1.28 s ± 31.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
1.85 ms ± 57.8 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
```

### Issue description

So when adding multiple columns, ``concat`` is way faster than the normal ``with_columns()`` method. Would be nice if these two methods were equally as fast.

### Expected behavior

That these functions become equally as fast

### Installed versions

<details>

```
--------Version info---------
Polars:              1.12.0
Index type:          UInt32
Platform:            Windows-11-10.0.22631-SP0
Python:              3.13.0 | packaged by Anaconda, Inc. | (main, Oct  7 2024, 21:21:52) [MSC v.1929 64 bit (AMD64)]
LTS CPU:             False

----Optional dependencies----
adbc_driver_manager  <not installed>
altair               <not installed>
cloudpickle          3.1.0
connectorx           <not installed>
deltalake            <not installed>
fastexcel            <not installed>
fsspec               <not installed>
gevent               <not installed>
great_tables         <not installed>
matplotlib           3.9.2
nest_asyncio         1.6.0
numpy                2.1.2
openpyxl             3.1.5
pandas               2.2.3
pyarrow              18.0.0
pydantic             <not installed>
pyiceberg            <not installed>
sqlalchemy           <not installed>
torch                <not installed>
xlsx2csv             <not installed>
xlsxwriter           <not installed>
```

</details>



Discussion / hints:

Fix patch:
diff --git a/crates/polars-core/src/frame/mod.rs b/crates/polars-core/src/frame/mod.rs
index aa434fb07df7..e3b969a81756 100644
--- a/crates/polars-core/src/frame/mod.rs
+++ b/crates/polars-core/src/frame/mod.rs
@@ -1382,12 +1382,24 @@ impl DataFrame {
         self
     }
 
+    // Note: Schema can be both input or output_schema
     fn add_column_by_schema(&mut self, c: Column, schema: &Schema) -> PolarsResult<()> {
         let name = c.name();
         if let Some((idx, _, _)) = schema.get_full(name.as_str()) {
-            // schema is incorrect fallback to search
             if self.columns.get(idx).map(|s| s.name()) != Some(name) {
-                self.add_column_by_search(c)?;
+                // Given schema is output_schema and we can push.
+                if idx == self.columns.len() {
+                    if self.width() == 0 {
+                        self.height = c.len();
+                    }
+
+                    self.columns.push(c);
+                }
+                // Schema is incorrect fallback to search
+                else {
+                    debug_assert!(false);
+                    self.add_column_by_search(c)?;
+                }
             } else {
                 self.replace_column(idx, c)?;
             }
@@ -1401,6 +1413,7 @@ impl DataFrame {
         Ok(())
     }
 
+    // Note: Schema can be both input or output_schema
     pub fn _add_series(&mut self, series: Vec<Series>, schema: &Schema) -> PolarsResult<()> {
         for (i, s) in series.into_iter().enumerate() {
             // we need to branch here
@@ -1430,6 +1443,8 @@ impl DataFrame {
     /// 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.
+    ///
+    /// Note: Schema can be both input or output_schema
     pub fn with_column_and_schema<C: IntoColumn>(
         &mut self,
         column: C,
diff --git a/crates/polars-mem-engine/src/executors/stack.rs b/crates/polars-mem-engine/src/executors/stack.rs
index ba6fa8111402..a93d4fc72d89 100644
--- a/crates/polars-mem-engine/src/executors/stack.rs
+++ b/crates/polars-mem-engine/src/executors/stack.rs
@@ -8,6 +8,7 @@ pub struct StackExec {
     pub(crate) has_windows: bool,
     pub(crate) exprs: Vec<Arc<dyn PhysicalExpr>>,
     pub(crate) input_schema: SchemaRef,
+    pub(crate) output_schema: SchemaRef,
     pub(crate) options: ProjectionOptions,
     // Can run all operations elementwise
     pub(crate) streamable: bool,
@@ -19,7 +20,7 @@ impl StackExec {
         state: &ExecutionState,
         mut df: DataFrame,
     ) -> PolarsResult<DataFrame> {
-        let schema = &*self.input_schema;
+        let schema = &*self.output_schema;
 
         // Vertical and horizontal parallelism.
         let df = if self.streamable
diff --git a/crates/polars-mem-engine/src/planner/lp.rs b/crates/polars-mem-engine/src/planner/lp.rs
index 3a5e525867fb..0d438b5f5bd1 100644
--- a/crates/polars-mem-engine/src/planner/lp.rs
+++ b/crates/polars-mem-engine/src/planner/lp.rs
@@ -629,7 +629,7 @@ fn create_physical_plan_impl(
         HStack {
             input,
             exprs,
-            schema: _schema,
+            schema: output_schema,
             options,
         } => {
             let input_schema = lp_arena.get(input).schema(lp_arena).into_owned();
@@ -659,6 +659,7 @@ fn create_physical_plan_impl(
                 has_windows: state.has_windows,
                 exprs: phys_exprs,
                 input_schema,
+                output_schema,
                 options,
                 streamable,
             }))


Test patch:
diff --git a/py-polars/tests/benchmark/test_with_columns.py b/py-polars/tests/benchmark/test_with_columns.py
new file mode 100644
index 000000000000..8ea3402ac696
--- /dev/null
+++ b/py-polars/tests/benchmark/test_with_columns.py
@@ -0,0 +1,21 @@
+import time
+
+import pytest
+
+import polars as pl
+
+
+# TODO: this is slow in streaming
+@pytest.mark.may_fail_auto_streaming
+def test_with_columns_quadratic_19503() -> None:
+    num_columns = 2000
+    data1 = {f"col_{i}": [0] for i in range(num_columns)}
+    df1 = pl.DataFrame(data1)
+
+    data2 = {f"feature_{i}": [0] for i in range(num_columns)}
+    df2 = pl.DataFrame(data2)
+
+    t0 = time.time()
+    df1.with_columns(df2)
+    t1 = time.time()
+    assert t1 - t0 < 0.2