FIX-modin-project#2442: fixed 'mask' arguments

Signed-off-by: Dmitry Chigarev <[email protected]>
dchigarev · Dec 9, 2020 · 24c92cc · 24c92cc
1 parent 2787ed9
commit 24c92cc
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/modin/backends/pandas/query_compiler.py b/modin/backends/pandas/query_compiler.py
@@ -2168,7 +2168,7 @@ def getitem_row_array(self, key):
         """
         return self.__constructor__(self._modin_frame.mask(row_numeric_idx=key))
 
-    def insert_item(self, axis, loc, value):
+    def insert_item(self, axis, loc, value, **kwargs):
         """
         Insert the column/row defined by `value` at the specified `loc`
 
@@ -2178,17 +2178,22 @@ def insert_item(self, axis, loc, value):
         """
         assert isinstance(value, type(self))
 
+        how = kwargs.get("join", "left")
+
         def execute_concat(left, middle, *right):
             return self.__constructor__(
                 left._concat(
-                    axis=axis, others=[middle, *right], how="outer", sort=False
+                    axis=axis, others=[middle, *right], how=how, sort=False
                 )
             )
+
+        def get_kwargs(value):
+            return {("col_numeric_idx" if axis else "row_numeric_idx"): value}
 
         if 0 < loc < len(self.get_axis(axis)) - 1:
-            first_mask = self._modin_frame.mask(col_numeric_idx=list(range(loc)))
+            first_mask = self._modin_frame.mask(**get_kwargs(list(range(loc))))
             second_mask = self._modin_frame.mask(
-                col_numeric_idx=list(range(loc, len(self.get_axis(axis))))
+                **get_kwargs(list(range(loc, len(self.get_axis(axis)))))
             )
             return execute_concat(first_mask, value._modin_frame, second_mask)
         else:

diff --git a/modin/pandas/dataframe.py b/modin/pandas/dataframe.py
@@ -921,16 +921,13 @@ def insert(self, loc, column, value, allow_duplicates=False):
         if isinstance(value, (DataFrame, pandas.DataFrame)):
             if len(value.columns) != 1:
                 raise ValueError("Wrong number of items passed 2, placement implies 1")
-            value = value.iloc[:, 0]
-
-        #        if isinstance(value, Series)
+            value = value.squeeze(axis=1)
 
         if isinstance(value, Series):
             # TODO: Remove broadcast of Series
             value = value._to_pandas()
 
         if not self._query_compiler.lazy_execution and len(self.index) == 0:
-
             try:
                 value = pandas.Series(value)
             except (TypeError, ValueError, IndexError):