From d65387e6dbb7b894f6c3eb2c4936cacf190ac9a2 Mon Sep 17 00:00:00 2001
From: Dmitry Chigarev <dmitry.chigarev@intel.com>
Date: Wed, 5 Jul 2023 10:48:32 +0000
Subject: [PATCH] Remove high-level imports in algebra operators

Signed-off-by: Dmitry Chigarev <dmitry.chigarev@intel.com>
---
 modin/core/dataframe/algebra/binary.py      |  4 +---
 modin/core/dataframe/algebra/operator.py    | 13 ++-----------
 modin/core/dataframe/algebra/reduce.py      |  7 ++-----
 modin/core/dataframe/algebra/tree_reduce.py |  6 ++----
 4 files changed, 7 insertions(+), 23 deletions(-)

diff --git a/modin/core/dataframe/algebra/binary.py b/modin/core/dataframe/algebra/binary.py
index c3600d385ce..8f86a6b18fa 100644
--- a/modin/core/dataframe/algebra/binary.py
+++ b/modin/core/dataframe/algebra/binary.py
@@ -442,8 +442,6 @@ 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
@@ -451,7 +449,7 @@ def apply(
         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
index 393b83f99a8..344d160ed97 100644
--- 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
index 25f8755f6a0..583533ce8dd 100644
--- 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
index 479bf217178..45ae29c4f6d 100644
--- 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)