Remove high-level imports in algebra operators

Signed-off-by: Dmitry Chigarev <dmitry.chigarev@intel.com>
dchigarev · Jul 5, 2023 · 9614b2b · 9614b2b
1 parent ee354fd
commit 9614b2b
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 23 deletions.
diff --git a/modin/core/dataframe/algebra/binary.py b/modin/core/dataframe/algebra/binary.py
@@ -442,16 +442,14 @@ def apply(
         -------
         The same type as `df`.
         """
-        from modin.pandas import Series
-
         operator = cls.register(func, **kwargs)
 
         func_args = tuple() if func_args is None else func_args
         func_kwargs = dict() if func_kwargs is None else func_kwargs
         qc_result = operator(
             left._query_compiler,
             right._query_compiler,
-            broadcast=isinstance(right, Series),
+            broadcast=right.ndim == 1,
             *func_args,
             axis=axis,
             **func_kwargs,

diff --git a/modin/core/dataframe/algebra/operator.py b/modin/core/dataframe/algebra/operator.py
@@ -62,9 +62,7 @@ def validate_axis(cls, axis: Optional[int]) -> int:
         return 0 if axis is None else axis
 
     @classmethod
-    def apply(
-        cls, df, func, func_args=None, func_kwargs=None, _return_type=None, **kwargs
-    ):
+    def apply(cls, df, func, func_args=None, func_kwargs=None, **kwargs):
         r"""
         Apply a function to a Modin DataFrame using the operators scheme.
 
@@ -78,9 +76,6 @@ def apply(
             Positional arguments to pass to the `func`.
         func_kwargs : dict, optional
             Keyword arguments to pass to the `func`.
-        _return_type : type, optional
-            A class that takes the ``query_compiler`` keyword argument. If not specified
-            will be identical to the type of the passed `df`.
         **kwargs : dict
             Aditional arguments to pass to the ``cls.register()``.
 
@@ -94,8 +89,4 @@ def apply(
         func_kwargs = dict() if func_kwargs is None else func_kwargs
 
         qc_result = operator(df._query_compiler, *func_args, **func_kwargs)
-
-        if _return_type is None:
-            _return_type = type(df)
-
-        return _return_type(query_compiler=qc_result)
+        return type(df)(query_compiler=qc_result)
diff --git a/modin/core/dataframe/algebra/reduce.py b/modin/core/dataframe/algebra/reduce.py
@@ -73,8 +73,5 @@ def apply(cls, df, func, axis=0, func_args=None, func_kwargs=None):
         -------
         modin.pandas.Series
         """
-        from modin.pandas import Series
-
-        return super().apply(
-            df, func, func_args, func_kwargs, axis=axis, _return_type=Series
-        )
+        result = super().apply(df, func, func_args, func_kwargs, axis=axis)
+        return result if result.ndim == 1 else result.squeeze(axis)
diff --git a/modin/core/dataframe/algebra/tree_reduce.py b/modin/core/dataframe/algebra/tree_reduce.py
@@ -81,14 +81,12 @@ def apply(
         -------
         modin.pandas.Series
         """
-        from modin.pandas import Series
-
-        return super().apply(
+        result = super().apply(
             df,
             map_function,
             func_args,
             func_kwargs,
             reduce_function=reduce_function,
             axis=axis,
-            _return_type=Series,
         )
+        return result if result.ndim == 1 else result.squeeze(axis)