code: def test_pandas_container(
      batch_dim_exc: t.Type[Exception], wrong_batch_dim: int
  ):
      cols = ["a", "b", "c"]
      arr1 = np.ones((3, 3))
      df1 = pd.DataFrame(arr1, columns=cols)
      arr2 = np.arange(6, dtype=np.float64).reshape(2, 3)
      df2 = pd.DataFrame(arr2, columns=cols)
      batches = [df1, df2]
      batch, indices = c.PandasDataFrameContainer.batches_to_batch(
          batches
      )
      assert batch.equals(pd.concat(batches, ignore_index=True))
  
      (
          restored_df1,
          restored_df2,
      ) = c.PandasDataFrameContainer.batch_to_batches(batch, indices)
      assert df1.equals(restored_df1)
      assert df2.equals(restored_df2)
  
      assert c.PandasDataFrameContainer.from_payload(
          c.PandasDataFrameContainer.to_payload(df1, batch_dim=0)
      ).equals(df1)
  
      (
          restored_batch,
          restored_indices,
      ) = c.PandasDataFrameContainer.from_batch_payloads(
          c.PandasDataFrameContainer.batch_to_payloads(batch, indices)
      )
      assert restored_indices == indices
      assert restored_batch.equals(batch)
  
      # PandasDataFrameContainer should only allow batch_dim = 0
  
      with pytest.raises(batch_dim_exc):
          c.PandasDataFrameContainer.batches_to_batch(
              batches, batch_dim=wrong_batch_dim
          )
  
      with pytest.raises(batch_dim_exc):
          c.PandasDataFrameContainer.batch_to_batches(
              batch, indices, batch_dim=wrong_batch_dim
          )
  
file_path: BentoML/tests/unit/_internal/runner/test_container.py
imports: from __future__ import annotations
  import typing
  import numpy
  import pandas
  import pytest
  import bentoml._internal.runner.container
label_desc: Use image node_id 3 for calling a global function with example usage: test_pandas_container(batch_dim_exc, wrong_batch_dim) without return types