code: def _balance_blocks_with_index(
      kvs,
      partition_order_mappings: dict = None,
      data_manager: DataManager = None,
  ):
      block_id = None
      block_num = data_manager.block_num
      ret_blocks = [[] for _ in range(block_num)]
      block_size = 0
      for _, (blocks, t) in kvs:
          if block_id is None:
              block_id = partition_order_mappings[_]["start_block_id"]
  
          flat_blocks = [
              Block.transform_block_to_list(block) for block in blocks
          ]
          for i, v in enumerate(t):
              v = v.item()
              if not v:
                  continue
  
              block_size += 1
              for j in range(block_num):
                  ret_blocks[j].append(flat_blocks[j][i])
  
              if block_size == data_manager.block_row_size:
                  yield block_id, data_manager.convert_to_blocks(
                      ret_blocks
                  )
                  block_size = 0
                  block_id += 1
                  ret_blocks = [[] for _ in range(block_num)]
  
      if block_size:
          yield block_id, data_manager.convert_to_blocks(ret_blocks)
  
imports: import copy
  import functools
  from typing import List, Union
  import torch
  from sklearn.utils import resample
  from .._dataframe import DataFrame
  from ..manager.data_manager import DataManager
  from ..manager.block_manager import Block
  from ._compress_block import compress_blocks
  from ._indexer import get_partition_order_by_raw_table, get_partition_order_mappings_by_block_table
  from ._promote_types import promote_partial_block_types
  from ._set_item import set_item
  from fate.arch.tensor import DTensor
label_desc: Use image node_id 9 for calling a global function with example usage: _balance_blocks_with_index(kvs, partition_order_mappings, data_manager) without return types