[data][1/n] Base + HashLinkedBundleQueue by iamjustinhsu · Pull Request #60017 · ray-project/ray

iamjustinhsu · 2026-01-09T22:34:25Z

Description

Old PR: #59093

Related issues

Additional information

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

gemini-code-assist

Code Review

This pull request introduces a new base class for bundle queues, BaseBundleQueue, and a HashLinkedQueue implementation that provides efficient removal and containment checks. This is a solid architectural improvement for better extensibility. The OutputSplitter is updated to leverage this new queue, which simplifies some of its logic. I've identified a critical issue where a method was not updated to use the new queue's API, which could lead to runtime errors. Additionally, I have some suggestions to improve the base class design and code clarity.

python/ray/data/_internal/execution/operators/output_splitter.py

python/ray/data/_internal/execution/bundle_queue/base.py

python/ray/data/_internal/execution/bundle_queue/hash_link.py

python/ray/data/tests/test_bundle_queue.py

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

iamjustinhsu · 2026-01-09T22:52:51Z

/gemini review

gemini-code-assist

Code Review

This pull request introduces a new base class BaseBundleQueue for bundle queues and a new implementation HashLinkedQueue. HashLinkedQueue is a doubly linked list with a hash map, which provides efficient O(1) for add, get_next, contains, and remove operations. This new queue is used to replace list in OutputSplitter for better performance, especially for bundle removal based on locality. The old FIFOBundleQueue is replaced by HashLinkedQueue in tests. The changes are well-structured and improve performance. I have a few minor suggestions for improving docstrings and fixing a bug with an @override decorator.

python/ray/data/_internal/execution/bundle_queue/hash_link.py

python/ray/data/_internal/execution/bundle_queue/base.py

python/ray/data/_internal/execution/operators/output_splitter.py

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

iamjustinhsu · 2026-01-09T23:09:14Z

/gemini review

gemini-code-assist

Code Review

This pull request introduces a significant and well-structured refactoring of the bundle queue implementation. It defines a new base class BaseBundleQueue and a HashLinkedQueue that provides efficient removal of elements. The changes are mostly clean and improve the modularity of the code. However, I've found a critical issue in the implementation of the remove method in HashLinkedQueue which can lead to incorrect behavior in OutputSplitter when dealing with bundles that have the same value but are distinct instances. This needs to be addressed to ensure correctness, especially for locality-aware scheduling.

gemini-code-assist · 2026-01-09T23:12:05Z

python/ray/data/_internal/execution/bundle_queue/hash_link.py

+    def remove(self, bundle: RefBundle) -> RefBundle:
+        # Case 1: The queue is empty.
+        if bundle not in self._bundle_to_nodes:
+            raise ValueError(f"The bundle {bundle} is not in the queue.")
+
+        node = self._bundle_to_nodes[bundle].popleft()
+        if not self._bundle_to_nodes[bundle]:
+            del self._bundle_to_nodes[bundle]
+
+        node = self._remove_node(node)
+        return node.value


This remove implementation is ambiguous and can lead to incorrect behavior. It removes the first-added bundle that is equal to the given bundle by using popleft() on the deque of nodes. However, a caller like OutputSplitter._pop_bundle_to_dispatch iterates through the queue to find a specific bundle instance that matches a condition (e.g., locality), and then calls remove(bundle) to remove it. If there are multiple bundle instances in the queue that are equal, the iterator might find one instance, but remove(bundle) might remove a different (older) instance.

Example of bug:

Two equal but distinct bundles, bundleA_1 and bundleA_2, are added to the queue. The queue state is [nodeA_1, nodeA_2]. The map is {bundleA: deque([nodeA_1, nodeA_2])}.

The iterator in OutputSplitter finds nodeA_2's bundle, which matches a locality requirement.

remove(bundleA) is called.

This method does popleft() on the deque, which gets nodeA_1.

nodeA_1 is removed from the queue, but nodeA_2 should have been removed.

This can lead to incorrect bundle dispatching based on locality.

To fix this, the remove operation needs to be unambiguous. Here are some possible solutions:

Change __iter__ to yield a handle (like the _Node object) along with the bundle. The caller can then pass this handle to a new remove_node method to remove the specific instance. This would require changes in OutputSplitter.

Add a method like find_and_remove(predicate) to HashLinkedQueue that finds the first bundle matching a predicate and removes it atomically.

Two equal but distinct bundles, bundleA_1 and bundleA_2

This is not possible. if two bundles are equal, they cannot be distinct. The hash of a refbundle is id(self), so if their ids are not identical, then according to python dict implementation, even if the ids hash collide, it's still a different ref bundle

