PERF-modin-project#7123: Preserve shape_hint for dropna (modin-projec…

…t#7124) Signed-off-by: Iaroslav Igoshev <iaroslav.igoshev@intel.com>
sfc-gh-mvashishtha · Mar 26, 2024 · a616e4c · a616e4c
1 parent 9055fd7
commit a616e4c
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/modin/core/storage_formats/pandas/query_compiler.py b/modin/core/storage_formats/pandas/query_compiler.py
@@ -3005,13 +3005,14 @@ def reduce(df: pandas.DataFrame, mask: pandas.DataFrame):
                 copartition=False,
                 labels="drop",
             )
-            return self.__constructor__(result)
+            return self.__constructor__(result, shape_hint=self._shape_hint)
 
         return self.__constructor__(
             self._modin_frame.filter(
                 kwargs.get("axis", 0) ^ 1,
                 lambda df: pandas.DataFrame.dropna(df, **kwargs),
-            )
+            ),
+            shape_hint=self._shape_hint,
         )
 
     def drop(self, index=None, columns=None, errors: str = "raise"):