class_docstr: Push-based shuffle merges intermediate map outputs on the reducer nodes
  while other map tasks are executing. The merged outputs are merged again
  during a final reduce stage. This works as follows:
  
  1. Submit rounds of concurrent map and merge tasks until all map inputs
  have been processed. In each round, we execute:
  
     M map tasks
       Each produces N outputs. Each output contains P blocks.
     N merge tasks
       Takes 1 output from each of M map tasks.
       Each produces P outputs.
     Where M and N are chosen to maximize parallelism across CPUs. Note that
     this assumes that all CPUs in the cluster will be dedicated to the
     shuffle job.
  
     Map and merge tasks are pipelined so that we always merge the previous
     round of map outputs while executing the next round of map tasks.
  
  2. In the final reduce stage:
     R reduce tasks
       Takes 1 output from one of the merge tasks from every round.
  
  Notes:
      N * P = R = total number of output blocks
      M / N = merge factor - the ratio of map : merge tasks is to improve
        pipelined parallelism. For example, if map takes twice as long to
        execute as merge, then we should set this to 2. If pipeline bubbles
        appear and the merge tasks are much longer than the map tasks, then
        the merge factor should be decreased, and vice versa.
      See paper at https://arxiv.org/abs/2203.05072 for more details.
class_docstr_tok: ["Push-based","shuffle","merges","intermediate","map","outputs","on","the","reducer","nodes","while","other","map","tasks","are","executing",".","The","merged","outputs","are","merged","again","during","a","final","reduce","stage",".","This","works","as","follows",":","1",".","Submit","rounds","of","concurrent","map","and","merge","tasks","until","all","map","inputs","have","been","processed",".","In","each","round",",","we","execute",":","M","map","tasks","Each","produces","N","outputs",".","Each","output","contains","P","blocks",".","N","merge","tasks","Takes","1","output","from","each","of","M","map","tasks",".","Each","produces","P","outputs",".","Where","M","and","N","are","chosen","to","maximize","parallelism","across","CPUs",".","Note","that","this","assumes","that","all","CPUs","in","the","cluster","will","be","dedicated","to","the","shuffle","job",".","Map","and","merge","tasks","are","pipelined","so","that","we","always","merge","the","previous","round","of","map","outputs","while","executing","the","next","round","of","map","tasks",".","2",".","In","the","final","reduce","stage",":","R","reduce","tasks","Takes","1","output","from","one","of","the","merge","tasks","from","every","round",".","Notes",":","N","*","P","=","R","=","total","number","of","output","blocks","M","\/","N","=","merge","factor","-","the","ratio","of","map",":","merge","tasks","is","to","improve","pipelined","parallelism",".","For","example",",","if","map","takes","twice","as","long","to","execute","as","merge",",","then","we","should","set","this","to","2",".","If","pipeline","bubbles","appear","and","the","merge","tasks","are","much","longer","than","the","map","tasks",",","then","the","merge","factor","should","be","decreased",",","and","vice","versa",".","See","paper","at","https",":","\/\/arxiv.org\/abs\/2203.05072","for","more","details","."]
docstr: Returns list of [BlockMetadata, O1, O2, O3, ...output_num_blocks].
docstr_tok: ["Returns","list","of","[","BlockMetadata",",","O1",",","O2",",","O3",",","...","output_num_blocks","]","."]
code: def _merge(
      reduce_fn,
      *all_mapper_outputs: List[List[Block]],
      reduce_args: Optional[List[Any]] = None,
  ) -> List[Union[BlockMetadata, Block]]:
      """
      Returns list of [BlockMetadata, O1, O2, O3, ...output_num_blocks].
      """
      assert (
          len(
              {
                  len(mapper_outputs)
                  for mapper_outputs in all_mapper_outputs
              }
          )
          == 1
      ), "Received different number of map inputs"
      stats = BlockExecStats.builder()
      if not reduce_args:
          reduce_args = []
  
      num_rows = 0
      size_bytes = 0
      schema = None
      for i, mapper_outputs in enumerate(zip(*all_mapper_outputs)):
          block, meta = reduce_fn(
              *reduce_args, *mapper_outputs, partial_reduce=True
          )
          yield block
  
          block = BlockAccessor.for_block(block)
          num_rows += block.num_rows()
          size_bytes += block.size_bytes()
          schema = block.schema()
          del block
  
      yield BlockMetadata(
          num_rows=num_rows,
          size_bytes=size_bytes,
          schema=schema,
          input_files=None,
          exec_stats=stats.build(),
      )
  
imports: import math
  from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union
  import ray
  from ray.data._internal.dataset_logger import DatasetLogger
  from ray.data._internal.execution.interfaces import RefBundle, TaskContext
  from ray.data._internal.planner.exchange.interfaces import ExchangeTaskScheduler, ExchangeTaskSpec
  from ray.data._internal.progress_bar import ProgressBar
  from ray.data._internal.remote_fn import cached_remote_fn
  from ray.data._internal.stats import StatsDict
  from ray.data.block import Block, BlockAccessor, BlockExecStats, BlockMetadata
  from ray.data.context import DataContext
  from ray.types import ObjectRef
  from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
label_desc: Use image node_id 3 for calling the PushBasedShuffleTaskScheduler obj's underlying member method code with example usage: obj._merge(reduce_fn) without return types