use squeeze_other in qc

noloerino · Sep 12, 2024 · 9116bbd · 9116bbd
1 parent ae5b571
commit 9116bbd
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 13 deletions.
diff --git a/modin/core/storage_formats/base/query_compiler.py b/modin/core/storage_formats/base/query_compiler.py
@@ -655,7 +655,7 @@ def series_eq(self, other, **kwargs):  # noqa: PR02
             self,
             other=other,
             squeeze_self=True,
-            squeeze_other=isinstance(other, pandas.DataFrame),
+            squeeze_other=kwargs.pop("squeeze_other", False),
             **kwargs,
         )
 
@@ -707,7 +707,7 @@ def series_ge(self, other, **kwargs):  # noqa: PR02
             self,
             other=other,
             squeeze_self=True,
-            squeeze_other=isinstance(other, pandas.DataFrame),
+            squeeze_other=kwargs.pop("squeeze_other", False),
             **kwargs,
         )
 
@@ -725,7 +725,7 @@ def series_gt(self, other, **kwargs):  # noqa: PR02
             self,
             other=other,
             squeeze_self=True,
-            squeeze_other=isinstance(other, pandas.DataFrame),
+            squeeze_other=kwargs.pop("squeeze_other", False),
             **kwargs,
         )
 
@@ -745,7 +745,7 @@ def series_le(self, other, **kwargs):  # noqa: PR02
             self,
             other=other,
             squeeze_self=True,
-            squeeze_other=isinstance(other, pandas.DataFrame),
+            squeeze_other=kwargs.pop("squeeze_other", False),
             **kwargs,
         )
 
@@ -763,7 +763,7 @@ def series_lt(self, other, **kwargs):  # noqa: PR02
             self,
             other=other,
             squeeze_self=True,
-            squeeze_other=isinstance(other, pandas.DataFrame),
+            squeeze_other=kwargs.pop("squeeze_other", False),
             **kwargs,
         )
 
@@ -890,7 +890,7 @@ def series_ne(self, other, **kwargs):  # noqa: PR02
             self,
             other=other,
             squeeze_self=True,
-            squeeze_other=isinstance(other, pandas.DataFrame),
+            squeeze_other=kwargs.pop("squeeze_other", False),
             **kwargs,
         )
 

diff --git a/modin/core/storage_formats/pandas/query_compiler.py b/modin/core/storage_formats/pandas/query_compiler.py
@@ -268,7 +268,7 @@ def _series_logical_binop(func):
     """
     return lambda x, y, **kwargs: func(
         x.squeeze(axis=1),
-        y.squeeze(axis=1) if isinstance(y, pandas.DataFrame) else y,
+        y.squeeze(axis=1) if kwargs.pop("squeeze_other", False) else y,
         **kwargs,
     ).to_frame()
 

diff --git a/modin/pandas/series.py b/modin/pandas/series.py
@@ -1032,7 +1032,12 @@ def eq(
         """
         new_self, new_other = self._prepare_inter_op(other)
         return new_self._binary_op(
-            "eq", new_other, level=level, fill_value=fill_value, axis=axis
+            "eq",
+            new_other,
+            level=level,
+            fill_value=fill_value,
+            axis=axis,
+            squeeze_other=isinstance(other, (pandas.Series, Series)),
         )
 
     def equals(self, other) -> bool:  # noqa: PR01, RT01, D200
@@ -1143,7 +1148,12 @@ def ge(
         """
         new_self, new_other = self._prepare_inter_op(other)
         return new_self._binary_op(
-            "ge", new_other, level=level, fill_value=fill_value, axis=axis
+            "ge",
+            new_other,
+            level=level,
+            fill_value=fill_value,
+            axis=axis,
+            squeeze_other=isinstance(other, (pandas.Series, Series)),
         )
 
     def groupby(
@@ -1193,7 +1203,12 @@ def gt(
         """
         new_self, new_other = self._prepare_inter_op(other)
         return new_self._binary_op(
-            "gt", new_other, level=level, fill_value=fill_value, axis=axis
+            "gt",
+            new_other,
+            level=level,
+            fill_value=fill_value,
+            axis=axis,
+            squeeze_other=isinstance(other, (pandas.Series, Series)),
         )
 
     def hist(
@@ -1313,7 +1328,12 @@ def le(
         """
         new_self, new_other = self._prepare_inter_op(other)
         return new_self._binary_op(
-            "le", new_other, level=level, fill_value=fill_value, axis=axis
+            "le",
+            new_other,
+            level=level,
+            fill_value=fill_value,
+            axis=axis,
+            squeeze_other=isinstance(other, (pandas.Series, Series)),
         )
 
     def lt(
@@ -1324,7 +1344,12 @@ def lt(
         """
         new_self, new_other = self._prepare_inter_op(other)
         return new_self._binary_op(
-            "lt", new_other, level=level, fill_value=fill_value, axis=axis
+            "lt",
+            new_other,
+            level=level,
+            fill_value=fill_value,
+            axis=axis,
+            squeeze_other=isinstance(other, (pandas.Series, Series)),
         )
 
     def map(self, arg, na_action=None) -> Series:  # noqa: PR01, RT01, D200
@@ -1453,7 +1478,12 @@ def ne(
         """
         new_self, new_other = self._prepare_inter_op(other)
         return new_self._binary_op(
-            "ne", new_other, level=level, fill_value=fill_value, axis=axis
+            "ne",
+            new_other,
+            level=level,
+            fill_value=fill_value,
+            axis=axis,
+            squeeze_other=isinstance(other, (pandas.Series, Series)),
         )
 
     def nlargest(self, n=5, keep="first") -> Series:  # noqa: PR01, RT01, D200