code: def test_map_batches_extra_args(shutdown_only, tmp_path):
      ray.shutdown()
      ray.init(num_cpus=3)
  
      def put(x):
          # We only support automatic deref in the legacy backend.
          if DataContext.get_current().new_execution_backend:
              return x
          else:
              return ray.put(x)
  
      # Test input validation
      ds = ray.data.range(5)
  
      class Foo:
          def __call__(self, df):
              return df
  
      with pytest.raises(ValueError):
          # fn_constructor_args and fn_constructor_kwargs only supported for actor
          # compute strategy.
          ds.map_batches(
              lambda x: x,
              fn_constructor_args=(1,),
              fn_constructor_kwargs={"a": 1},
          )
  
      with pytest.raises(ValueError):
          # fn_constructor_args and fn_constructor_kwargs only supported for callable
          # class UDFs.
          ds.map_batches(
              lambda x: x,
              fn_constructor_args=(1,),
              fn_constructor_kwargs={"a": 1},
          )
  
      # Set up.
      df = pd.DataFrame({"one": [1, 2, 3], "two": [2, 3, 4]})
      table = pa.Table.from_pandas(df)
      pq.write_table(table, os.path.join(tmp_path, "test1.parquet"))
  
      # Test extra UDF args.
      # Test positional.
      def udf(batch, a):
          assert a == 1
          return batch + a
  
      ds = ray.data.read_parquet(str(tmp_path))
      ds2 = ds.map_batches(
          udf,
          batch_size=1,
          batch_format="pandas",
          fn_args=(put(1),),
      )
      ds_list = ds2.take()
      values = sorted([s["one"] for s in ds_list])
      assert values == [2, 3, 4]
      values = sorted([s["two"] for s in ds_list])
      assert values == [3, 4, 5]
  
      # Test kwargs.
      def udf(batch, b=None):
          assert b == 2
          return b * batch
  
      ds = ray.data.read_parquet(str(tmp_path))
      ds2 = ds.map_batches(
          udf,
          batch_size=1,
          batch_format="pandas",
          fn_kwargs={"b": put(2)},
      )
      ds_list = ds2.take()
      values = sorted([s["one"] for s in ds_list])
      assert values == [2, 4, 6]
      values = sorted([s["two"] for s in ds_list])
      assert values == [4, 6, 8]
  
      # Test both.
      def udf(batch, a, b=None):
          assert a == 1
          assert b == 2
          return b * batch + a
  
      ds = ray.data.read_parquet(str(tmp_path))
      ds2 = ds.map_batches(
          udf,
          batch_size=1,
          batch_format="pandas",
          fn_args=(put(1),),
          fn_kwargs={"b": put(2)},
      )
      ds_list = ds2.take()
      values = sorted([s["one"] for s in ds_list])
      assert values == [3, 5, 7]
      values = sorted([s["two"] for s in ds_list])
      assert values == [5, 7, 9]
  
      # Test constructor UDF args.
      # Test positional.
      class CallableFn:
          def __init__(self, a):
              assert a == 1
              self.a = a
  
          def __call__(self, x):
              return x + self.a
  
      ds = ray.data.read_parquet(str(tmp_path))
      ds2 = ds.map_batches(
          CallableFn,
          concurrency=1,
          batch_size=1,
          batch_format="pandas",
          fn_constructor_args=(put(1),),
      )
      ds_list = ds2.take()
      values = sorted([s["one"] for s in ds_list])
      assert values == [2, 3, 4]
      values = sorted([s["two"] for s in ds_list])
      assert values == [3, 4, 5]
  
      # Test kwarg.
      class CallableFn:
          def __init__(self, b=None):
              assert b == 2
              self.b = b
  
          def __call__(self, x):
              return self.b * x
  
      ds = ray.data.read_parquet(str(tmp_path))
      ds2 = ds.map_batches(
          CallableFn,
          concurrency=1,
          batch_size=1,
          batch_format="pandas",
          fn_constructor_kwargs={"b": put(2)},
      )
      ds_list = ds2.take()
      values = sorted([s["one"] for s in ds_list])
      assert values == [2, 4, 6]
      values = sorted([s["two"] for s in ds_list])
      assert values == [4, 6, 8]
  
      # Test both.
      class CallableFn:
          def __init__(self, a, b=None):
              assert a == 1
              assert b == 2
              self.a = a
              self.b = b
  
          def __call__(self, x):
              return self.b * x + self.a
  
      ds = ray.data.read_parquet(str(tmp_path))
      ds2 = ds.map_batches(
          CallableFn,
          concurrency=1,
          batch_size=1,
          batch_format="pandas",
          fn_constructor_args=(put(1),),
          fn_constructor_kwargs={"b": put(2)},
      )
      ds_list = ds2.take()
      values = sorted([s["one"] for s in ds_list])
      assert values == [3, 5, 7]
      values = sorted([s["two"] for s in ds_list])
      assert values == [5, 7, 9]
  
      # Test callable chain.
      ds = ray.data.read_parquet(str(tmp_path))
      fn_constructor_args = (put(1),)
      fn_constructor_kwargs = {"b": put(2)}
      ds2 = ds.map_batches(
          CallableFn,
          concurrency=1,
          batch_size=1,
          batch_format="pandas",
          fn_constructor_args=fn_constructor_args,
          fn_constructor_kwargs=fn_constructor_kwargs,
      ).map_batches(
          CallableFn,
          concurrency=1,
          batch_size=1,
          batch_format="pandas",
          fn_constructor_args=fn_constructor_args,
          fn_constructor_kwargs=fn_constructor_kwargs,
      )
      ds_list = ds2.take()
      values = sorted([s["one"] for s in ds_list])
      assert values == [7, 11, 15]
      values = sorted([s["two"] for s in ds_list])
      assert values == [11, 15, 19]
  
      # Test function + callable chain.
      ds = ray.data.read_parquet(str(tmp_path))
      fn_constructor_args = (put(1),)
      fn_constructor_kwargs = {"b": put(2)}
      ds2 = ds.map_batches(
          lambda df, a, b=None: b * df + a,
          batch_size=1,
          batch_format="pandas",
          fn_args=(put(1),),
          fn_kwargs={"b": put(2)},
      ).map_batches(
          CallableFn,
          concurrency=1,
          batch_size=1,
          batch_format="pandas",
          fn_constructor_args=fn_constructor_args,
          fn_constructor_kwargs=fn_constructor_kwargs,
      )
      ds_list = ds2.take()
      values = sorted([s["one"] for s in ds_list])
      assert values == [7, 11, 15]
      values = sorted([s["two"] for s in ds_list])
      assert values == [11, 15, 19]
  
imports: import itertools
  import math
  import os
  import signal
  import threading
  import time
  from typing import Iterator
  import numpy
  import pandas
  import pyarrow
  import pyarrow.parquet
  import pytest
  import ray
  from ray._private.test_utils import wait_for_condition
  from ray.data.block import BlockAccessor
  from ray.data.context import DataContext
  from ray.data.tests.conftest import *
  from ray.data.tests.util import column_udf, column_udf_class, extract_values
  from ray.tests.conftest import *
label_desc: Use image node_id 11 for calling a global function with example usage: test_map_batches_extra_args(shutdown_only, tmp_path) and returns: unknown, unknown, unknown, x, ray, df, unknown, unknown, unknown