code: def generate_sort_fn(
      sort_key: SortKey,
      _debug_limit_shuffle_execution_to_num_blocks: Optional[
          int
      ] = None,
  ) -> AllToAllTransformFn:
      """Generate function to sort blocks by the specified key column or key function."""
  
      def fn(
          sort_key: SortKey,
          refs: List[RefBundle],
          ctx: TaskContext,
      ) -> Tuple[List[RefBundle], StatsDict]:
          blocks = []
          metadata = []
          for ref_bundle in refs:
              for block, block_metadata in ref_bundle.blocks:
                  blocks.append(block)
                  metadata.append(block_metadata)
          if len(blocks) == 0:
              return (blocks, {})
          sort_key.validate_schema(
              unify_block_metadata_schema(metadata)
          )
  
          num_mappers = len(blocks)
          # Use same number of output partitions.
          num_outputs = num_mappers
  
          # Sample boundaries for sort key.
          if not sort_key.boundaries:
              boundaries = SortTaskSpec.sample_boundaries(
                  blocks, sort_key, num_outputs
              )
          else:
              boundaries = [(b,) for b in sort_key.boundaries]
              num_outputs = len(boundaries) + 1
          _, ascending = sort_key.to_pandas_sort_args()
          if not ascending:
              boundaries.reverse()
          sort_spec = SortTaskSpec(
              boundaries=boundaries, sort_key=sort_key
          )
  
          if DataContext.get_current().use_push_based_shuffle:
              scheduler = PushBasedShuffleTaskScheduler(sort_spec)
          else:
              scheduler = PullBasedShuffleTaskScheduler(sort_spec)
  
          return scheduler.execute(
              refs,
              num_outputs,
              ctx,
              _debug_limit_execution_to_num_blocks=(
                  _debug_limit_shuffle_execution_to_num_blocks
              ),
          )
  
      # NOTE: use partial function to pass parameters to avoid error like
      # "UnboundLocalError: local variable ... referenced before assignment",
      # because `key` and `descending` variables are reassigned in `fn()`.
      return partial(fn, sort_key)
  
file_path: ray/python/ray/data/_internal/planner/sort.py
imports: from functools import partial
  from typing import List, Optional, Tuple
  from ray.data._internal.execution.interfaces import AllToAllTransformFn, RefBundle, TaskContext
  from ray.data._internal.planner.exchange.pull_based_shuffle_task_scheduler import PullBasedShuffleTaskScheduler
  from ray.data._internal.planner.exchange.push_based_shuffle_task_scheduler import PushBasedShuffleTaskScheduler
  from ray.data._internal.planner.exchange.sort_task_spec import SortTaskSpec
  from ray.data._internal.sort import SortKey
  from ray.data._internal.stats import StatsDict
  from ray.data._internal.util import unify_block_metadata_schema
  from ray.data.context import DataContext
label_desc: Use image node_id 1 for calling a global function with example usage: generate_sort_fn(sort_key, _debug_limit_shuffle_execution_to_num_blocks) and returns: partial, scheduler, blocks,  dict