Repository: modin-project/modin
Issue category: Feature Request

Issue:
FEAT: Do not put binary functions to the Ray storage multiple times.
Currently, the binary operations are wrapped into lambdas which are put into the Ray storage on each operation.


Hints:

Patch:
diff --git a/modin/core/dataframe/algebra/binary.py b/modin/core/dataframe/algebra/binary.py
index f19040cc104..af0c6ee7e8e 100644
--- a/modin/core/dataframe/algebra/binary.py
+++ b/modin/core/dataframe/algebra/binary.py
@@ -415,7 +415,9 @@ def caller(
                     ):
                         shape_hint = "column"
                     new_modin_frame = query_compiler._modin_frame.map(
-                        lambda df: func(df, other, *args, **kwargs),
+                        func,
+                        func_args=(other, *args),
+                        func_kwargs=kwargs,
                         dtypes=dtypes,
                     )
                 return query_compiler.__constructor__(
diff --git a/modin/core/dataframe/pandas/dataframe/dataframe.py b/modin/core/dataframe/pandas/dataframe/dataframe.py
index 7da33f52284..029de373f75 100644
--- a/modin/core/dataframe/pandas/dataframe/dataframe.py
+++ b/modin/core/dataframe/pandas/dataframe/dataframe.py
@@ -2093,6 +2093,8 @@ def map(
         func: Callable,
         dtypes: Optional[str] = None,
         new_columns: Optional[pandas.Index] = None,
+        func_args=None,
+        func_kwargs=None,
     ) -> "PandasDataframe":
         """
         Perform a function that maps across the entire dataset.
@@ -2108,13 +2110,19 @@ def map(
         new_columns : pandas.Index, optional
             New column labels of the result, its length has to be identical
             to the older columns. If not specified, old column labels are preserved.
+        func_args : iterable, optional
+            Positional arguments for the 'func' callable.
+        func_kwargs : dict, optional
+            Keyword arguments for the 'func' callable.
 
         Returns
         -------
         PandasDataframe
             A new dataframe.
         """
-        new_partitions = self._partition_mgr_cls.map_partitions(self._partitions, func)
+        new_partitions = self._partition_mgr_cls.map_partitions(
+            self._partitions, func, func_args, func_kwargs
+        )
         if new_columns is not None and self.has_materialized_columns:
             assert len(new_columns) == len(
                 self.columns
diff --git a/modin/core/dataframe/pandas/partitioning/partition_manager.py b/modin/core/dataframe/pandas/partitioning/partition_manager.py
index 3a1dd63e555..0e9d35cf545 100644
--- a/modin/core/dataframe/pandas/partitioning/partition_manager.py
+++ b/modin/core/dataframe/pandas/partitioning/partition_manager.py
@@ -566,7 +566,13 @@ def broadcast_axis_partitions(
 
     @classmethod
     @wait_computations_if_benchmark_mode
-    def map_partitions(cls, partitions, map_func):
+    def map_partitions(
+        cls,
+        partitions,
+        map_func,
+        func_args=None,
+        func_kwargs=None,
+    ):
         """
         Apply `map_func` to every partition in `partitions`.
 
@@ -576,6 +582,10 @@ def map_partitions(cls, partitions, map_func):
             Partitions housing the data of Modin Frame.
         map_func : callable
             Function to apply.
+        func_args : iterable, optional
+            Positional arguments for the 'map_func'.
+        func_kwargs : dict, optional
+            Keyword arguments for the 'map_func'.
 
         Returns
         -------
@@ -585,14 +595,23 @@ def map_partitions(cls, partitions, map_func):
         preprocessed_map_func = cls.preprocess_func(map_func)
         return np.array(
             [
-                [part.apply(preprocessed_map_func) for part in row_of_parts]
+                [
+                    part.apply(
+                        preprocessed_map_func,
+                        *func_args if func_args is not None else (),
+                        **func_kwargs if func_kwargs is not None else {},
+                    )
+                    for part in row_of_parts
+                ]
                 for row_of_parts in partitions
             ]
         )
 
     @classmethod
     @wait_computations_if_benchmark_mode
-    def lazy_map_partitions(cls, partitions, map_func, func_args=None):
+    def lazy_map_partitions(
+        cls, partitions, map_func, func_args=None, func_kwargs=None
+    ):
         """
         Apply `map_func` to every partition in `partitions` *lazily*.
 
@@ -604,6 +623,8 @@ def lazy_map_partitions(cls, partitions, map_func, func_args=None):
             Function to apply.
         func_args : iterable, optional
             Positional arguments for the 'map_func'.
+        func_kwargs : dict, optional
+            Keyword arguments for the 'map_func'.
 
         Returns
         -------
@@ -616,7 +637,8 @@ def lazy_map_partitions(cls, partitions, map_func, func_args=None):
                 [
                     part.add_to_apply_calls(
                         preprocessed_map_func,
-                        *(tuple() if func_args is None else func_args),
+                        *func_args if func_args is not None else (),
+                        **func_kwargs if func_kwargs is not None else {},
                     )
                     for part in row
                 ]
diff --git a/modin/core/execution/ray/common/engine_wrapper.py b/modin/core/execution/ray/common/engine_wrapper.py
index 8e20033d20d..e274d28c764 100644
--- a/modin/core/execution/ray/common/engine_wrapper.py
+++ b/modin/core/execution/ray/common/engine_wrapper.py
@@ -18,10 +18,14 @@
 """
 
 import asyncio
+import os
+from types import FunctionType
 
 import ray
 from ray.util.client.common import ClientObjectRef
 
+from modin.error_message import ErrorMessage
+
 
 @ray.remote
 def _deploy_ray_func(func, *args, **kwargs):  # pragma: no cover
@@ -48,6 +52,8 @@ def _deploy_ray_func(func, *args, **kwargs):  # pragma: no cover
 class RayWrapper:
     """Mixin that provides means of running functions remotely and getting local results."""
 
+    _func_cache = {}
+
     @classmethod
     def deploy(cls, func, f_args=None, f_kwargs=None, num_returns=1):
         """
@@ -127,6 +133,19 @@ def put(cls, data, **kwargs):
         ray.ObjectID
             Ray object identifier to get the value by.
         """
+        if isinstance(data, FunctionType):
+            qname = data.__qualname__
+            if "<locals>" not in qname and "<lambda>" not in qname:
+                ref = cls._func_cache.get(data, None)
+                if ref is None:
+                    if len(cls._func_cache) < 1024:
+                        ref = ray.put(data)
+                        cls._func_cache[data] = ref
+                    else:
+                        msg = "To many functions in the RayWrapper cache!"
+                        assert "MODIN_GITHUB_CI" not in os.environ, msg
+                        ErrorMessage.warn(msg)
+                return ref
         return ray.put(data, **kwargs)
 
     @classmethod


Test patch: