FIX-#6774: Modify conditions for loc to get similar behavior to pan…

…das (#6798)

Signed-off-by: JignyasAnand <[email protected]>

modin/pandas/indexing.py

@@ -813,20 +813,20 @@ def _setitem_with_new_columns(self, row_loc, col_loc, item):
         ----------
         row_loc : scalar, slice, list, array or tuple
             Row locator.
-        col_loc : scalar, slice, list, array or tuple
+        col_loc : list, array or tuple
             Columns locator.
         item : modin.pandas.DataFrame, modin.pandas.Series or scalar
             Value that should be assigned to located dataset.
         """
         if is_list_like(item) and not isinstance(item, (DataFrame, Series)):
             item = np.array(item)
             if len(item.shape) == 1:
-                if item.shape[0] != len(col_loc):
+                if len(col_loc) != 1:
                     raise ValueError(
                         "Must have equal len keys and value when setting with an iterable"
                     )
             else:
-                if item.shape != (len(row_loc), len(col_loc)):
+                if item.shape[-1] != len(col_loc):
                     raise ValueError(
                         "Must have equal len keys and value when setting with an iterable"
                     )

modin/pandas/test/dataframe/test_indexing.py

@@ -515,6 +515,31 @@ def test_loc_4456(
     eval_loc(modin_df, pandas_df, (mdf_value, pdf_value), key)
 
 
+def test_loc_6774():
+    modin_df, pandas_df = create_test_dfs(
+        {"a": [1, 2, 3, 4, 5], "b": [10, 20, 30, 40, 50]}
+    )
+    pandas_df.loc[:, "c"] = [10, 20, 30, 40, 51]
+    modin_df.loc[:, "c"] = [10, 20, 30, 40, 51]
+    df_equals(modin_df, pandas_df)
+
+    pandas_df.loc[2:, "y"] = [30, 40, 51]
+    modin_df.loc[2:, "y"] = [30, 40, 51]
+    df_equals(modin_df, pandas_df)
+
+    pandas_df.loc[:, ["b", "c", "d"]] = (
+        pd.DataFrame([[10, 20, 30, 40, 50], [10, 20, 30, 40], [10, 20, 30]])
+        .transpose()
+        .values
+    )
+    modin_df.loc[:, ["b", "c", "d"]] = (
+        pd.DataFrame([[10, 20, 30, 40, 50], [10, 20, 30, 40], [10, 20, 30]])
+        .transpose()
+        .values
+    )
+    df_equals(modin_df, pandas_df)
+
+
 def test_loc_5829():
     data = {"a": [1, 2, 3, 4, 5], "b": [11, 12, 13, 14, 15]}
     modin_df = pd.DataFrame(data, dtype=object)