Repository: modin-project/modin
Pull request: FIX-#4635: allow pass modin functions to `apply`
<!--
Thank you for your contribution!
Please review the contributing docs: https://modin.readthedocs.io/en/latest/development/contributing.html
if you have questions about contributing.
-->

## What do these changes do?

<!-- Please give a short brief about these changes. -->

- [x] first commit message and PR title follow format outlined [here](https://modin.readthedocs.io/en/latest/development/contributing.html#commit-message-formatting)
  > **_NOTE:_**  If you edit the PR title to match this format, you need to add another commit (even if it's empty) or amend your last commit for the CI job that checks the PR title to pick up the new PR title.
- [x] passes `flake8 modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [x] passes `black --check modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [x] signed commit with `git commit -s` <!-- you can amend your commit with a signature via `git commit -amend -s` -->
- [x] Resolves #4635 <!-- issue must be created for each patch -->
- [ ] tests added and passing
- [x] module layout described at `docs/development/architecture.rst` is up-to-date <!-- if you have added, renamed or removed files or directories please update the documentation accordingly -->


Problem statement:
BUG: Can't pass modin functions to modin functions that take a callable parameter
### System information
OS X 11.6.4
Modin version '0.15.2'
Python 3.9.12

### Describe the problem
When I try to use `DataFrame.apply` with `DataFrame.sample`, modin throws an error. Code runs without issues with pandas.

### Source code / logs
```
>>> import modin.pandas as pd
>>> import modin.config as cfg
>>> cfg.Engine.put('Python')
>>> x = pd.DataFrame(data={"a":[1,2,3], "b":[1,2,3], "c":[1,2,3]})
UserWarning: Distributing <class 'dict'> object. This may take some time.
>>> x = x.set_index(['a'])
>>> x.groupby('a', group_keys=False).apply(pd.DataFrame.sample, n=1)
Traceback (most recent call last):
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/pandas/core/generic.py", line 550, in _get_axis_number
    return cls._AXIS_TO_AXIS_NUMBER[axis]
KeyError: None

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/core/storage_formats/pandas/query_compiler.py", line 2805, in groupby_agg_builder
    return compute_groupby(df, drop, partition_idx)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/core/storage_formats/pandas/query_compiler.py", line 2754, in compute_groupby
    result = partition_agg_func(grouped_df, *agg_args, **agg_kwargs)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/pandas/core/groupby/groupby.py", line 1423, in apply
    result = self._python_apply_general(f, self._selected_obj)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/pandas/core/groupby/groupby.py", line 1464, in _python_apply_general
    values, mutated = self.grouper.apply(f, data, self.axis)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/pandas/core/groupby/ops.py", line 761, in apply
    res = f(group)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/pandas/core/groupby/groupby.py", line 1397, in f
    return func(g, *args, **kwargs)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/utils.py", line 521, in wrapper
    result = func(*args, **kwargs)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/logging/logger_metaclass.py", line 68, in log_wrap
    return method(*args, **kwargs)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/pandas/base.py", line 2567, in sample
    axis = self._get_axis_number(axis)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/pandas/core/generic.py", line 552, in _get_axis_number
    raise ValueError(f"No axis named {axis} for object type {cls.__name__}")
ValueError: No axis named None for object type DataFrame

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/pandas/core/generic.py", line 550, in _get_axis_number
    return cls._AXIS_TO_AXIS_NUMBER[axis]
KeyError: None

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/logging/logger_metaclass.py", line 68, in log_wrap
    return method(*args, **kwargs)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/pandas/groupby.py", line 327, in apply
    self._wrap_aggregation(
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/logging/logger_metaclass.py", line 68, in log_wrap
    return method(*args, **kwargs)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/pandas/groupby.py", line 1082, in _wrap_aggregation
    query_compiler=qc_method(
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/logging/logger_metaclass.py", line 68, in log_wrap
    return method(*args, **kwargs)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/core/storage_formats/pandas/query_compiler.py", line 2813, in groupby_agg
    new_modin_frame = self._modin_frame.broadcast_apply_full_axis(
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/logging/logger_metaclass.py", line 68, in log_wrap
    return method(*args, **kwargs)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/core/dataframe/pandas/dataframe/dataframe.py", line 115, in run_f_on_minimally_updated_metadata
    result = f(self, *args, **kwargs)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/core/dataframe/pandas/dataframe/dataframe.py", line 2330, in broadcast_apply_full_axis
    new_partitions = self._partition_mgr_cls.broadcast_axis_partitions(
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/core/dataframe/pandas/partitioning/partition_manager.py", line 463, in broadcast_axis_partitions
    [
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/core/dataframe/pandas/partitioning/partition_manager.py", line 464, in <listcomp>
    left_partitions[i].apply(
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/core/dataframe/pandas/partitioning/axis_partition.py", line 99, in apply
    return self._wrap_partitions(self.deploy_axis_func(*args, **kwargs))
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/core/dataframe/pandas/partitioning/axis_partition.py", line 162, in deploy_axis_func
    result = func(dataframe, *kwargs.pop("args", ()), **kwargs)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/core/dataframe/pandas/dataframe/dataframe.py", line 1393, in _tree_reduce_func
    series_result = func(df, *args, **kwargs)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/core/storage_formats/pandas/query_compiler.py", line 2815, in <lambda>
    func=lambda df, by=None, partition_idx=None: groupby_agg_builder(
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/core/storage_formats/pandas/query_compiler.py", line 2809, in groupby_agg_builder
    return compute_groupby(df.copy(), drop, partition_idx)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/core/storage_formats/pandas/query_compiler.py", line 2754, in compute_groupby
    result = partition_agg_func(grouped_df, *agg_args, **agg_kwargs)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/pandas/core/groupby/groupby.py", line 1423, in apply
    result = self._python_apply_general(f, self._selected_obj)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/pandas/core/groupby/groupby.py", line 1464, in _python_apply_general
    values, mutated = self.grouper.apply(f, data, self.axis)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/pandas/core/groupby/ops.py", line 761, in apply
    res = f(group)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/pandas/core/groupby/groupby.py", line 1397, in f
    return func(g, *args, **kwargs)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/utils.py", line 521, in wrapper
    result = func(*args, **kwargs)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/logging/logger_metaclass.py", line 68, in log_wrap
    return method(*args, **kwargs)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/modin/pandas/base.py", line 2567, in sample
    axis = self._get_axis_number(axis)
  File "/usr/local/Caskroom/miniforge/base/envs/enobase3/lib/python3.9/site-packages/pandas/core/generic.py", line 552, in _get_axis_number
    raise ValueError(f"No axis named {axis} for object type {cls.__name__}")
ValueError: No axis named None for object type DataFrame
```

Hints from issue discussion:
@meijiu thanks for reporting this issue! I was able to reproduce this issue on my Mac. I also observed some other interesting behavior while playing around with other apply functions that may be related to a deeper issue:
```python
x.groupby('a', group_keys=False).apply(pd.DataFrame.sample, n=1) # This FAILS
x.groupby('a', group_keys=False).apply(pandas.DataFrame.sample, n=1) # This works
x.groupby('a', group_keys=False).apply(lambda df: df.sample(n=1)) 3 # This works
```
The root of the error is that the type of the DataFrame we get when applying the function is a pandas.DataFrame. We see this in the logs you've posted and can be further confirmed if we try to run something like the following:
```python
x.groupby('a', group_keys=False).apply(pd.DataFrame.mean)
``` 
This throws the following error:
```
File ~/Documents/modin/modin/pandas/base.py:1946, in BasePandasDataset.mean(self, axis, skipna, level, numeric_only, **kwargs)
   1935 def mean(
   1936     self,
   1937     axis: "int | None | NoDefault" = no_default,
   (...)
   1941     **kwargs,
   1942 ):  # noqa: PR01, RT01, D200
   1943     """
   1944     Return the mean of the values over the requested axis.
   1945     """
-> 1946     return self._stat_operation("mean", axis, skipna, level, numeric_only, **kwargs)

File ~/opt/anaconda3/envs/modin/lib/python3.9/site-packages/pandas/core/generic.py:5575, in NDFrame.__getattr__(self, name)
   5568 if (
   5569     name not in self._internal_names_set
   5570     and name not in self._metadata
   5571     and name not in self._accessors
   5572     and self._info_axis._can_hold_identifiers_and_holds_name(name)
   5573 ):
   5574     return self[name]
-> 5575 return object.__getattribute__(self, name)

AttributeError: 'DataFrame' object has no attribute '_stat_operation'
```
I'm not sure what the exact solution would be here, so I'm thinking off the top of my head here. Since we're applying functions at a partition level at this point, should we just default any Modin functions to pandas if they are being called via `apply`? 

edit: looks like this is a known issue with GroupBy!
@meijiu FYI, this seems to be a known issue and @RehanSD is working on merging in a couple of changes that should address this as well within the next 2-3 weeks. 
@pyrito do you have a link to existing duplicate issues? Could we de-duplicate this one?
Yeah I think we could de-duplicate it. @RehanSD sent me this issue earlier: https://github.com/modin-project/modin/issues/4165
@pyrito this looks like a distinct issue from #4165, which looks like a duplicate of #3435: in those issues, the bug is that the function in the `apply` can't access data from other column partitions. Here, the problem is, as you said, that `apply` gets applied to the inner pandas dataframe rather than to a Modin dataframe. Given that Modin is meant to be a drop-in replacement for pandas, I think it's reasonable for users to expect `apply` functions like `modin.pandas.DataFrame.sample` to work as well as `pandas.DataFrame.sample`. 

Maybe for now we can find a way to replace functions like `modin.pandas.DataFrame.sample` with the pandas equivalents. I don't know a good way to do this, though. @modin-project/modin-core @modin-project/modin-contributors is that reasonable?

Note that this problem isn't limited to `DataFrameGroupBy.Apply`. For example, we can apply `pandas.Series.sum` on each column of a pandas dataframe, but we can't apply the modin Series sum on each column of a Modin dataframe. The following script works at pandas  1.4.3 but fails at Modin 05933a5f27fb96f5a7ff6025ae2573d033a31b11 if I replace `pandas as pd` with `modin.pandas as pd`:

```python
import pandas as pd
df = pd.DataFrame([1])
print(df.apply(pd.Series.sum))
```
Hmm - that's a great point @mvashishtha - perhaps we can do some arg pre processing that checks if an argument is callable, and if it is + it's from the `modin.pandas` module, we try and swap it out with the equivalent from the pandas module?

For context, this issue results from how Modin internally handles its partitions. Partitions are (usually) backed by pandas dataframes (where one big Modin DataFrame is usually chunked into multiple smaller pandas DataFrames) - when we call `modin.pandas.DataFrame.sum` on our Modin DataFrame, everything is fine, since it's type is `modin.pandas.DataFrame`, and the API knows how to handle it.

If a user tried to do something like `df.apply(pd.sum)`, where `pd` is `modin.pandas`, that `pd.sum` is applied to our partitions, which is basically the same as doing the following:
```python
getattr(modin.pandas.DataFrame, "sum")(pandas.DataFrame())
```
Which doesn't work, since Modin is, by design, **not interoperable with pandas**. 

Ultimately though, I think Mahesh is right that users **should** be able to do things like this and just have it work - we could either add the processing I suggested at the API layer (which may result in a lot of code, since the API layer is quite broad) or at the dataframe (or maybe even the partition layer (which should hopefully involve less work). 
@mvashishtha @RehanSD yes my first intuition was to have any Modin methods being applied default to pandas instead. I'm not sure if there's a better solution though. 
So I wouldn't say it's "defaulting to pandas" in this case - the functions need to be from pandas in order to work, but will still be parallel since they'll be applied to partitions. In fact, this is how Modin implements many functions eg count or sum - we map the pandas function across the partitions. 
Ah right, yeah we're not exactly defaulting to pandas, we're just applying the pandas function specifically across all the partitions.

Commit message: move code into a separate function

Signed-off-by: Anatoly Myachev <anatoly.myachev@intel.com>

Patch to review:
diff --git a/modin/pandas/dataframe.py b/modin/pandas/dataframe.py
index 76e2327a577..462781b484f 100644
--- a/modin/pandas/dataframe.py
+++ b/modin/pandas/dataframe.py
@@ -57,6 +57,7 @@
     from_pandas,
     from_non_pandas,
     broadcast_item,
+    cast_function_modin2pandas,
     SET_DATAFRAME_ATTRIBUTE_WARNING,
 )
 from .iterator import PartitionIterator
@@ -383,6 +384,7 @@ def apply(
         """
         Apply a function along an axis of the ``DataFrame``.
         """
+        func = cast_function_modin2pandas(func)
         axis = self._get_axis_number(axis)
         query_compiler = super(DataFrame, self).apply(
             func,
diff --git a/modin/pandas/groupby.py b/modin/pandas/groupby.py
index 4caff18cfdb..037a1e27e65 100644
--- a/modin/pandas/groupby.py
+++ b/modin/pandas/groupby.py
@@ -36,6 +36,7 @@
     wrap_into_list,
     MODIN_UNNAMED_SERIES_LABEL,
 )
+from modin.pandas.utils import cast_function_modin2pandas
 from modin.core.storage_formats.base.query_compiler import BaseQueryCompiler
 from modin.core.dataframe.algebra.default2pandas.groupby import GroupBy
 from modin.config import IsExperimental
@@ -433,6 +434,7 @@ def cummax(self, axis=0, **kwargs):
         )
 
     def apply(self, func, *args, **kwargs):
+        func = cast_function_modin2pandas(func)
         if not isinstance(func, BuiltinFunctionType):
             func = wrap_udf_function(func)
 
diff --git a/modin/pandas/resample.py b/modin/pandas/resample.py
index bbbfb283f49..b69af48270f 100644
--- a/modin/pandas/resample.py
+++ b/modin/pandas/resample.py
@@ -21,6 +21,7 @@
 from typing import Optional
 from modin.logging import ClassLogger
 from modin.utils import _inherit_docstrings
+from modin.pandas.utils import cast_function_modin2pandas
 
 
 @_inherit_docstrings(pandas.core.resample.Resampler)
@@ -142,6 +143,7 @@ def get_group(self, name, obj=None):
         return group if self.axis == 0 else group.T
 
     def apply(self, func, *args, **kwargs):
+        func = cast_function_modin2pandas(func)
         from .dataframe import DataFrame
 
         if isinstance(self._dataframe, DataFrame):
diff --git a/modin/pandas/series.py b/modin/pandas/series.py
index d7af8c49ba5..1121a75fe13 100644
--- a/modin/pandas/series.py
+++ b/modin/pandas/series.py
@@ -37,7 +37,7 @@
 from modin.config import IsExperimental, PersistentPickle
 from .base import BasePandasDataset, _ATTRS_NO_LOOKUP
 from .iterator import PartitionIterator
-from .utils import from_pandas, is_scalar, _doc_binary_op
+from .utils import from_pandas, is_scalar, _doc_binary_op, cast_function_modin2pandas
 from .accessor import CachedAccessor, SparseAccessor
 
 
@@ -637,6 +637,7 @@ def apply(
         """
         Invoke function on values of Series.
         """
+        func = cast_function_modin2pandas(func)
         self._validate_function(func)
         # apply and aggregate have slightly different behaviors, so we have to use
         # each one separately to determine the correct return type. In the case of
diff --git a/modin/pandas/test/dataframe/test_udf.py b/modin/pandas/test/dataframe/test_udf.py
index 32a61e2b161..2bafb004ab4 100644
--- a/modin/pandas/test/dataframe/test_udf.py
+++ b/modin/pandas/test/dataframe/test_udf.py
@@ -245,6 +245,22 @@ def test_apply_udf(data, func):
     )
 
 
+def test_apply_modin_func_4635():
+    data = [1]
+    modin_df, pandas_df = create_test_dfs(data)
+    df_equals(modin_df.apply(pd.Series.sum), pandas_df.apply(pandas.Series.sum))
+
+    data = data = {"a": [1, 2, 3], "b": [1, 2, 3], "c": [1, 2, 3]}
+    modin_df, pandas_df = create_test_dfs(data)
+    modin_df = modin_df.set_index(["a"])
+    pandas_df = pandas_df.set_index(["a"])
+
+    df_equals(
+        modin_df.groupby("a", group_keys=False).apply(pd.DataFrame.sample, n=1),
+        pandas_df.groupby("a", group_keys=False).apply(pandas.DataFrame.sample, n=1),
+    )
+
+
 def test_eval_df_use_case():
     frame_data = {"a": random_state.randn(10), "b": random_state.randn(10)}
     df = pandas.DataFrame(frame_data)
diff --git a/modin/pandas/utils.py b/modin/pandas/utils.py
index 6875359a695..abef69318be 100644
--- a/modin/pandas/utils.py
+++ b/modin/pandas/utils.py
@@ -138,6 +138,29 @@ def from_dataframe(df):
     return DataFrame(query_compiler=FactoryDispatcher.from_dataframe(df))
 
 
+def cast_function_modin2pandas(func):
+    """
+    Replace Modin functions with pandas functions if `func` is callable.
+
+    Parameters
+    ----------
+    func : object
+
+    Returns
+    -------
+    object
+    """
+    if callable(func):
+        if func.__module__ == "modin.pandas.series":
+            func = getattr(pandas.Series, func.__name__)
+        elif func.__module__ in ("modin.pandas.dataframe", "modin.pandas.base"):
+            # FIXME: when the method is defined in `modin.pandas.base` file, then the
+            # type cannot be determined, in general there may be an error, but at the
+            # moment it is better.
+            func = getattr(pandas.DataFrame, func.__name__)
+    return func
+
+
 def is_scalar(obj):
     """
     Return True if given object is scalar.
diff --git a/modin/pandas/window.py b/modin/pandas/window.py
index 2ad74c0b302..c6aed4d3668 100644
--- a/modin/pandas/window.py
+++ b/modin/pandas/window.py
@@ -19,6 +19,7 @@
 
 from modin.logging import ClassLogger
 from modin.utils import _inherit_docstrings
+from modin.pandas.utils import cast_function_modin2pandas
 
 
 @_inherit_docstrings(pandas.core.window.rolling.Window)
@@ -224,6 +225,7 @@ def apply(
         args=None,
         kwargs=None,
     ):
+        func = cast_function_modin2pandas(func)
         return self._dataframe.__constructor__(
             query_compiler=self._query_compiler.rolling_apply(
                 self.axis,


Code review comments:
diff_hunk: @@ -245,6 +245,22 @@ def test_apply_udf(data, func):
     )
 
 
+def test_apply_modin_func_4635():
+    data = [1]
+    modin_df, pandas_df = create_test_dfs(data)
+    df_equals(modin_df.apply(pd.Series.sum), pandas_df.apply(pandas.Series.sum))
+
+    data = data = {"a": [1, 2, 3], "b": [1, 2, 3], "c": [1, 2, 3]}
path: modin/pandas/test/dataframe/test_udf.py
text: @author:
```suggestion
    data = {"a": [1, 2, 3], "b": [1, 2, 3], "c": [1, 2, 3]}
```