code: def test_map_batches_basic(
      ray_start_regular_shared, tmp_path, restore_data_context
  ):
      ctx = DataContext.get_current()
      ctx.execution_options.preserve_order = True
  
      # Test input validation
      ds = ray.data.range(5)
      with pytest.raises(ValueError):
          ds.map_batches(
              column_udf("id", lambda x: x + 1),
              batch_format="pyarrow",
              batch_size=-1,
          ).take()
  
      # 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 pandas
      ds = ray.data.read_parquet(str(tmp_path))
      ds2 = ds.map_batches(
          lambda df: df + 1, batch_size=1, batch_format="pandas"
      )
      ds_list = ds2.take()
      values = [s["one"] for s in ds_list]
      assert values == [2, 3, 4]
      values = [s["two"] for s in ds_list]
      assert values == [3, 4, 5]
  
      # Test Pyarrow
      ds = ray.data.read_parquet(str(tmp_path))
      ds2 = ds.map_batches(
          lambda pa: pa, batch_size=1, batch_format="pyarrow"
      )
      ds_list = ds2.take()
      values = [s["one"] for s in ds_list]
      assert values == [1, 2, 3]
      values = [s["two"] for s in ds_list]
      assert values == [2, 3, 4]
  
      # Test batch
      size = 300
      ds = ray.data.range(size)
      ds2 = ds.map_batches(
          lambda df: df + 1, batch_size=17, batch_format="pandas"
      )
      ds_list = ds2.take_all()
      for i in range(size):
          # The pandas column is "value", and it originally has rows from 0~299.
          # After the map batch, it should have 1~300.
          row = ds_list[i]
          assert row["id"] == i + 1
      assert ds.count() == 300
  
      # Test the lambda returns different types than the batch_format
      # pandas => list block
      ds = ray.data.read_parquet(str(tmp_path))
      ds2 = ds.map_batches(
          lambda df: {"id": np.array([1])}, batch_size=1
      )
      ds_list = extract_values("id", ds2.take())
      assert ds_list == [1, 1, 1]
      assert ds.count() == 3
  
      # pyarrow => list block
      ds = ray.data.read_parquet(str(tmp_path))
      ds2 = ds.map_batches(
          lambda df: {"id": np.array([1])},
          batch_size=1,
          batch_format="pyarrow",
      )
      ds_list = extract_values("id", ds2.take())
      assert ds_list == [1, 1, 1]
      assert ds.count() == 3
  
      # Test the wrong return value raises an exception.
      ds = ray.data.read_parquet(str(tmp_path))
      with pytest.raises(ValueError):
          ds_list = ds.map_batches(
              lambda df: 1, batch_size=2, batch_format="pyarrow"
          ).take()
  
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 10 for calling a global function with example usage: test_map_batches_basic(ray_start_regular_shared, tmp_path, restore_data_context) without return types