From 8b8806ec25d6c483ceb3b66d378872f71c48ec0a Mon Sep 17 00:00:00 2001 From: Jonathan Shi Date: Fri, 20 Sep 2024 11:46:23 -0700 Subject: [PATCH] REFACTOR-#7315: Refactor axis checks in squeeze (#7400) Signed-off-by: Jonathan Shi --- modin/pandas/dataframe.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modin/pandas/dataframe.py b/modin/pandas/dataframe.py index de96ea0ab26..2ce83913ebb 100644 --- a/modin/pandas/dataframe.py +++ b/modin/pandas/dataframe.py @@ -2074,12 +2074,12 @@ 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): + if axis is None and (len(self.columns) == 1 or len(self) == 1): return Series(query_compiler=self._query_compiler).squeeze() if axis == 1 and len(self.columns) == 1: self._query_compiler._shape_hint = "column" return Series(query_compiler=self._query_compiler) - if axis == 0 and len(self.index) == 1: + if axis == 0 and len(self) == 1: qc = self.T._query_compiler qc._shape_hint = "column" return Series(query_compiler=qc)