wehub-resource-sync
563 lines
17 KiB
Python
563 lines
17 KiB
Python
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__]))
|