class Test: def __hash__(self): return 42 a = Test() b = Test() c = {} c[a] = 1 c[b] = 2 print(c)

This creates a dictionary with length 2

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

…/unify-internal-queue-abstraction-add-base

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

python/ray/data/_internal/execution/operators/output_splitter.py

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

python/ray/data/_internal/execution/bundle_queue/base.py

alexeykudinkin · 2026-01-13T01:05:27Z

python/ray/data/_internal/execution/bundle_queue/base.py

+        """Return the total # of rows across all bundles."""
+        return self._num_rows
+
+    def _reset_metrics(self):


Hold on, why do we need this?

We need this when we call queue.clear(), because we incrementally keep track of queued and dequeued bundles. I could just this to an absolute calculation, but would want to wait until all PRs get merged to do that. The other alternative is to just instantiate a new instance. @alexeykudinkin thoughts?

on second thought, I realize this is used for these methods: clear_internal_*() so those methods can be moved to the base class.

alexeykudinkin · 2026-01-13T01:07:40Z

python/ray/data/_internal/execution/bundle_queue/base.py

+        )
+
+
+class BaseBundleQueue(_QueueMetricRecorder):


Is every subclass gonna inherit from Base? If so, why not just merging mixin into it?

oh this was a remnant of an old PR where I was creating more base queue methods. I just split up the PRs and didn't change any behavior. I will move it up into the base class

alexeykudinkin · 2026-01-13T01:08:40Z

python/ray/data/_internal/execution/bundle_queue/hash_link.py

+    @override
+    def finalize(self, key: Optional[int] = None):
+        pass


Make it optional in the base

alexeykudinkin · 2026-01-13T01:11:08Z

python/ray/data/_internal/execution/operators/union_operator.py

        # Only used when `self._prserve_order` is True.
