[data] Unify bundle queue abstraction by iamjustinhsu · Pull Request #59093 · ray-project/ray

iamjustinhsu · 2025-12-01T19:30:07Z

Description

Currently we support many types of bundle queues inside operators:

BlockRefBundler
FIFOBundleQueue
UnorderedOutputQueue
OrderedOutputQueue
List[RefBundle]
Deque[RefBundle]
StreamingRepartitionRefBundler

All these independent classes makes it very awkward to record queue metrics within operators. In fact, operators themselves have to manage it themselves. This PR introduces a base class called BaseBundleQueue that hopes to unify all the independent classes together.

class BaseBundleQueue:
    """Base class for storing bundles."""

    @abc.abstractmethod
    def add(self, bundle: RefBundle, **kwargs: Any) -> None:
        ...

    @abc.abstractmethod
    def get_next(self) -> RefBundle:
        ...

    @abc.abstractmethod
    def peek_next(self) -> Optional[RefBundle]:
        ...

    @abc.abstractmethod
    def has_next(self) -> bool:
        ...

    @abc.abstractmethod
    def clear(self):
        ...

    @abc.abstractmethod
    def finalize(self, **kwargs: Any):
        ...

NOTE:

I mostly refractored and moved stuff into new files, so that's why the lines are long. However, I did make some changes to all the existing bundle queues, namely their get_next() method to support a peak method in the base class. I will add tests for these once finalized.

With this approach, I hope to

Simplify design of future operators with InternalQueueOperatorMixin, that will automatically get size of queues, as well as clearing queues for free in the mixin
Dedicated queue metrics, so that operators themselves don't need to track
Lower cognitive overhead of trying to understand what type of queue is being used within each operator. You now know it's always a BaseBundleQueue that stores RefBundles.

Future Work

This PR is a stepping stone to completely refractor how metrics are collected, stored, and exported. I'd like to refractor OpRuntimeMetrics because it assumes that all operators are MapOperators, which isn't true. After this PR, I have a few proposals to refractor OpRuntimeMetrics

Create a BaseOpMetrics class, so that individual operators can override BaseOpMetrics. For operators with internal queues, we can automatically hook in QueueOpMetrics(BaseOpMetrics) too
Allow classes themselves to report to _StatsActor directly through a something like a MetricsRegistry interface. This will require a lot of work, and I will have to repurpose the _StatsActor to read from that Registry.
I'm also thinking we can just completely remove _StatsActor and _StatsManager, but I'm not sure if that is very clean.

Update on 12/2

Here is a summary of extensions

Update on 12/29

New summary of extensions

Tests

TBD

Related issues

None

Additional information

None

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

gemini-code-assist

Code Review

This pull request introduces a BaseBundleQueue abstraction to unify the various queue implementations within Ray Data operators. This is a significant and valuable refactoring that simplifies operator design and centralizes queue metrics. The overall approach is solid, and the new queue implementations in fifo.py, hash_link.py, and ordered.py are well-designed.

However, I've found several issues with the implementation that need to be addressed:

There are a few critical correctness bugs, including a missing import in bundle_queue/__init__.py and incorrect buffer handling in OutputSplitter.
The rebundling queues (RebundleQueue and ExactRebundleQueue) have issues: one violates the peek_next contract by having side effects, and the other has a metric leak in get_next_with_original.
The transition to queue-based metrics is incomplete in MapOperator, leading to redundant and confusing metric tracking.
There are also some minor issues with incorrect type hints and misleading docstrings.

I've left detailed comments on each of these points. Addressing them will ensure the new abstraction is robust and correctly implemented.

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

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

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

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

python/ray/data/tests/test_actor_pool_map_operator.py

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

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

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

python/ray/data/tests/test_actor_pool_map_operator.py

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

omatthew98

Mostly left comments surrounding the high level method names / APIs. I think solid start but some opportunities for tightening the naming up since we are already doing a big shake up. LMK if you need another review.

omatthew98 · 2025-12-02T19:54:27Z

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

@@ -1,75 +0,0 @@
-import abc


Is there a reason to fully scrap BundleQueue in favor of BaseBundleQueue?

Oh, it's mostly a rename into the base.py (i'm gonna make clear what is renamed and what is not after this comment), the only difference is that the base also contains a method called done_adding_bundles. This method is used for BlockRefBundler, StreamingRepartitionRefBunlder, and OrderedOutputQueue

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

omatthew98 · 2025-12-02T20:22:24Z

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

+    """Protocol for storing bundles AND supporting remove(bundle)
+    and contains(bundle) operations quickly."""
+
+    def __contains__(self, bundle: RefBundle) -> bool:


