FEAT-modin-project#2375: 'squeeze' warning fixed

Signed-off-by: Dmitry Chigarev <dmitry.chigarev@intel.com>
dchigarev · Dec 8, 2020 · fad766d · fad766d
1 parent bc921c8
commit fad766d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/modin/engines/base/frame/data.py b/modin/engines/base/frame/data.py
@@ -1736,7 +1736,7 @@ def broadcast_apply_full_axis(
             None,
             dtypes,
             validate_axes="all"
-            if any(o is None for o in [new_index, new_columns])
+            if any(o is not None for o in [new_index, new_columns])
             else False,
         )
 

diff --git a/modin/pandas/groupby.py b/modin/pandas/groupby.py
@@ -91,8 +91,8 @@ def __init__(
             "sort": sort,
             "as_index": as_index,
             "group_keys": group_keys,
-            "squeeze": squeeze,
         }
+        self._squeeze = squeeze
         self._kwargs.update(kwargs)
 
     _index_grouped_cache = None
@@ -283,7 +283,7 @@ def backfill(self, limit=None):
         return self.bfill(limit)
 
     def __getitem__(self, key):
-        kwargs = self._kwargs.copy()
+        kwargs = {**self._kwargs.copy(), "squeeze": self._squeeze}
         # Most of time indexing DataFrameGroupBy results in another DataFrameGroupBy object unless circumstances are
         # special in which case SeriesGroupBy has to be returned. Such circumstances are when key equals to a single
         # column name and is not a list of column names or list of one column name.
@@ -484,6 +484,7 @@ def size(self):
                     self._axis,
                     drop=False,
                     idx_name=None,
+                    squeeze=self._squeeze,
                     **self._kwargs,
                 )
                 result = work_object._wrap_aggregation(
@@ -506,6 +507,7 @@ def size(self):
                 self._axis,
                 drop=False,
                 idx_name=None,
+                squeeze=self._squeeze,
                 **self._kwargs,
             )
             result = work_object._wrap_aggregation(
@@ -532,6 +534,7 @@ def size(self):
                 0,
                 drop=self._drop,
                 idx_name=self._idx_name,
+                squeeze=self._squeeze,
                 **self._kwargs,
             ).size()
 
@@ -871,7 +874,7 @@ def _wrap_aggregation(
                 drop=self._drop,
             )
         )
-        if self._kwargs.get("squeeze", False):
+        if self._squeeze:
             return result.squeeze()
         return result
 
@@ -917,7 +920,7 @@ def _apply_agg_function(self, f, drop=True, *args, **kwargs):
         result = type(self._df)(query_compiler=new_manager)
         if result._query_compiler.get_index_name() == "__reduced__":
             result._query_compiler.set_index_name(None)
-        if self._kwargs.get("squeeze", False):
+        if self._squeeze:
             return result.squeeze()
         return result
 
@@ -950,7 +953,11 @@ def _default_to_pandas(self, f, *args, **kwargs):
 
         def groupby_on_multiple_columns(df, *args, **kwargs):
             return f(
-                df.groupby(by=by, axis=self._axis, **self._kwargs), *args, **kwargs
+                df.groupby(
+                    by=by, axis=self._axis, squeeze=self._squeeze, **self._kwargs
+                ),
+                *args,
+                **kwargs,
             )
 
         return self._df._default_to_pandas(groupby_on_multiple_columns, *args, **kwargs)