Repository: modin-project/modin

Issue:
PERF: `__setitem__` on multiple columns should be evaluated lazily
NOTE: I haven't done any clear benchmarking for this.

In the following reproducer:
```python
import modin.pandas as pd
import pandas as vpd
import numpy as np

cols = [f"feature_{i}" for i in range(4)]
cols.append('labels')
df = pd.read_csv('data.txt', header=None)
df.columns = cols
val_df = df.sample(frac=0.2)
train_df = df.drop(val_df.index)

train_means = train_df[cols[:-1]].mean()
train_std = train_df[cols[:-1]].std()
train_df[cols[:-1]] = (train_df[cols[:-1]]- train_means)/train_std
```

data: [data.txt](https://github.com/modin-project/modin/files/11969394/data.txt)

`__setitem__` on multiple columns should not materialize partitions immediately.


Discussion:

Resolving patch:
diff --git a/modin/pandas/dataframe.py b/modin/pandas/dataframe.py
--- a/modin/pandas/dataframe.py
+++ b/modin/pandas/dataframe.py
@@ -2570,12 +2570,45 @@ def __setitem__(self, key, value):
                         value = np.array(value)
                     if len(key) != value.shape[-1]:
                         raise ValueError("Columns must be same length as key")
-                new_qc = self._query_compiler.write_items(
-                    slice(None),
-                    self.columns.get_indexer_for(key),
-                    value,
-                    need_columns_reindex=False,
-                )
+                if isinstance(value, type(self)):
+                    # importing here to avoid circular import
+                    from .general import concat
+
+                    if not value.columns.equals(pandas.Index(key)):
+                        # we only need to change the labels, so shallow copy here
+                        value = value.copy(deep=False)
+                        value.columns = key
+
+                    # here we iterate over every column in the 'self' frame, then check if it's in the 'key'
+                    # and so has to be taken from either from the 'value' or from the 'self'. After that,
+                    # we concatenate those mixed column chunks and get a dataframe with updated columns
+                    to_concat = []
+                    # columns to take for this chunk
+                    to_take = []
+                    # whether columns in this chunk are in the 'key' and has to be taken from the 'value'
+                    get_cols_from_value = False
+                    # an object to take columns from for this chunk
+                    src_obj = self
+                    for col in self.columns:
+                        if (col in key) != get_cols_from_value:
+                            if len(to_take):
+                                to_concat.append(src_obj[to_take])
+                            to_take = [col]
+                            get_cols_from_value = not get_cols_from_value
+                            src_obj = value if get_cols_from_value else self
+                        else:
+                            to_take.append(col)
+                    if len(to_take):
+                        to_concat.append(src_obj[to_take])
+
+                    new_qc = concat(to_concat, axis=1)._query_compiler
+                else:
+                    new_qc = self._query_compiler.write_items(
+                        slice(None),
+                        self.columns.get_indexer_for(key),
+                        value,
+                        need_columns_reindex=False,
+                    )
                 self._update_inplace(new_qc)
                 # self.loc[:, key] = value
                 return


Test patch:
diff --git a/modin/pandas/test/dataframe/test_indexing.py b/modin/pandas/test/dataframe/test_indexing.py
--- a/modin/pandas/test/dataframe/test_indexing.py
+++ b/modin/pandas/test/dataframe/test_indexing.py
@@ -2388,6 +2388,44 @@ def build_value_picker(modin_value, pandas_value):
     )
 
 
+@pytest.mark.parametrize("does_value_have_different_columns", [True, False])
+def test_setitem_2d_update(does_value_have_different_columns):
+    def test(dfs, iloc):
+        """Update columns on the given numeric indices."""
+        df1, df2 = dfs
+        cols1 = df1.columns[iloc].tolist()
+        cols2 = df2.columns[iloc].tolist()
+        df1[cols1] = df2[cols2]
+        return df1
+
+    modin_df, pandas_df = create_test_dfs(test_data["int_data"])
+    modin_df2, pandas_df2 = create_test_dfs(test_data["int_data"])
+    modin_df2 *= 10
+    pandas_df2 *= 10
+
+    if does_value_have_different_columns:
+        new_columns = [f"{col}_new" for col in modin_df.columns]
+        modin_df2.columns = new_columns
+        pandas_df2.columns = new_columns
+
+    modin_dfs = (modin_df, modin_df2)
+    pandas_dfs = (pandas_df, pandas_df2)
+
+    eval_general(modin_dfs, pandas_dfs, test, iloc=[0, 1, 2])
+    eval_general(modin_dfs, pandas_dfs, test, iloc=[0, -1])
+    eval_general(
+        modin_dfs, pandas_dfs, test, iloc=slice(1, None)
+    )  # (start=1, stop=None)
+    eval_general(
+        modin_dfs, pandas_dfs, test, iloc=slice(None, -2)
+    )  # (start=None, stop=-2)
+    eval_general(modin_dfs, pandas_dfs, test, iloc=[0, 1, 5, 6, 9, 10, -2, -1])
+    eval_general(modin_dfs, pandas_dfs, test, iloc=[5, 4, 0, 10, 1, -1])
+    eval_general(
+        modin_dfs, pandas_dfs, test, iloc=slice(None, None, 2)
+    )  # (start=None, stop=None, step=2)
+
+
 def test___setitem__single_item_in_series():
     # Test assigning a single item in a Series for issue
     # https://github.com/modin-project/modin/issues/3860