[Data] Fix args and kwargs passed to ActorPool map_batches

python/ray/data/_internal/planner/plan_udf_map_op.py

@@ -54,10 +54,10 @@ def _plan_udf_map_op(
 
         fn_ = make_callable_class_concurrent(op._fn)
 
-        def fn(item: Any) -> Any:
+        def fn(item: Any, *args, **kwargs) -> Any:
             assert ray.data._cached_fn is not None
             assert ray.data._cached_cls == fn_
-            return ray.data._cached_fn(item)
+            return ray.data._cached_fn(item, *args, **kwargs)
 
         def init_fn():
             if ray.data._cached_fn is None:

python/ray/data/tests/test_map.py

@@ -167,6 +167,26 @@ def __call__(self, x):
     actor_reuse = ds.filter(StatefulFn, compute=ray.data.ActorPoolStrategy()).take()
     assert len(actor_reuse) == 9, actor_reuse
 
+    class StatefulFnWithArgs:
+        def __init__(self, arg, kwarg):
+            assert arg == 1
+            assert kwarg == 2
+
+        def __call__(self, x, arg, kwarg):
+            assert arg == 1
+            assert kwarg == 2
+            return x
+
+    # map_batches with args & kwargs
+    ds.map_batches(
+        StatefulFnWithArgs,
+        compute=ray.data.ActorPoolStrategy(),
+        fn_args=(1,),
+        fn_kwargs={"kwarg": 2},
+        fn_constructor_args=(1,),
+        fn_constructor_kwargs={"kwarg": 2},
+    ).take() == list(range(10))
+
 
 def test_concurrent_callable_classes(shutdown_only):
     """Test that concurrenct actor pool runs user UDF in a separate thread."""