Repository: modin-project/modin

Issue:
BUG: passing a function returning a series to Series.apply raises KeyError: 'DataFrame'
### System information
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: macOS Monterey
- **Modin version** (`modin.__version__`): 05bf659b67d2b7153d8aaa004dd95252377eed51
- **Python version**: 3.10.4
- **Code we can use to reproduce**:

```python
import modin.pandas as pd
import pandas

pd.Series(1).apply(lambda x: pandas.Series([x, x]))
```

### Describe the problem
In pandas the above script gives a dataframe:

```
   0  1
0  1  1
```

In modin I get an error from `result = globals()[return_type](query_compiler=result)`. The string `"DataFrame"` isn't in `globals()`.

### Source code / logs
<details>

<summary>Stack trace</summary>


```
UserWarning: Distributing <class 'int'> object. This may take some time.
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Input In [3], in <cell line: 4>()
      1 import modin.pandas as pd
      2 import pandas
----> 4 pd.Series(1).apply(lambda x: pandas.Series([x, x]))

File ~/software_sources/modin/modin/logging/logger_decorator.py:128, in enable_logging.<locals>.decorator.<locals>.run_and_log(*args, **kwargs)
    113 """
    114 Compute function with logging if Modin logging is enabled.
    115
   (...)
    125 Any
    126 """
    127 if LogMode.get() == "disable":
--> 128     return obj(*args, **kwargs)
    130 logger = get_logger()
    131 logger_level = getattr(logger, log_level)

File ~/software_sources/modin/modin/_compat/pandas_api/latest/series.py:28, in LatestCompatibleSeries.apply(self, func, convert_dtype, args, **kwargs)
     27 def apply(self, func, convert_dtype=True, args=(), **kwargs):
---> 28     return self._apply(func, convert_dtype=convert_dtype, args=args, **kwargs)

File ~/software_sources/modin/modin/logging/logger_decorator.py:128, in enable_logging.<locals>.decorator.<locals>.run_and_log(*args, **kwargs)
    113 """
    114 Compute function with logging if Modin logging is enabled.
    115
   (...)
    125 Any
    126 """
    127 if LogMode.get() == "disable":
--> 128     return obj(*args, **kwargs)
    130 logger = get_logger()
    131 logger_level = getattr(logger, log_level)

File ~/software_sources/modin/modin/pandas/series.py:663, in Series._apply(self, func, convert_dtype, args, **kwargs)
    661         return result
    662 else:
--> 663     result = globals()[return_type](query_compiler=result)
    664     if result.name == self.index[0]:
    665         result.name = None

KeyError: 'DataFrame'
```

</details>


Discussion:
This is an interesting bug. The reason why we are not seeing `DataFrame` in globals is because it's not being imported properly. If we do import and check the type properly by doing something like this:
```python
     ...
    else:
            from .dataframe import DataFrame
            
            if return_type == "DataFrame":
                result = DataFrame(query_compiler=result)
                
            elif return_type == "Series":
                result = Series(query_compiler=result)
                if result.name == self.index[0]:
                    result.name = None
            else:
                raise TypeError("return type is somehow not DataFrame or Series.")
            
            return result
```
I get an output like this:
```Out[3]: 
                  __reduced__
0  0    1
1    1
dtype: int64
```
Which doesn't seem fully correct. IIRC, `__reduced__` means the Series column does not have a name.
@pyrito yes, your fix does get the type we want! However, I see that the result is a 1x1 frame whose only value is `pandas.Series([1, 1])`. We get that result because when the `apply` function returns a `DataFrame`, we return the result of `Series.map`:

https://github.com/modin-project/modin/blob/3f985ed6864cc1b5b587094d75ca5b2695e4139f/modin/pandas/series.py#L651-L663

but mapping every element of the series `pands.Series(1)` gives the 1x1 dataframe.

I don't know yet how to solve this problem.
After investigating this a bit more I think what we are doing is not _fundamentally_ incorrect, besides the fact that we have `__reduced__` was the name for one of the DataFrame columns. If we don't set it as that, we end up getting this as an output:
```
                            0
0  0    1
1    1
dtype: int64
```
which is also incorrect as you have stated @mvashishtha . However, after reading the code, it looks like we are doing exactly this:
```python
y = lambda x: pandas.Series([x, x])
print(pandas.DataFrame(pandas.Series(1)).applymap(lambda s: y(s)))
```
This gives us the same output as the one mentioned above. I tried some other different apply combinations with pandas and I was getting some weird results, so I'm not sure what the best approach is here. For example if I try the following:

```python
y = lambda x: pandas.Series([x, x])
print(pandas.DataFrame(pandas.Series(1)).apply(lambda s: y(s), axis=0))
```
I get this as the output:
```
                              0
0  0    1
Name: 0, dtype: int64
1  0    1
Name: 0, dtype: int64
```