-        self._input_buffers: List[BundleQueue] = [
-            FIFOBundleQueue() for _ in range(len(input_ops))
+        self._input_buffers: List["BaseBundleQueue"] = [


Suggested change

self._input_buffers: List["BaseBundleQueue"] = [

self._input_buffers: List[HashLinkedQueue] = [

alexeykudinkin · 2026-01-13T01:11:37Z

python/ray/data/tests/test_output_splitter.py

 from ray.tests.conftest import *  # noqa


-@pytest.mark.parametrize("equal", [False, True])


Why changing these?

goutamvenkat-anyscale · 2026-01-09T23:28:48Z

python/ray/data/_internal/execution/bundle_queue/hash_link.py

+        self._tail: Optional[_Node] = None
+        # We use a dictionary to keep track of the nodes corresponding to each bundle.
+        # This allows us to remove a bundle from the queue in O(1) time. We need a list
+        # because a bundle can be added to the queue multiple times. Nodes in each list


When would a bundle be added multiple times? If concurrency > 1 or because of retries?

I'm actually not so sure when this can happen, I can take a look later

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

alexeykudinkin · 2026-01-14T21:12:43Z

python/ray/data/_internal/execution/bundle_queue/hash_link.py

+        # Remove will dequeue the bundle, therefore, we override
+        # get_next(), _get_next_inner()


Comment is confusing, this method should be overriding _get_next_inner, right?

lol yea that comment makes no sense. What i meant to say is that remove() and get_next() both use remove_node() which will dequeue bundle for you. I just moved the dequeuing logic up the stack to make it clearer, so now we override get_next_inner()

python/ray/data/_internal/execution/bundle_queue/hash_link.py

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

alexeykudinkin · 2026-01-16T02:53:18Z

python/ray/data/_internal/execution/bundle_queue/base.py

+        """Remove all bundles from the queue."""
+        ...
+
+    def finalize(self, **kwargs: Any):


Please add return type

## Description Old PR: ray-project#59093 ## Related issues ## Additional information --------- Signed-off-by: iamjustinhsu <jhsu@anyscale.com> Signed-off-by: Limark Dcunha <limarkdcunha@gmail.com>

## Description Old PR: ray-project#59093 ## Related issues ## Additional information --------- Signed-off-by: iamjustinhsu <jhsu@anyscale.com> Signed-off-by: jeffery4011 <jefferyshen1015@gmail.com>

## Description After #60017 got merged, I forgot to update the `test_bundle_queue` test suite. This PR adds more tests for `num_blocks`, `num_rows`, `estimate_size_bytes`, and `len(queue)` ## Related issues ## Additional information --------- Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

…eue` (#60538) #60017 and #60228 refactored the `FIFOBundleQueue` interface and renamed `FIFOBundleQueue.popleft` with `FIFOBundleQueue.get_next`. However, this name change wasn't reflected in the `UnionOperator` implementation, and as a result the operator can error when it clears its output queue. This change also fixes the flaky `test_union.py`. Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu>

## Description After ray-project#60017 got merged, I forgot to update the `test_bundle_queue` test suite. This PR adds more tests for `num_blocks`, `num_rows`, `estimate_size_bytes`, and `len(queue)` ## Related issues ## Additional information --------- Signed-off-by: iamjustinhsu <jhsu@anyscale.com> Signed-off-by: jinbum-kim <jinbum9958@gmail.com>

…eue` (ray-project#60538) ray-project#60017 and ray-project#60228 refactored the `FIFOBundleQueue` interface and renamed `FIFOBundleQueue.popleft` with `FIFOBundleQueue.get_next`. However, this name change wasn't reflected in the `UnionOperator` implementation, and as a result the operator can error when it clears its output queue. This change also fixes the flaky `test_union.py`. Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu> Signed-off-by: jinbum-kim <jinbum9958@gmail.com>

…eue` (ray-project#60538) ray-project#60017 and ray-project#60228 refactored the `FIFOBundleQueue` interface and renamed `FIFOBundleQueue.popleft` with `FIFOBundleQueue.get_next`. However, this name change wasn't reflected in the `UnionOperator` implementation, and as a result the operator can error when it clears its output queue. This change also fixes the flaky `test_union.py`. Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu>

## Description After ray-project#60017 got merged, I forgot to update the `test_bundle_queue` test suite. This PR adds more tests for `num_blocks`, `num_rows`, `estimate_size_bytes`, and `len(queue)` ## Related issues ## Additional information --------- Signed-off-by: iamjustinhsu <jhsu@anyscale.com> Signed-off-by: 400Ping <jiekaichang@apache.org>

…eue` (ray-project#60538) ray-project#60017 and ray-project#60228 refactored the `FIFOBundleQueue` interface and renamed `FIFOBundleQueue.popleft` with `FIFOBundleQueue.get_next`. However, this name change wasn't reflected in the `UnionOperator` implementation, and as a result the operator can error when it clears its output queue. This change also fixes the flaky `test_union.py`. Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu> Signed-off-by: 400Ping <jiekaichang@apache.org>

## Description Old PR: ray-project#59093 ## Related issues ## Additional information --------- Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

## Description After ray-project#60017 got merged, I forgot to update the `test_bundle_queue` test suite. This PR adds more tests for `num_blocks`, `num_rows`, `estimate_size_bytes`, and `len(queue)` ## Related issues ## Additional information --------- Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

…eue` (ray-project#60538) ray-project#60017 and ray-project#60228 refactored the `FIFOBundleQueue` interface and renamed `FIFOBundleQueue.popleft` with `FIFOBundleQueue.get_next`. However, this name change wasn't reflected in the `UnionOperator` implementation, and as a result the operator can error when it clears its output queue. This change also fixes the flaky `test_union.py`. Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu> Signed-off-by: Adel Nour <ans9868@nyu.edu>

## Description Old PR: ray-project#59093 ## Related issues ## Additional information --------- Signed-off-by: iamjustinhsu <jhsu@anyscale.com> Signed-off-by: peterxcli <peterxcli@gmail.com>

## Description After ray-project#60017 got merged, I forgot to update the `test_bundle_queue` test suite. This PR adds more tests for `num_blocks`, `num_rows`, `estimate_size_bytes`, and `len(queue)` ## Related issues ## Additional information --------- Signed-off-by: iamjustinhsu <jhsu@anyscale.com> Signed-off-by: peterxcli <peterxcli@gmail.com>

…eue` (ray-project#60538) ray-project#60017 and ray-project#60228 refactored the `FIFOBundleQueue` interface and renamed `FIFOBundleQueue.popleft` with `FIFOBundleQueue.get_next`. However, this name change wasn't reflected in the `UnionOperator` implementation, and as a result the operator can error when it clears its output queue. This change also fixes the flaky `test_union.py`. Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu> Signed-off-by: peterxcli <peterxcli@gmail.com>

## Description Old PR: ray-project#59093 ## Related issues ## Additional information --------- Signed-off-by: iamjustinhsu <jhsu@anyscale.com> Signed-off-by: peterxcli <peterxcli@gmail.com>

## Description After ray-project#60017 got merged, I forgot to update the `test_bundle_queue` test suite. This PR adds more tests for `num_blocks`, `num_rows`, `estimate_size_bytes`, and `len(queue)` ## Related issues ## Additional information --------- Signed-off-by: iamjustinhsu <jhsu@anyscale.com> Signed-off-by: peterxcli <peterxcli@gmail.com>

…eue` (ray-project#60538) ray-project#60017 and ray-project#60228 refactored the `FIFOBundleQueue` interface and renamed `FIFOBundleQueue.popleft` with `FIFOBundleQueue.get_next`. However, this name change wasn't reflected in the `UnionOperator` implementation, and as a result the operator can error when it clears its output queue. This change also fixes the flaky `test_union.py`. Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu> Signed-off-by: peterxcli <peterxcli@gmail.com>

[data][1/n] Base + HashLinkedBundleQueue

0a23374

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

iamjustinhsu requested a review from a team as a code owner January 9, 2026 22:34

clean

e6aa407

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

gemini-code-assist bot reviewed Jan 9, 2026

View reviewed changes

iamjustinhsu added 5 commits January 9, 2026 14:37

clean again

2786713

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

fix test

24571b2

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

fix

1c47693

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

update test

12bd273

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

raise error docstring

3d221c9

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

gemini-code-assist bot reviewed Jan 9, 2026

View reviewed changes

python/ray/data/_internal/execution/bundle_queue/hash_link.py Show resolved Hide resolved

python/ray/data/_internal/execution/bundle_queue/base.py Show resolved Hide resolved

python/ray/data/_internal/execution/bundle_queue/base.py Show resolved Hide resolved

cursor bot reviewed Jan 9, 2026

View reviewed changes

python/ray/data/_internal/execution/operators/output_splitter.py Show resolved Hide resolved

fix add_to_front

05daa16

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

gemini-code-assist bot reviewed Jan 9, 2026

View reviewed changes

iamjustinhsu added 2 commits January 9, 2026 15:30

correct type annotations

a7c6ae3

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

remove obselete file

e5ccfdb

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

iamjustinhsu force-pushed the jhsu/unify-internal-queue-abstraction-add-base branch from 2920cce to e5ccfdb Compare January 9, 2026 23:31

iamjustinhsu added 5 commits January 9, 2026 16:13

use HashLinkedQueue for now

6e0d741

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

lint

9f0520d

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

one more

906aa0d

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

Merge branch 'master' of https://github.com/ray-project/ray into jhsu…

ca67546

…/unify-internal-queue-abstraction-add-base

remove add_to_front

8cb837c

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

cursor bot reviewed Jan 10, 2026

View reviewed changes

python/ray/data/_internal/execution/operators/output_splitter.py Outdated Show resolved Hide resolved

python/ray/data/_internal/execution/operators/output_splitter.py Show resolved Hide resolved

ray-gardener bot added python Pull requests that update Python code data Ray Data-related issues labels Jan 10, 2026

fix

70a8536

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

alexeykudinkin reviewed Jan 13, 2026

View reviewed changes

goutamvenkat-anyscale reviewed Jan 13, 2026

View reviewed changes

address feedback + tidy up

3682450

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

alexeykudinkin approved these changes Jan 14, 2026

View reviewed changes

cursor bot reviewed Jan 14, 2026

View reviewed changes

python/ray/data/_internal/execution/bundle_queue/hash_link.py Outdated Show resolved Hide resolved

python/ray/data/_internal/execution/bundle_queue/hash_link.py Outdated Show resolved Hide resolved

add _remove_inner

8510969

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

iamjustinhsu force-pushed the jhsu/unify-internal-queue-abstraction-add-base branch from 5b702ff to 8510969 Compare January 15, 2026 08:03

iamjustinhsu added the go add ONLY when ready to merge, run all tests label Jan 15, 2026

alexeykudinkin approved these changes Jan 16, 2026

View reviewed changes

alexeykudinkin merged commit f6d9b28 into ray-project:master Jan 16, 2026
6 checks passed

iamjustinhsu mentioned this pull request Jan 20, 2026

[data] Update bundle queue tests #60326

Merged

bveeramani mentioned this pull request Jan 27, 2026

[Data] Fix attribute error in UnionOperator.clear_internal_output_queue #60538

Merged

	self._input_buffers: List["BaseBundleQueue"] = [
	self._input_buffers: List[HashLinkedQueue] = [

		from ray.tests.conftest import * # noqa


		@pytest.mark.parametrize("equal", [False, True])

		# Remove will dequeue the bundle, therefore, we override
		# get_next(), _get_next_inner()

Conversation

iamjustinhsu commented Jan 9, 2026

Description

Related issues

Additional information

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

iamjustinhsu commented Jan 9, 2026

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

iamjustinhsu commented Jan 9, 2026

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jan 9, 2026

Choose a reason for hiding this comment

Uh oh!

iamjustinhsu Jan 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

iamjustinhsu Jan 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

iamjustinhsu Jan 9, 2026 •

edited

Loading

iamjustinhsu Jan 13, 2026 •

edited

Loading