import time import pyarrow as pa import pytest import ray from ray.data._internal.arrow_block import ArrowBlockAccessor from ray.data._internal.arrow_ops.transform_pyarrow import try_combine_chunked_columns from ray.data._internal.batcher import ( SHUFFLE_BUFFER_COMPACTION_THRESHOLD, Batcher, ShufflingBatcher, ) from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder from ray.data.block import BlockAccessor def gen_block(num_rows): return pa.table({"foo": [1] * num_rows}) 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, ) total_added = 0 total_yielded = 0 def add_and_check(num_rows, expect_has_batch): """Add a block and verify has_batch() matches expectation.""" nonlocal total_added batcher.add(gen_block(num_rows)) total_added += num_rows assert batcher.has_batch() == expect_has_batch, ( f"after adding {num_rows}: has_batch()={batcher.has_batch()}, " f"expected {expect_has_batch} " f"(compacted={batcher._num_compacted_rows()}, " f"uncompacted={batcher._num_uncompacted_rows()}, " f"total={batcher._num_rows()})" ) def next_and_check( expect_full_batch=True, expect_has_batch_after=True, ): """Consume one batch and verify size and post-state.""" nonlocal total_yielded batch = batcher.next_batch() total_yielded += len(batch) if expect_full_batch: assert ( len(batch) == batch_size ), f"expected full batch of {batch_size}, got {len(batch)}" else: assert len(batch) <= batch_size assert batcher.has_batch() == expect_has_batch_after, ( f"after next_batch: has_batch()={batcher.has_batch()}, " f"expected {expect_has_batch_after} " f"(compacted={batcher._num_compacted_rows()}, " f"uncompacted={batcher._num_uncompacted_rows()}, " f"total={batcher._num_rows()})" ) # Before any data is added, there should be no batches. assert not batcher.has_batch() assert not batcher.has_any() # Add blocks incrementally. All rows go into the pending buffer, # and has_batch will return False until enough rows accumulate. add_and_check(3, expect_has_batch=False) # total=3 add_and_check(7, expect_has_batch=False) # total=10 add_and_check(10, expect_has_batch=False) # total=20 # After adding 15 more (total=35), total - batch_size = 30 >= min_rows_to_trigger. add_and_check(15, expect_has_batch=True) # total=35 # All 35 rows are still uncompacted since no next_batch() has been called. assert batcher._shuffle_buffer is None assert batcher._builder.num_rows() == 35 # Consume one batch — this triggers the first compaction. next_and_check(expect_full_batch=True, expect_has_batch_after=True) assert batcher._shuffle_buffer is not None # compaction happened assert batcher._builder.num_rows() == 0 # all rows moved to compacted buffer # Add more data while consuming. add_and_check(20, expect_has_batch=True) # total grows # Consume batches. Each must be full since we're still streaming. while batcher.has_batch(): batch = batcher.next_batch() assert len(batch) == batch_size total_yielded += batch_size # Streaming exhausted: remaining rows <= batch_size (not enough to trigger # has_batch without more data or done_adding). assert batcher._num_rows() <= batch_size # Add a partial amount and signal done. batcher.add(gen_block(8)) total_added += 8 batcher.done_adding() # Drain remaining full batches via next_and_check. while batcher.has_batch(): remaining_after = batcher._num_rows() - batch_size next_and_check( expect_full_batch=True, expect_has_batch_after=remaining_after >= batch_size, ) # Final partial batch. if batcher.has_any(): next_and_check(expect_full_batch=False, expect_has_batch_after=False) # All rows must be accounted for. assert total_yielded == total_added assert not batcher.has_any() def test_batching_pyarrow_table_with_many_chunks(): """Make sure batching a pyarrow table with many chunks is fast. See https://github.com/ray-project/ray/issues/31108 for more details. """ num_chunks = 5000 batch_size = 1024 batches = [] for _ in range(num_chunks): batch = {} for i in range(10): batch[str(i)] = list(range(batch_size)) batches.append(pa.Table.from_pydict(batch)) block = pa.concat_tables(batches, promote=True) start = time.perf_counter() batcher = Batcher(batch_size, ensure_copy=False) batcher.add(block) batcher.done_adding() while batcher.has_any(): batcher.next_batch() duration = time.perf_counter() - start assert duration < 10 start = time.perf_counter() shuffling_batcher = ShufflingBatcher( batch_size=batch_size, shuffle_buffer_min_size=batch_size ) shuffling_batcher.add(block) shuffling_batcher.done_adding() while shuffling_batcher.has_any(): shuffling_batcher.next_batch() duration = time.perf_counter() - start assert duration < 30 @pytest.mark.parametrize( "batch_size,local_shuffle_buffer_size", [(1, 1), (10, 1), (1, 10), (10, 1000), (1000, 10)], ) def test_shuffling_batcher_grid(batch_size, local_shuffle_buffer_size): ds = ray.data.range_tensor(10000, shape=(130,)) start = time.time() count = 0 for batch in ds.iter_batches( batch_size=batch_size, local_shuffle_buffer_size=local_shuffle_buffer_size ): count += len(batch["data"]) print((ds.size_bytes() / 1e9) / (time.time() - start), "GB/s") assert count == 10000 @pytest.mark.parametrize( "batch_size,local_shuffle_buffer_size", [(1, 1), (10, 1), (1, 10), (10, 100), (100, 10)], ) def test_local_shuffle_determinism(batch_size, local_shuffle_buffer_size): # Preserve order so that the blocks are in the same order prior to shuffling. ctx = ray.data.DataContext.get_current() ctx.execution_options.preserve_order = True TEST_ITERATIONS = 10 ds = ray.data.range(1000) batch_map = {} for i in range(TEST_ITERATIONS): for batch in ds.iter_batches( batch_size=batch_size, local_shuffle_buffer_size=local_shuffle_buffer_size, local_shuffle_seed=0, ): if i == 0: batch_map[batch["id"][0]] = batch else: # Check that batch is the same as the first dataset's batch assert all(batch_map[batch["id"][0]]["id"] == batch["id"]) def test_local_shuffle_buffer_warns_if_too_large(shutdown_only): ray.shutdown() ray.init(object_store_memory=128 * 1024 * 1024) # Each row is 16 MiB * 8 = 128 MiB ds = ray.data.range_tensor(2, shape=(16, 1024, 1024)) # Test that Ray Data emits a warning if the local shuffle buffer size would cause # spilling. with pytest.warns(UserWarning, match="shuffle buffer"): # Each row is 128 MiB and the shuffle buffer size is 2 rows, so expect at least # 256 MiB of memory usage > 128 MiB total on node. batches = ds.iter_batches(batch_size=1, local_shuffle_buffer_size=2) next(iter(batches)) def _collect_rows_full_method(blocks, batch_size, buffer_size, seed): """Reference implementation using the old full-shuffle method. Materializes a fully shuffled copy of the buffer on each compaction, then yields contiguous slices. Used to validate the incremental index method. """ shuffle_buffer_min_size = max(buffer_size, batch_size) min_rows_to_yield_batch = max( 1, int(shuffle_buffer_min_size * SHUFFLE_BUFFER_COMPACTION_THRESHOLD) ) builder = DelegatingBlockBuilder() shuffle_buffer = None batch_head = 0 shuffle_seed = seed for block in blocks: if BlockAccessor.for_block(block).num_rows() > 0: builder.add_block(block) done_adding = True rows = [] while True: compacted = 0 if shuffle_buffer is not None: compacted = max( 0, BlockAccessor.for_block(shuffle_buffer).num_rows() - batch_head ) uncompacted = builder.num_rows() num_rows = compacted + uncompacted has_batch = num_rows >= batch_size has_any = num_rows > 0 if not (has_batch or (done_adding and has_any)): break # Compaction: merge uncompacted rows into shuffle buffer. if uncompacted > 0 and (done_adding or compacted <= min_rows_to_yield_batch): if shuffle_buffer is not None: if batch_head > 0: block_acc = BlockAccessor.for_block(shuffle_buffer) shuffle_buffer = block_acc.slice(batch_head, block_acc.num_rows()) builder.add_block(shuffle_buffer) shuffle_buffer = builder.build() shuffle_buffer = BlockAccessor.for_block(shuffle_buffer).random_shuffle( shuffle_seed ) if shuffle_seed is not None: shuffle_seed += 1 if isinstance(BlockAccessor.for_block(shuffle_buffer), ArrowBlockAccessor): shuffle_buffer = try_combine_chunked_columns(shuffle_buffer) builder = DelegatingBlockBuilder() batch_head = 0 buf_size = BlockAccessor.for_block(shuffle_buffer).num_rows() bs = min(batch_size, buf_size - batch_head) batch = BlockAccessor.for_block(shuffle_buffer).slice( batch_head, batch_head + bs ) batch_head += bs rows.extend(batch.column("val").to_pylist()) return rows @pytest.mark.parametrize( "batch_size,buffer_size,num_blocks,block_size", [ (5, 20, 10, 10), (1, 10, 5, 20), (10, 10, 3, 50), (7, 30, 8, 15), (100, 100, 2, 200), ], ) def test_incremental_index_matches_full_method( batch_size, buffer_size, num_blocks, block_size ): """Verify that the incremental index method yields the same multiset of rows as the old full-shuffle reference implementation.""" seed = 42 blocks = [ pa.table({"val": list(range(i * block_size, (i + 1) * block_size))}) for i in range(num_blocks) ] # Incremental index method (current implementation). batcher = ShufflingBatcher( batch_size=batch_size, shuffle_buffer_min_size=buffer_size, shuffle_seed=seed, ) for block in blocks: batcher.add(block) batcher.done_adding() rows_index = [] while batcher.has_batch() or batcher.has_any(): batch = batcher.next_batch() rows_index.extend(batch.column("val").to_pylist()) # Full-shuffle reference. rows_full = _collect_rows_full_method(blocks, batch_size, buffer_size, seed) total_rows = num_blocks * block_size assert len(rows_index) == total_rows assert len(rows_full) == total_rows assert sorted(rows_index) == sorted(rows_full) == list(range(total_rows)) def test_no_partial_batch_mid_stream(): """has_batch() must not return True when total rows < batch_size. With SHUFFLE_BUFFER_COMPACTION_THRESHOLD < 1.0, _min_rows_to_yield_batch can be less than batch_size. If we drain the compacted buffer below batch_size while no uncompacted rows are available, has_batch() must return False — otherwise next_batch() would return a partial batch mid-stream. """ batch_size = 10 buffer_size = 10 # common case: equal to batch_size batcher = ShufflingBatcher( batch_size=batch_size, shuffle_buffer_min_size=buffer_size, shuffle_seed=0, ) # Add enough rows to trigger compaction and yield some batches. batcher.add(gen_block(35)) # Consume batches until the compacted buffer is partially drained. batches = [] while batcher.has_batch(): batch = batcher.next_batch() batches.append(batch) # Every batch returned mid-stream must be full. assert ( len(batch) == batch_size ), f"got partial batch of {len(batch)} rows mid-stream" # At this point has_batch() is False. There may be leftover rows # (< batch_size) but they should not be yielded until done_adding. leftover = batcher._num_rows() assert leftover < batch_size # After done_adding, the remaining rows should drain as a partial batch. batcher.done_adding() assert batcher.has_any() final_batch = batcher.next_batch() assert len(final_batch) == leftover total = sum(len(b) for b in batches) + len(final_batch) assert total == 35 if __name__ == "__main__": import sys sys.exit(pytest.main(["-v", __file__]))