Is __contains__ not supported for the other queues? Should we maybe specify in the comment that this should both implement the method and be efficient or there might be degraded performance?

yea it can be supported for most queues, but I opted to not move it into the base implementation because

our existing code infrequently uses __contains__

BlockRefBundler and StreamingRefBundler have weird semantics for __contains__, because one moment the bundle is there, and the next it's gone due to "rebundling".
On that note of bullet # 2, rather than have SupportsRebundling be a protocol, we can have it become SupportsRebundling(BaseBundleQueue), which only supports rebundling, and won't contain a __contains__ method. Haven't thought too deeply into this approach

Ok yeah the current approach seems reasonable to me then. I think fine to keep the protocols as is rather than creating a more piecemeal approach where we have very tiny protocols.

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

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

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

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

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

cursor · 2026-01-07T18:53:37Z

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

+                "This is a bug in the Ray Data code."
+            )
+            remaining_bundle: Optional[RefBundle] = None
+            if rows_needed < last_pending_bundle.num_rows():


TypeError when comparing with None num_rows result

Low Severity

The comparison rows_needed < last_pending_bundle.num_rows() at line 183 and the assertion rows_needed == last_pending_bundle.num_rows() at line 187 directly call num_rows() which can return Optional[int]. If num_rows() returns None, these comparisons would raise a TypeError. While the built-in strategies (EstimateSize, ExactMultipleSize) protect against this via earlier assertions, custom RebundlingStrategy implementations could trigger this error. The comparisons could defensively handle None using or 0 like other parts of the codebase.

Additional Locations (1)

python/ray/data/_internal/execution/bundle_queue/bundler.py#L186-L187

if it's None, that means we default to 0. The bundle can't have 0 rows because we know it triggered a ready bundle.

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

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

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

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

## 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 Moves OrderedBundleQueue to `bundle_queue/` directory, and rebases off of the `BaseBundleQueue` class. Creates a `FIFOBundleQueue` for simple ref bundle queues. Old PR: #59093 There are about ~400 lines of tests, so not actually that big ## Related issues ## Additional information --------- Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

## Description Moves OrderedBundleQueue to `bundle_queue/` directory, and rebases off of the `BaseBundleQueue` class. Creates a `FIFOBundleQueue` for simple ref bundle queues. Old PR: ray-project#59093 There are about ~400 lines of tests, so not actually that big ## Related issues ## Additional information --------- Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

## Description Moves OrderedBundleQueue to `bundle_queue/` directory, and rebases off of the `BaseBundleQueue` class. Creates a `FIFOBundleQueue` for simple ref bundle queues. Old PR: ray-project#59093 There are about ~400 lines of tests, so not actually that big ## Related issues ## Additional information --------- Signed-off-by: iamjustinhsu <jhsu@anyscale.com> Signed-off-by: jinbum-kim <jinbum9958@gmail.com>

## Description Moves OrderedBundleQueue to `bundle_queue/` directory, and rebases off of the `BaseBundleQueue` class. Creates a `FIFOBundleQueue` for simple ref bundle queues. Old PR: ray-project#59093 There are about ~400 lines of tests, so not actually that big ## Related issues ## Additional information --------- Signed-off-by: iamjustinhsu <jhsu@anyscale.com> 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 Old PR: ray-project#59093 ## Related issues ## Additional information --------- Signed-off-by: iamjustinhsu <jhsu@anyscale.com> Signed-off-by: peterxcli <peterxcli@gmail.com>

## Description Moves OrderedBundleQueue to `bundle_queue/` directory, and rebases off of the `BaseBundleQueue` class. Creates a `FIFOBundleQueue` for simple ref bundle queues. Old PR: ray-project#59093 There are about ~400 lines of tests, so not actually that big ## Related issues ## Additional information --------- Signed-off-by: iamjustinhsu <jhsu@anyscale.com> 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 Moves OrderedBundleQueue to `bundle_queue/` directory, and rebases off of the `BaseBundleQueue` class. Creates a `FIFOBundleQueue` for simple ref bundle queues. Old PR: ray-project#59093 There are about ~400 lines of tests, so not actually that big ## Related issues ## Additional information --------- Signed-off-by: iamjustinhsu <jhsu@anyscale.com> Signed-off-by: peterxcli <peterxcli@gmail.com>

[data] Unify bundle queue abstraction

8334317

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

iamjustinhsu requested a review from a team as a code owner December 1, 2025 19:30

gemini-code-assist bot reviewed Dec 1, 2025

View reviewed changes

cursor bot reviewed Dec 1, 2025

