Skip to content

Commit

Permalink
remove extra squeeze length check
Browse files Browse the repository at this point in the history
  • Loading branch information
noloerino committed Sep 19, 2024
1 parent 6cf3ca2 commit d6c78ae
Showing 1 changed file with 11 additions and 8 deletions.
19 changes: 11 additions & 8 deletions modin/pandas/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -2074,16 +2074,19 @@ def squeeze(
Squeeze 1 dimensional axis objects into scalars.
"""
axis = self._get_axis_number(axis) if axis is not None else None
if axis is None and (len(self.columns) == 1 or len(self.index) == 1):
return Series(query_compiler=self._query_compiler).squeeze()
if axis == 1 and len(self.columns) == 1:
len_columns = len(self.columns)
if axis == 1 and len_columns == 1:
self._query_compiler._shape_hint = "column"
return Series(query_compiler=self._query_compiler)
if axis == 0 and len(self.index) == 1:
qc = self.T._query_compiler
qc._shape_hint = "column"
return Series(query_compiler=qc)
else:
if axis in [0, None]:
# Only compute the length of the index if axis is 0 or None.
len_index = len(self)
if axis is None and (len_columns == 1 or len_index == 1):
return Series(query_compiler=self._query_compiler).squeeze()
if axis == 0 and len_index == 1:
qc = self.T._query_compiler
qc._shape_hint = "column"
return Series(query_compiler=qc)
return self.copy()

def stack(
Expand Down

0 comments on commit d6c78ae

Please sign in to comment.