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

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.


Functions to edit:
modin/core/dataframe/algebra/binary.py:Binary.register
modin/core/dataframe/pandas/dataframe/dataframe.py:PandasDataframe.map
modin/core/dataframe/pandas/partitioning/partition_manager.py:PandasDataframePartitionManager.map_partitions
modin/core/dataframe/pandas/partitioning/partition_manager.py:PandasDataframePartitionManager.lazy_map_partitions
modin/core/execution/ray/common/engine_wrapper.py:RayWrapper
modin/core/execution/ray/common/engine_wrapper.py:RayWrapper.put

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