View reviewed changes

python/ray/data/tests/test_actor_pool_map_operator.py Outdated Show resolved Hide resolved

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

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

fix based on AI

1f70039

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

cursor bot reviewed Dec 1, 2025

View reviewed changes

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

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

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

iamjustinhsu added 3 commits December 1, 2025 12:17

override clear method

0f4c843

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

fix

f970f2d

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

docstring

2bf9128

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

cursor bot reviewed Dec 1, 2025

View reviewed changes

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

iamjustinhsu added 2 commits December 1, 2025 13:06

reset metrics on clear

ae0a56d

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

doc

b4a35df

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

cursor bot reviewed Dec 1, 2025

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

lint

e24b424

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

cursor bot reviewed Dec 1, 2025

View reviewed changes

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

iamjustinhsu added 2 commits December 1, 2025 14:42

don't fix the streaming repartition builder in this PR

ad1a792

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

remove random code

e5a1398

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

cursor bot reviewed Dec 1, 2025

View reviewed changes

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

iamjustinhsu added 2 commits December 1, 2025 16:24

deque does not support slicing

6ca2ce8

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

fix docstring

e82e4a5

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

cursor bot reviewed Dec 2, 2025

View reviewed changes

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

cursor

af39037

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

ray-gardener bot added the data Ray Data-related issues label Dec 2, 2025

omatthew98 reviewed Dec 2, 2025

View reviewed changes

address mowen awesome feedback :)

0f60342

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

iamjustinhsu requested review from a team as code owners December 2, 2025 23:11

iamjustinhsu added 2 commits December 2, 2025 15:21

address mowen awesome feedback :) pt2

3ad6b5f

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

remove random code

1a58f57

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

iamjustinhsu force-pushed the jhsu/unify-internal-queue-abstraction branch from 2ddb01b to 1a58f57 Compare December 2, 2025 23:32

iamjustinhsu added 2 commits December 30, 2025 17:45

fix

bc54d3b

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

fix test_stats

3690b7d

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

iamjustinhsu mentioned this pull request Jan 6, 2026

[data] Fix StreamingRepartition hang with empty upstream results #59848

Merged

iamjustinhsu added 2 commits January 6, 2026 14:07

revert the bundle queue metrics tracking change

d478472

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

rebase

cdab451

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

cursor bot reviewed Jan 6, 2026

View reviewed changes

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

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

iamjustinhsu added 2 commits January 6, 2026 15:50

fix

452b436

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

use mod since streaming ref partition bundler actually needs multiples

27e1580

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

cursor bot reviewed Jan 7, 2026

View reviewed changes

iamjustinhsu added 2 commits January 7, 2026 13:53

fix bundle consumption

a28f641

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

lint

acf1c2c

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

cursor bot reviewed Jan 7, 2026

View reviewed changes

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

reassign empty list

712ad55

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

This was referenced Jan 9, 2026

[data][1/n] Base + HashLinkedBundleQueue #60017

Merged

[data][2/n] FIFO + Ordered Bundle Queue iamjustinhsu/ray#5

Closed

alexeykudinkin pushed a commit that referenced this pull request Jan 16, 2026

[data][1/n] Base + HashLinkedBundleQueue (#60017)

f6d9b28

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

iamjustinhsu mentioned this pull request Jan 16, 2026

[data][2/n] Ordered + FIFO bundle queue #60228

Merged

iamjustinhsu added the unstale A PR that has been marked unstale. It will not get marked stale again if this label is on it. label Jan 21, 2026

iamjustinhsu mentioned this pull request Jan 26, 2026

[data][3/N] Add rebundle queue #60487

Open

Conversation

iamjustinhsu commented Dec 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

NOTE:

Future Work

Update on 12/2

Update on 12/29

Tests

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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

omatthew98 left a comment

Choose a reason for hiding this comment

Uh oh!

omatthew98 Dec 2, 2025

Choose a reason for hiding this comment

Uh oh!

iamjustinhsu Dec 2, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

omatthew98 Dec 2, 2025

Choose a reason for hiding this comment

Uh oh!

iamjustinhsu Dec 2, 2025 • 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!

omatthew98 Dec 3, 2025 • 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!

Uh oh!

Uh oh!

Uh oh!

cursor bot Jan 7, 2026

Choose a reason for hiding this comment

TypeError when comparing with None num_rows result

Uh oh!

iamjustinhsu Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

iamjustinhsu commented Dec 1, 2025 •

edited

Loading

iamjustinhsu Dec 2, 2025 •

edited

Loading

omatthew98 Dec 3, 2025 •

edited

Loading