code: def test_sort_arrow(
      ray_start_regular,
      num_items,
      parallelism,
      use_push_based_shuffle,
      use_polars,
  ):
      ctx = ray.data.context.DataContext.get_current()
  
      try:
          original_use_polars = ctx.use_polars
          ctx.use_polars = use_polars
  
          a = list(reversed(range(num_items)))
          b = [f"{x:03}" for x in range(num_items)]
          shard = int(np.ceil(num_items / parallelism))
          offset = 0
          dfs = []
          while offset < num_items:
              dfs.append(
                  pd.DataFrame(
                      {
                          "a": a[offset : offset + shard],
                          "b": b[offset : offset + shard],
                      }
                  )
              )
              offset += shard
          if offset < num_items:
              dfs.append(
                  pd.DataFrame({"a": a[offset:], "b": b[offset:]})
              )
          ds = ray.data.from_pandas(dfs).map_batches(
              lambda t: t, batch_format="pyarrow", batch_size=None
          )
  
          def assert_sorted(sorted_ds, expected_rows):
              assert [
                  tuple(row.values()) for row in sorted_ds.iter_rows()
              ] == list(expected_rows)
  
          assert_sorted(ds.sort(key="a"), zip(reversed(a), reversed(b)))
          # Make sure we have rows in each block.
          assert (
              len(
                  [
                      n
                      for n in ds.sort(key="a")._block_num_rows()
                      if n > 0
                  ]
              )
              == parallelism
          )
          assert_sorted(ds.sort(key="b"), zip(a, b))
          assert_sorted(ds.sort(key="a", descending=True), zip(a, b))
      finally:
          ctx.use_polars = original_use_polars
  
imports: import random
  from collections import defaultdict
  import numpy
  import pandas
  import pyarrow
  import pytest
  import ray
  from ray.data import Dataset
  from ray.data._internal.planner.exchange.push_based_shuffle_task_scheduler import PushBasedShuffleTaskScheduler
  from ray.data._internal.sort import SortKey
  from ray.data.block import BlockAccessor
  from ray.data.context import DataContext
  from ray.data.tests.conftest import *
  from ray.data.tests.util import extract_values
  from ray.tests.conftest import *
label_desc: Use image node_id 4 for calling a global function with example usage: test_sort_arrow(ray_start_regular, num_items, parallelism, use_push_based_shuffle, use_polars) without return types