chore: import upstream snapshot with attribution

python/ray/data/tests/test_bundle_queue.py

@@ -0,0 +1,149 @@
+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__]))