code: def test_shuffling_batcher():
      batch_size = 5
      buffer_size = 20
  
      with pytest.raises(
          ValueError,
          match="Must specify a batch_size if using a local shuffle.",
      ):
          ShufflingBatcher(
              batch_size=None, shuffle_buffer_min_size=buffer_size
          )
  
      # Should not raise error.
      ShufflingBatcher(
          batch_size=batch_size, shuffle_buffer_min_size=batch_size - 1
      )
  
      batcher = ShufflingBatcher(
          batch_size=batch_size,
          shuffle_buffer_min_size=buffer_size,
      )
  
      def add_and_check(
          num_rows,
          materialized_buffer_size,
          pending_buffer_size,
          expect_has_batch=False,
          no_nexting_yet=True,
      ):
          block = gen_block(num_rows)
          batcher.add(block)
          if expect_has_batch:
              assert batcher.has_batch()
          else:
              assert not batcher.has_batch()
  
          if no_nexting_yet:
              # Check that no shuffle buffer has been materialized yet.
              assert batcher._shuffle_buffer is None
              assert batcher._batch_head == 0
  
          assert batcher._builder.num_rows() == pending_buffer_size
          assert (
              batcher._materialized_buffer_size()
              == materialized_buffer_size
          )
  
      def next_and_check(
          current_cursor,
          materialized_buffer_size,
          pending_buffer_size,
          should_batch_be_full=True,
          should_have_batch_after=True,
          new_data_added=False,
      ):
          if should_batch_be_full:
              assert batcher.has_batch()
          else:
              batcher.has_any()
          if new_data_added:
              # If new data was added, the builder should be non-empty.
              assert batcher._builder.num_rows() > 0
          batch = batcher.next_batch()
  
          if should_batch_be_full:
              assert len(batch) == batch_size
  
          assert batcher._builder.num_rows() == pending_buffer_size
          assert (
              batcher._materialized_buffer_size()
              == materialized_buffer_size
          )
  
          if should_have_batch_after:
              assert batcher.has_batch()
          else:
              assert not batcher.has_batch()
  
      # Add less than a batch.
      # Buffer not full and no batch slack.
      add_and_check(
          3, materialized_buffer_size=0, pending_buffer_size=3
      )
  
      # Add to more than a batch (total=10).
      # Buffer not full and no batch slack.
      add_and_check(
          7, materialized_buffer_size=0, pending_buffer_size=10
      )
  
      # Fill up to buffer (total=20).
      # Buffer is full but no batch slack.
      add_and_check(
          10, materialized_buffer_size=0, pending_buffer_size=20
      )
  
      # Fill past buffer and over 1.5 * min buffer size. A batch is now available, but
      # compaction still doesn't happen until a next().
      add_and_check(
          15,
          materialized_buffer_size=0,
          pending_buffer_size=35,
          expect_has_batch=True,
      )
  
      # Consume only available batch.
      next_and_check(
          0,
          materialized_buffer_size=30,
          pending_buffer_size=0,
          should_have_batch_after=True,
          new_data_added=True,
      )
  
      # Add 4 batches-worth to the already-full buffer.
      add_and_check(
          20,
          materialized_buffer_size=30,
          pending_buffer_size=20,
          expect_has_batch=True,
          no_nexting_yet=False,
      )
  
      # Consume 4 batches from the buffer.
      next_and_check(
          0,
          materialized_buffer_size=25,
          pending_buffer_size=20,
          new_data_added=True,
      )
      next_and_check(
          batch_size,
          materialized_buffer_size=20,
          pending_buffer_size=20,
      )
      # Triggers materialization of pending batches to avoid falling below min buf size.
      next_and_check(
          2 * batch_size,
          materialized_buffer_size=35,
          pending_buffer_size=0,
      )
      next_and_check(
          3 * batch_size,
          materialized_buffer_size=30,
          pending_buffer_size=0,
          should_have_batch_after=True,
      )
  
      # Add a full batch + a partial batch to the buffer.
      add_and_check(
          8,
          materialized_buffer_size=30,
          pending_buffer_size=8,
          expect_has_batch=True,
          no_nexting_yet=False,
      )
      next_and_check(
          0,
          materialized_buffer_size=25,
          pending_buffer_size=8,
          should_have_batch_after=True,
          new_data_added=True,
      )
  
      # Indicate to the batcher that we're done adding blocks.
      batcher.done_adding()
  
      # Consume 6 full batches and one partial batch, fully draining the buffer.
      next_and_check(
          batch_size, materialized_buffer_size=28, pending_buffer_size=0
      )
      next_and_check(
          2 * batch_size,
          materialized_buffer_size=23,
          pending_buffer_size=0,
      )
      next_and_check(
          3 * batch_size,
          materialized_buffer_size=18,
          pending_buffer_size=0,
      )
      next_and_check(
          4 * batch_size,
          materialized_buffer_size=13,
          pending_buffer_size=0,
      )
      next_and_check(
          5 * batch_size,
          materialized_buffer_size=8,
          pending_buffer_size=0,
      )
      next_and_check(
          6 * batch_size,
          materialized_buffer_size=3,
          pending_buffer_size=0,
          should_have_batch_after=False,
      )
      next_and_check(
          7 * batch_size,
          materialized_buffer_size=0,
          pending_buffer_size=0,
          should_batch_be_full=False,
          should_have_batch_after=False,
      )
  
file_path: ray/python/ray/data/tests/test_batcher.py
imports: import time
  import pyarrow
  import pytest
  import ray
  from ray.data._internal.batcher import Batcher, ShufflingBatcher
label_desc: Use image node_id 2 for calling a global function with example usage: test_shuffling_batcher() without return types