Repository: modin-project/modin

Issue:
Implement 2D insertions of Modin DataFrames in `DataFrame.__setitem__`
Currently, attempting to insert a 2D modin DataFrame via `.__setitem__` results in falling back to pandas:
```python
import modin.pandas as pd

df = pd.DataFrame({"a": [1, 2, 3], "b": [3, 4, 5]})
df2 = pd.DataFrame({"new_col1": [10, 20, 30], "new_col2": [30, 40, 50]})

df[["new_col1", "new_col2"]] = df2
print(df)
```
2D insertion is already implemented in lower levels ([`query_compiler.insert_item`](https://github.com/modin-project/modin/blob/08ec2e1397cf48fd26cb2a69fb378bb2f38c5d16/modin/core/storage_formats/base/query_compiler.py#L3109)) so it should be relatively easy to add support for this at the high-level API layer.



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
@@ -2527,6 +2527,27 @@ def __setitem__(self, key, value):
                 self._update_inplace(new_qc)
                 # self.loc[:, key] = value
                 return
+            elif (
+                isinstance(key, list)
+                and isinstance(value, type(self))
+                # Mixed case is more complicated, it's defaulting to pandas for now
+                and all((x not in self.columns for x in key))
+            ):
+                if len(key) != len(value.columns):
+                    raise ValueError("Columns must be same length as key")
+
+                # Aligning the value's columns with the key
+                if not np.array_equal(value.columns, key):
+                    value = value.set_axis(key, axis=1)
+
+                new_qc = self._query_compiler.insert_item(
+                    axis=1,
+                    loc=len(self.columns),
+                    value=value._query_compiler,
+                    how="left",
+                )
+                self._update_inplace(new_qc)
+                return
 
             def setitem_unhashable_key(df, value):
                 df[key] = value


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
@@ -2047,6 +2047,57 @@ def _make_copy(df1, df2):
         eval_setitem(modin_df, pandas_df, df_value[["value_col1"]], key)
 
 
+def test_setitem_2d_insertion():
+    def build_value_picker(modin_value, pandas_value):
+        """Build a function that returns either Modin or pandas DataFrame depending on the passed frame."""
+        return (
+            lambda source_df, *args, **kwargs: modin_value
+            if isinstance(source_df, (pd.DataFrame, pd.Series))
+            else pandas_value
+        )
+
+    modin_df, pandas_df = create_test_dfs(test_data["int_data"])
+
+    # Easy case - key and value.columns are equal
+    modin_value, pandas_value = create_test_dfs(
+        {"new_value1": np.arange(len(modin_df)), "new_value2": np.arange(len(modin_df))}
+    )
+    eval_setitem(
+        modin_df,
+        pandas_df,
+        build_value_picker(modin_value, pandas_value),
+        col=["new_value1", "new_value2"],
+    )
+
+    # Key and value.columns have equal values but in different order
+    new_columns = ["new_value3", "new_value4"]
+    modin_value.columns, pandas_value.columns = new_columns, new_columns
+    eval_setitem(
+        modin_df,
+        pandas_df,
+        build_value_picker(modin_value, pandas_value),
+        col=["new_value4", "new_value3"],
+    )
+
+    # Key and value.columns have different values
+    new_columns = ["new_value5", "new_value6"]
+    modin_value.columns, pandas_value.columns = new_columns, new_columns
+    eval_setitem(
+        modin_df,
+        pandas_df,
+        build_value_picker(modin_value, pandas_value),
+        col=["__new_value5", "__new_value6"],
+    )
+
+    # Key and value.columns have different lengths, testing that both raise the same exception
+    eval_setitem(
+        modin_df,
+        pandas_df,
+        build_value_picker(modin_value.iloc[:, [0]], pandas_value.iloc[:, [0]]),
+        col=["new_value7", "new_value8"],
+    )
+
+
 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