wehub-resource-sync
150 lines
4.0 KiB
Python
150 lines
4.0 KiB
Python
from typing import Any
|
|
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data._internal.execution.bundle_queue import create_bundle_queue
|
|
from ray.data._internal.execution.interfaces import BlockEntry, RefBundle
|
|
from ray.data.block import BlockAccessor
|
|
|
|
|
|
def _create_bundle(data: Any) -> RefBundle:
|
|
"""Create a RefBundle with a single row with the given data."""
|
|
block = pa.Table.from_pydict({"data": [data]})
|
|
block_ref = ray.put(block)
|
|
metadata = BlockAccessor.for_block(block).get_metadata()
|
|
schema = BlockAccessor.for_block(block).schema()
|
|
return RefBundle(
|
|
[BlockEntry(block_ref, metadata)], owns_blocks=False, schema=schema
|
|
)
|
|
|
|
|
|
# CVGA-start
|
|
def test_add_and_length():
|
|
queue = create_bundle_queue()
|
|
queue.add(_create_bundle("test1"))
|
|
queue.add(_create_bundle("test2"))
|
|
assert len(queue) == 2
|
|
|
|
|
|
def test_get_next():
|
|
queue = create_bundle_queue()
|
|
bundle1 = _create_bundle("test1")
|
|
queue.add(bundle1)
|
|
bundle2 = _create_bundle("test11")
|
|
queue.add(bundle2)
|
|
|
|
popped_bundle = queue.get_next()
|
|
assert popped_bundle is bundle1
|
|
assert len(queue) == 1
|
|
assert queue.num_blocks() == 1
|
|
assert queue.num_rows() == 1
|
|
assert queue.estimate_size_bytes() == bundle2.size_bytes()
|
|
|
|
|
|
def test_peek_next():
|
|
queue = create_bundle_queue()
|
|
bundle1 = _create_bundle("test1")
|
|
queue.add(bundle1)
|
|
bundle2 = _create_bundle("test11")
|
|
queue.add(bundle2)
|
|
|
|
peeked_bundle = queue.peek_next()
|
|
assert peeked_bundle is bundle1
|
|
assert len(queue) == 2 # Length should remain unchanged
|
|
assert queue.num_blocks() == 2
|
|
assert queue.num_rows() == 2
|
|
assert queue.estimate_size_bytes() == bundle1.size_bytes() + bundle2.size_bytes()
|
|
|
|
|
|
def test_get_next_empty_queue():
|
|
queue = create_bundle_queue()
|
|
with pytest.raises(IndexError):
|
|
queue.get_next()
|
|
|
|
|
|
def test_get_next_does_not_leak_objects():
|
|
queue = create_bundle_queue()
|
|
bundle1 = _create_bundle("test11")
|
|
queue.add(bundle1)
|
|
queue.get_next()
|
|
assert len(queue) == 0
|
|
assert queue.estimate_size_bytes() == 0
|
|
assert queue.num_rows() == 0
|
|
assert queue.num_blocks() == 0
|
|
|
|
|
|
def test_peek_next_empty_queue():
|
|
queue = create_bundle_queue()
|
|
assert queue.peek_next() is None
|
|
assert len(queue) == 0
|
|
assert queue.num_blocks() == 0
|
|
assert queue.estimate_size_bytes() == 0
|
|
assert queue.num_rows() == 0
|
|
|
|
|
|
def test_remove():
|
|
queue = create_bundle_queue()
|
|
bundle1 = _create_bundle("test1")
|
|
bundle2 = _create_bundle("test11")
|
|
queue.add(bundle1)
|
|
queue.add(bundle2)
|
|
|
|
queue.remove(bundle1)
|
|
assert len(queue) == 1
|
|
assert queue.num_blocks() == 1
|
|
assert queue.peek_next() is bundle2
|
|
assert queue.estimate_size_bytes() == bundle2.size_bytes()
|
|
assert queue.num_rows() == bundle2.num_rows()
|
|
|
|
|
|
def test_remove_does_not_leak_objects():
|
|
queue = create_bundle_queue()
|
|
bundle1 = _create_bundle("test1")
|
|
queue.add(bundle1)
|
|
queue.remove(bundle1)
|
|
assert len(queue) == 0
|
|
assert queue.num_blocks() == 0
|
|
assert queue.estimate_size_bytes() == 0
|
|
assert queue.num_rows() == 0
|
|
|
|
|
|
def test_add_and_remove_duplicates():
|
|
queue = create_bundle_queue()
|
|
bundle1 = _create_bundle("test1")
|
|
bundle2 = _create_bundle("test11")
|
|
queue.add(bundle1)
|
|
queue.add(bundle2)
|
|
queue.add(bundle1)
|
|
|
|
assert len(queue) == 3
|
|
assert queue.num_rows() == 3
|
|
assert queue.num_blocks() == 3
|
|
queue.remove(bundle1)
|
|
assert len(queue) == 2
|
|
|
|
assert queue.estimate_size_bytes() == bundle1.size_bytes() + bundle2.size_bytes()
|
|
assert queue.num_rows() == 2
|
|
assert queue.num_blocks() == 2
|
|
assert queue.peek_next() is bundle2
|
|
|
|
|
|
def test_clear():
|
|
queue = create_bundle_queue()
|
|
queue.add(_create_bundle("test1"))
|
|
queue.add(_create_bundle("test11"))
|
|
queue.clear()
|
|
assert len(queue) == 0
|
|
assert queue.estimate_size_bytes() == 0
|
|
assert queue.num_blocks() == 0
|
|
assert queue.num_rows() == 0
|
|
|
|
|
|
# CVGA-end
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|