code: def test_sort_partition_same_key_to_same_block(
      ray_start_regular, use_push_based_shuffle
  ):
      num_items = 100
      xs = [1] * num_items
      ds = ray.data.from_items(xs)
      sorted_ds = ds.repartition(num_items).sort("item")
  
      # We still have 100 blocks
      assert len(sorted_ds._block_num_rows()) == num_items
      # Only one of them is non-empty
      count = sum(1 for x in sorted_ds._block_num_rows() if x > 0)
      assert count == 1
      # That non-empty block contains all rows
      total = sum(x for x in sorted_ds._block_num_rows() if x > 0)
      assert total == num_items
  
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 3 for calling a global function with example usage: test_sort_partition_same_key_to_same_block(ray_start_regular, use_push_based_shuffle) without return types