code: def _test_equal_split_balanced(block_sizes, num_splits):
      blocks = []
      metadata = []
      ref_bundles = []
      total_rows = 0
      for block_size in block_sizes:
          block = pd.DataFrame(
              {"id": list(range(total_rows, total_rows + block_size))}
          )
          blocks.append(ray.put(block))
          metadata.append(
              BlockAccessor.for_block(block).get_metadata(None, None)
          )
          blk = (blocks[-1], metadata[-1])
          ref_bundles.append(RefBundle((blk,), owns_blocks=True))
          total_rows += block_size
      block_list = BlockList(blocks, metadata, owned_by_consumer=True)
  
      logical_plan = LogicalPlan(InputData(input_data=ref_bundles))
      ds = Dataset(
          ExecutionPlan(
              block_list, DatasetStats.TODO(), run_by_consumer=True
          ),
          logical_plan,
      )
  
      splits = ds.split(num_splits, equal=True)
      split_counts = [split.count() for split in splits]
      assert len(split_counts) == num_splits
      expected_block_size = total_rows // num_splits
      # Check that all splits are the expected size.
      assert all(
          [count == expected_block_size for count in split_counts]
      )
      expected_total_rows = sum(split_counts)
      # Check that the expected number of rows were dropped.
      assert total_rows - expected_total_rows == total_rows % num_splits
      # Check that all rows are unique (content check).
      split_rows = [
          row for split in splits for row in split.take(total_rows)
      ]
      assert len(set(extract_values("id", split_rows))) == len(
          split_rows
      )
  
imports: import itertools
  import math
  import random
  import time
  from unittest.mock import patch
  import numpy
  import pandas
  import pytest
  import ray
  from ray.data._internal.block_list import BlockList
  from ray.data._internal.equalize import _equalize
  from ray.data._internal.execution.interfaces import RefBundle
  from ray.data._internal.logical.interfaces import LogicalPlan
  from ray.data._internal.logical.operators.input_data_operator import InputData
  from ray.data._internal.plan import ExecutionPlan
  from ray.data._internal.split import _drop_empty_block_split, _generate_global_split_results, _generate_per_block_split_indices, _generate_valid_indices, _split_at_indices, _split_single_block
  from ray.data._internal.stats import DatasetStats
  from ray.data.block import BlockAccessor, BlockMetadata
  from ray.data.dataset import Dataset
  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_equal_split_balanced(block_sizes, num_splits) without return types