Resolving patch:
diff --git a/modin/core/storage_formats/base/query_compiler.py b/modin/core/storage_formats/base/query_compiler.py
--- a/modin/core/storage_formats/base/query_compiler.py
+++ b/modin/core/storage_formats/base/query_compiler.py
@@ -2254,6 +2254,32 @@ def apply(self, func, axis, raw=False, result_type=None, *args, **kwargs):
             **kwargs,
         )
 
+    def apply_on_series(self, func, *args, **kwargs):
+        """
+        Apply passed function on underlying Series.
+
+        Parameters
+        ----------
+        func : callable(pandas.Series) -> scalar, str, list or dict of such
+            The function to apply to each row.
+        *args : iterable
+            Positional arguments to pass to `func`.
+        **kwargs : dict
+            Keyword arguments to pass to `func`.
+
+        Returns
+        -------
+        BaseQueryCompiler
+        """
+        assert self.is_series_like()
+
+        return SeriesDefault.register(pandas.Series.apply)(
+            self,
+            func=func,
+            *args,
+            **kwargs,
+        )
+
     def explode(self, column):
         """
         Explode the given columns.
diff --git a/modin/core/storage_formats/pandas/query_compiler.py b/modin/core/storage_formats/pandas/query_compiler.py
--- a/modin/core/storage_formats/pandas/query_compiler.py
+++ b/modin/core/storage_formats/pandas/query_compiler.py
@@ -2393,6 +2393,23 @@ def apply(self, func, axis, *args, **kwargs):
         else:
             return self._callable_func(func, axis, *args, **kwargs)
 
+    def apply_on_series(self, func, *args, **kwargs):
+        args = try_cast_to_pandas(args)
+        kwargs = try_cast_to_pandas(kwargs)
+
+        assert self.is_series_like()
+
+        # We use apply_full_axis here instead of map since the latter assumes that the
+        # shape of the DataFrame does not change. However, it is possible for functions
+        # applied to Series objects to end up creating DataFrames. It is possible that
+        # using apply_full_axis is much less performant compared to using a variant of
+        # map.
+        return self.__constructor__(
+            self._modin_frame.apply_full_axis(
+                1, lambda df: df.squeeze(axis=1).apply(func, *args, **kwargs)
+            )
+        )
+
     def _dict_func(self, func, axis, *args, **kwargs):
         """
         Apply passed functions to the specified rows/columns.
diff --git a/modin/pandas/series.py b/modin/pandas/series.py
--- a/modin/pandas/series.py
+++ b/modin/pandas/series.py
@@ -671,19 +671,29 @@ def f(x):
             with np.errstate(all="ignore"):
                 if isinstance(f, np.ufunc):
                     return f(self)
-                result = self.map(f)._query_compiler
-        if return_type not in ["DataFrame", "Series"]:
-            # sometimes result can be not a query_compiler, but scalar (for example
-            # for sum or count functions)
-            if isinstance(result, type(self._query_compiler)):
-                return result.to_pandas().squeeze()
-            else:
-                return result
-        else:
-            result = globals()[return_type](query_compiler=result)
+
+                # The return_type is only a DataFrame when we have a function
+                # return a Series object. This is a very particular case that
+                # has to be handled by the underlying pandas.Series apply
+                # function and not our default applymap call.
+                if return_type == "DataFrame":
+                    result = self._query_compiler.apply_on_series(f)
+                else:
+                    result = self.map(f)._query_compiler
+
+        if return_type == "DataFrame":
+            from .dataframe import DataFrame
+
+            result = DataFrame(query_compiler=result)
+        elif return_type == "Series":
+            result = Series(query_compiler=result)
             if result.name == self.index[0]:
                 result.name = None
-            return result
+        elif isinstance(result, type(self._query_compiler)):
+            # sometimes result can be not a query_compiler, but scalar (for example
+            # for sum or count functions)
+            return result.to_pandas().squeeze()
+        return result
 
     def argmax(self, axis=None, skipna=True, *args, **kwargs):  # noqa: PR01, RT01, D200
         """


Test patch:
diff --git a/modin/pandas/test/test_series.py b/modin/pandas/test/test_series.py
--- a/modin/pandas/test/test_series.py
+++ b/modin/pandas/test/test_series.py
@@ -4711,3 +4711,15 @@ def func(val):
     modin_series = modin_df["col"].apply(func)
 
     df_equals(modin_series, pandas_series)
+
+
+@pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)
+def test_apply_return_df(data):
+    modin_series, pandas_series = create_test_series(data)
+    eval_general(
+        modin_series,
+        pandas_series,
+        lambda series: series.apply(
+            lambda x: pandas.Series([x + i for i in range(100)])
+        ),
+    )