from dataclasses import replace from unittest.mock import patch import pyarrow as pa import pytest from ray import ObjectRef from ray.data._internal.execution.interfaces import BlockSlice, RefBundle from ray.data._internal.execution.interfaces.ref_bundle import BlockEntry from ray.data.block import BlockMetadata _TEST_SCHEMA = pa.schema([("col", pa.int64())]) def test_get_preferred_locations(): first_block_ref = ObjectRef(b"1" * 28) second_block_ref = ObjectRef(b"2" * 28) third_block_ref = ObjectRef(b"3" * 28) meta = BlockMetadata(num_rows=None, size_bytes=1, exec_stats=None, input_files=None) bundle = RefBundle( blocks=[ BlockEntry(first_block_ref, meta), BlockEntry(second_block_ref, meta), BlockEntry(third_block_ref, meta), ], owns_blocks=True, schema=None, ) def _get_obj_locs(obj_refs): assert obj_refs == [first_block_ref, second_block_ref, third_block_ref] return { first_block_ref: { "object_size": 1024, "did_spill": False, "node_ids": ["1", "2", "3"], }, second_block_ref: { "object_size": 2048, "did_spill": False, "node_ids": ["2", "3"], }, third_block_ref: { "object_size": 4096, "did_spill": False, "node_ids": ["2"], }, } with patch("ray.experimental.get_local_object_locations", _get_obj_locs): preferred_object_locs = bundle.get_preferred_object_locations() assert { "1": 1024, # first_block_ref] "2": 7168, # first_block_ref, second_block_ref, third_block_ref "3": 3072, # first_block_ref, second_block_ref } == preferred_object_locs def test_ref_bundle_num_rows_size_bytes(): block_ref_one = ObjectRef(b"1" * 28) block_ref_two = ObjectRef(b"2" * 28) meta_one = BlockMetadata( num_rows=10, size_bytes=100, exec_stats=None, input_files=None ) meta_two = BlockMetadata( num_rows=5, size_bytes=50, exec_stats=None, input_files=None ) # Before slice bundle = RefBundle( blocks=[ BlockEntry(block_ref_one, meta_one), BlockEntry(block_ref_two, meta_two), ], owns_blocks=True, schema=None, ) assert bundle.num_rows() == 15 assert bundle.size_bytes() == 150 # After slice sliced = replace( bundle, slices=( BlockSlice(start_offset=2, end_offset=6), # 4 rows BlockSlice(start_offset=0, end_offset=2), # 2 rows ), ) assert sliced.num_rows() == 6 assert sliced.size_bytes() == 60 @pytest.mark.parametrize( "start_offset, end_offset", [ (-1, 0), # Negative start_offset (0, 11), # end_offset > num_rows (1, 0), # start_offset > end_offset ], ) def test_ref_bundle_with_invalid_slices(start_offset, end_offset): block_ref = ObjectRef(b"1" * 28) metadata = BlockMetadata( num_rows=10, size_bytes=100, exec_stats=None, input_files=None ) with pytest.raises(AssertionError): RefBundle( blocks=[BlockEntry(block_ref, metadata)], owns_blocks=True, schema=None, slices=[ BlockSlice(start_offset=start_offset, end_offset=end_offset), ], ) def test_slice_ref_bundle_basic(): block_ref_one = ObjectRef(b"1" * 28) block_ref_two = ObjectRef(b"2" * 28) meta_one = BlockMetadata( num_rows=6, size_bytes=60, exec_stats=None, input_files=None ) meta_two = BlockMetadata( num_rows=4, size_bytes=40, exec_stats=None, input_files=None ) bundle = RefBundle( blocks=[ BlockEntry(block_ref_one, meta_one), BlockEntry(block_ref_two, meta_two), ], owns_blocks=True, schema=_TEST_SCHEMA, ) consumed, remaining = bundle.slice(8) assert consumed.num_rows() == 8 assert remaining.num_rows() == 2 assert consumed.slices == ( BlockSlice(start_offset=0, end_offset=6), BlockSlice(start_offset=0, end_offset=2), ) assert remaining.slices == (BlockSlice(start_offset=2, end_offset=4),) def test_slice_ref_bundle_should_raise_error_if_needed_rows_is_not_less_than_num_rows(): block_ref = ObjectRef(b"1" * 28) metadata = BlockMetadata( num_rows=5, size_bytes=50, exec_stats=None, input_files=None ) bundle = RefBundle( blocks=[BlockEntry(block_ref, metadata)], owns_blocks=True, schema=None, ) with pytest.raises(AssertionError): bundle.slice(5) def test_slice_ref_bundle_with_existing_slices(): block_ref_one = ObjectRef(b"1" * 28) block_ref_two = ObjectRef(b"2" * 28) meta_one = BlockMetadata( num_rows=10, size_bytes=100, exec_stats=None, input_files=None ) meta_two = BlockMetadata( num_rows=5, size_bytes=50, exec_stats=None, input_files=None ) bundle = RefBundle( blocks=[ BlockEntry(block_ref_one, meta_one), BlockEntry(block_ref_two, meta_two), ], owns_blocks=True, schema=_TEST_SCHEMA, slices=[ BlockSlice(start_offset=2, end_offset=10), BlockSlice(start_offset=0, end_offset=3), ], ) consumed, remaining = bundle.slice(7) assert consumed.num_rows() == 7 assert consumed.slices == (BlockSlice(start_offset=2, end_offset=9),) assert consumed.size_bytes() == 70 assert remaining.num_rows() == 4 assert remaining.slices == ( BlockSlice(start_offset=9, end_offset=10), BlockSlice(start_offset=0, end_offset=3), ) assert remaining.size_bytes() == 40 @pytest.mark.parametrize( "num_rows,slice_rows", [ (5, 0), # Zero rows requested (5, 5), # Equal to total (must be less than) (5, 6), # More than available ], ) def test_slice_ref_bundle_invalid_rows(num_rows, slice_rows): """Test that slicing with invalid row counts raises appropriate errors.""" block_ref = ObjectRef(b"1" * 28) metadata = BlockMetadata( num_rows=num_rows, size_bytes=num_rows * 10, exec_stats=None, input_files=None ) bundle = RefBundle( blocks=[BlockEntry(block_ref, metadata)], owns_blocks=True, schema=None, ) with pytest.raises(AssertionError): bundle.slice(slice_rows) def test_ref_bundle_with_none_slices(): """Test that None can be used to represent full blocks in slices.""" block_ref_one = ObjectRef(b"1" * 28) block_ref_two = ObjectRef(b"2" * 28) meta_one = BlockMetadata( num_rows=10, size_bytes=100, exec_stats=None, input_files=None ) meta_two = BlockMetadata( num_rows=5, size_bytes=50, exec_stats=None, input_files=None ) # Test with all None slices (representing full blocks) bundle = RefBundle( blocks=[ BlockEntry(block_ref_one, meta_one), BlockEntry(block_ref_two, meta_two), ], owns_blocks=True, schema=None, slices=[None, None], ) assert bundle.num_rows() == 15 # 10 + 5 assert bundle.size_bytes() == 150 # 100 + 50 def test_ref_bundle_with_mixed_none_and_explicit_slices(): """Test mixing None and explicit BlockSlice objects.""" block_ref_one = ObjectRef(b"1" * 28) block_ref_two = ObjectRef(b"2" * 28) block_ref_three = ObjectRef(b"3" * 28) meta_one = BlockMetadata( num_rows=10, size_bytes=100, exec_stats=None, input_files=None ) meta_two = BlockMetadata( num_rows=8, size_bytes=80, exec_stats=None, input_files=None ) meta_three = BlockMetadata( num_rows=6, size_bytes=60, exec_stats=None, input_files=None ) # Mix None (full block) with explicit slices bundle = RefBundle( blocks=[ BlockEntry(block_ref_one, meta_one), BlockEntry(block_ref_two, meta_two), BlockEntry(block_ref_three, meta_three), ], owns_blocks=True, schema=None, slices=[ None, # Full block: 10 rows, 100 bytes BlockSlice(start_offset=2, end_offset=6), # 4 rows, ~40 bytes None, # Full block: 6 rows, 60 bytes ], ) assert bundle.num_rows() == 20 # 10 + 4 + 6 assert bundle.size_bytes() == 200 # 100 + 40 + 60 def test_slice_ref_bundle_with_none_slices(): """Test slicing a bundle that has None slices.""" block_ref_one = ObjectRef(b"1" * 28) block_ref_two = ObjectRef(b"2" * 28) meta_one = BlockMetadata( num_rows=6, size_bytes=60, exec_stats=None, input_files=None ) meta_two = BlockMetadata( num_rows=4, size_bytes=40, exec_stats=None, input_files=None ) # Start with None slices (full blocks) bundle = RefBundle( blocks=[ BlockEntry(block_ref_one, meta_one), BlockEntry(block_ref_two, meta_two), ], owns_blocks=True, schema=_TEST_SCHEMA, slices=[None, None], ) # Slice it to get first 8 rows consumed, remaining = bundle.slice(8) assert consumed.num_rows() == 8 assert remaining.num_rows() == 2 # The None slices should be converted to explicit BlockSlice objects assert consumed.slices == ( BlockSlice(start_offset=0, end_offset=6), BlockSlice(start_offset=0, end_offset=2), ) assert remaining.slices == (BlockSlice(start_offset=2, end_offset=4),) def test_ref_bundle_str(): """Test the __str__ method returns a readable representation.""" block_ref_one = ObjectRef(b"1" * 28) block_ref_two = ObjectRef(b"2" * 28) block_ref_three = ObjectRef(b"3" * 28) meta_one = BlockMetadata( num_rows=10, size_bytes=100, exec_stats=None, input_files=None ) meta_two = BlockMetadata( num_rows=5, size_bytes=50, exec_stats=None, input_files=None ) meta_three = BlockMetadata( num_rows=3, size_bytes=30, exec_stats=None, input_files=None ) slice_three = BlockSlice(start_offset=0, end_offset=3) bundle = RefBundle( blocks=[ BlockEntry(block_ref_one, meta_one), BlockEntry(block_ref_two, meta_two), BlockEntry(block_ref_three, meta_three), ], owns_blocks=True, schema=_TEST_SCHEMA, slices=[None, None, slice_three], ) expected = f"""RefBundle(3 blocks, 18 rows, schema={_TEST_SCHEMA}, owns_blocks=True, blocks=( 0: 10 rows, 100 bytes, slice=None (full block) 1: 5 rows, 50 bytes, slice=None (full block) 2: 3 rows, 30 bytes, slice=BlockSlice(start_offset=0, end_offset=3) ) )""" assert str(bundle) == expected def test_merge_ref_bundles(): block_ref_one = ObjectRef(b"1" * 28) block_ref_two = ObjectRef(b"2" * 28) metadata_one = BlockMetadata( num_rows=10, size_bytes=100, exec_stats=None, input_files=None ) metadata_two = BlockMetadata( num_rows=10, size_bytes=10, exec_stats=None, input_files=None ) bundle_one = RefBundle( blocks=[ BlockEntry(block_ref_one, metadata_one), BlockEntry(block_ref_one, metadata_one), ], owns_blocks=True, schema=_TEST_SCHEMA, slices=[ BlockSlice(start_offset=0, end_offset=1), BlockSlice(start_offset=1, end_offset=2), ], ) bundle_two = RefBundle( blocks=[ BlockEntry(block_ref_two, metadata_two), BlockEntry(block_ref_two, metadata_two), ], owns_blocks=False, schema=_TEST_SCHEMA, slices=[ BlockSlice(start_offset=2, end_offset=3), BlockSlice(start_offset=3, end_offset=4), ], ) merged = RefBundle.merge_ref_bundles([bundle_one, bundle_two]) assert merged.schema == _TEST_SCHEMA # The merged bundle should own the blocks if all input bundles own their blocks. # Since bundle_two doesn't own its blocks, the merged bundle should not own its # blocks. assert merged.owns_blocks is False assert len(merged.blocks) == 4 assert merged.slices == ( BlockSlice(start_offset=0, end_offset=1), BlockSlice(start_offset=1, end_offset=2), BlockSlice(start_offset=2, end_offset=3), BlockSlice(start_offset=3, end_offset=4), ) def test_ref_bundle_eq_and_hash(): """Tests that `__eq__` and `__hash__` use field-based comparison, so that copies (e.g. from dataclasses.replace) are equal to the original and usable as dict keys.""" ref_a = ObjectRef(b"1" * 28) ref_b = ObjectRef(b"2" * 28) meta = BlockMetadata(num_rows=10, size_bytes=100, exec_stats=None, input_files=None) bundle = RefBundle( blocks=[BlockEntry(ref_a, meta)], owns_blocks=True, schema=_TEST_SCHEMA, slices=[BlockSlice(start_offset=0, end_offset=5)], output_split_idx=1, ) # Identity: same object is always equal assert bundle == bundle # A dataclasses.replace copy should be equal and have the same hash copy = replace(bundle) assert copy is not bundle assert copy == bundle assert hash(copy) == hash(bundle) # Equal bundles work as dict keys d = {bundle: "value"} assert d[copy] == "value" # Non-RefBundle comparison returns False (not TypeError) assert bundle != "not a bundle" assert bundle != 42 assert bundle != None # noqa: E711 # Differing any field should break equality and hash diff_bundle_1 = replace(bundle, owns_blocks=False) assert bundle != diff_bundle_1 and hash(bundle) != hash(diff_bundle_1) diff_bundle_2 = replace(bundle, output_split_idx=2) assert bundle != diff_bundle_2 and hash(bundle) != hash(diff_bundle_2) diff_bundle_3 = replace(bundle, schema=pa.schema([("other", pa.string())])) assert bundle != diff_bundle_3 and hash(bundle) != hash(diff_bundle_3) diff_bundle_4 = replace(bundle, slices=(None,)) assert bundle != diff_bundle_4 and hash(bundle) != hash(diff_bundle_4) # Differing block ref diff_ref_bundle = replace(bundle, blocks=[BlockEntry(ref_b, meta)]) assert bundle != diff_ref_bundle # Differing block metadata diff_meta_bundle = replace( bundle, blocks=[BlockEntry(ref_a, replace(meta, num_rows=11))] ) assert bundle != diff_meta_bundle class TestBlockEntryAndSchedulingHints: """BlockEntry construction, legacy 2-tuple normalization, hints accessor.""" def _meta(self, num_rows=1, size_bytes=10): return BlockMetadata( num_rows=num_rows, size_bytes=size_bytes, exec_stats=None, input_files=None ) def test_blockentry_direct_construction(self): ref = ObjectRef(b"1" * 28) meta = self._meta() bundle = RefBundle( blocks=[BlockEntry(ref=ref, metadata=meta)], owns_blocks=True, schema=None, ) assert bundle.blocks[0].ref is ref assert bundle.blocks[0].metadata is meta def test_rejects_legacy_two_tuple_construction(self): # Strict construction: RefBundle no longer accepts 2-tuples. Callers # must wrap each entry in BlockEntry explicitly. This makes the # data shape visible at every call site instead of relying on # post-init magic. ref = ObjectRef(b"1" * 28) meta = self._meta() with pytest.raises(AssertionError, match="BlockEntry"): RefBundle(blocks=[(ref, meta)], owns_blocks=True, schema=None) def test_accessors_return_parallel_lists(self): ref_a = ObjectRef(b"1" * 28) ref_b = ObjectRef(b"2" * 28) meta = self._meta() bundle = RefBundle( blocks=[ BlockEntry(ref=ref_a, metadata=meta), BlockEntry(ref=ref_b, metadata=meta), ], owns_blocks=True, schema=None, ) assert bundle.block_refs == [ref_a, ref_b] assert bundle.metadata == [meta, meta] def test_size_bytes_still_works_with_blockentry(self): # Regression guard: BlockEntry refactor must not break size_bytes # / num_rows aggregation. bundle = RefBundle( blocks=[ BlockEntry(ObjectRef(b"1" * 28), self._meta(num_rows=3, size_bytes=30)), BlockEntry(ObjectRef(b"2" * 28), self._meta(num_rows=5, size_bytes=50)), ], owns_blocks=False, schema=None, ) assert bundle.num_rows() == 8 assert bundle.size_bytes() == 80 if __name__ == "__main__": import sys import pytest sys.exit(pytest.main(["-v", __file__]))