code: def test_map_batches_batch_mutation(
      ray_start_regular_shared,
      num_rows,
      num_blocks,
      batch_size,
      restore_data_context,
  ):
      ctx = DataContext.get_current()
      ctx.execution_options.preserve_order = True
  
      # Test that batch mutation works without encountering a read-only error (e.g. if the
      # batch is a zero-copy view on data in the object store).
      def mutate(df):
          df["id"] += 1
          return df
  
      ds = ray.data.range(num_rows, parallelism=num_blocks).repartition(
          num_blocks
      )
      # Convert to Pandas blocks.
      ds = ds.map_batches(
          lambda df: df, batch_format="pandas", batch_size=None
      )
  
      # Apply UDF that mutates the batches.
      ds = ds.map_batches(mutate, batch_size=batch_size)
      assert [row["id"] for row in ds.iter_rows()] == list(
          range(1, num_rows + 1)
      )
  
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 14 for calling a global function with example usage: test_map_batches_batch_mutation(ray_start_regular_shared, num_rows, num_blocks, batch_size, restore_data_context) and returns: df