code: def test_batch_mapper_arrow_data_format(
      ray_start_regular_shared, ds, expected_df, expected_numpy_df
  ):
      """Tests batch mapper functionality for arrow data format.
  
      Note:
          For single column pandas dataframes, we automatically convert it to
          single column tensor with column name as `__value__`.
      """
  
      def add_and_modify_udf_pandas(df: "pd.DataFrame"):
          col_name = "column_1"
          if len(df.columns) == 1:
              col_name = list(df.columns)[0]
          df[col_name] = df[col_name] + 1
          if "column_2" in df:
              df["column_2"] *= 2
          return df
  
      def add_and_modify_udf_numpy(data: Dict[str, np.ndarray]):
          data["column_1"] = data["column_1"] + 1
          if "column_2" in data:
              data["column_2"] = data["column_2"] * 2
          return data
  
      # Test map_batches
      transformed_ds = ds.map_batches(
          add_and_modify_udf_pandas, batch_format="pandas"
      )
      out_df_map_batches = transformed_ds.to_pandas()
      assert_frame_equal(out_df_map_batches, expected_df)
  
      transformed_ds = ds.map_batches(
          add_and_modify_udf_numpy, batch_format="numpy"
      )
      out_df_map_batches = transformed_ds.to_pandas()
      assert_frame_equal(out_df_map_batches, expected_numpy_df)
  
      # Test BatchMapper
      batch_mapper = BatchMapper(
          fn=add_and_modify_udf_pandas, batch_format="pandas"
      )
      batch_mapper.fit(ds)
      transformed_ds = batch_mapper.transform(ds)
      out_df = transformed_ds.to_pandas()
      assert_frame_equal(out_df, expected_df)
  
      batch_mapper = BatchMapper(
          fn=add_and_modify_udf_numpy, batch_format="numpy"
      )
      batch_mapper.fit(ds)
      transformed_ds = batch_mapper.transform(ds)
      out_df = transformed_ds.to_pandas()
      assert_frame_equal(out_df, expected_numpy_df)
  
file_path: ray/python/ray/data/tests/preprocessors/test_batch_mapper.py
imports: from typing import Dict
  import numpy
  import pandas
  import pytest
  from pandas.testing import assert_frame_equal
  from pytest_lazyfixture import lazy_fixture
  import ray
  from ray.data.preprocessors import BatchMapper
  from ray.tests.conftest import *
label_desc: Use image node_id 4 for calling a global function with example usage: test_batch_mapper_arrow_data_format(ray_start_regular_shared, ds, expected_df, expected_numpy_df) and returns: df, data