[Data] Make Zip operator streaming by owenowenisme · Pull Request #56504 · ray-project/ray

owenowenisme · 2025-09-13T09:25:49Z

Why are these changes needed?

Related issue number

Checks

Test script to show streaming works:

import time
import ray

ray.init()

def slow_batch(batch):
    # Sleep 1s per row in this batch.
    time.sleep(len(batch))
    return batch

# Two inputs of equal length
n = 100

ds1 = ray.data.range(n)
ds2 = ray.data.range(n).map_batches(slow_batch, batch_size=1)

t0 = time.time()
out = ds1.zip(ds2).take(3)
elapsed = time.time() - t0

print("Output:", out)
print(f"Elapsed: {elapsed:.2f}s")

With original batching zip we get:

Output: [{'id': 0, 'id_1': 0}, {'id': 1, 'id_1': 1}, {'id': 2, 'id_1': 2}]                                                                                      
Elapsed: 8.41s

With streaming zip we get :

Output: [{'id': 0, 'id_1': 0}, {'id': 1, 'id_1': 1}, {'id': 2, 'id_1': 2}]                                                                                      
Elapsed: 3.07s

Note that now zip operator would need at least N + 1 cpus for ds1.zip(ds2,ds3,...,dsN) in order to streaming.

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
I've run scripts/format.sh to lint the changes in this PR.
I've included any doc changes needed for https://docs.ray.io/en/master/.
- I've added any new APIs to the API Reference. For example, if I added a
  method in Tune, I've added it in doc/source/tune/api/ under the
  corresponding .rst file.
I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
Testing Strategy
- Unit tests
- Release tests
- This PR is not tested :(

Note

Refactors ZipOperator to stream and zip inputs on-the-fly across N inputs with leftover handling, adjusts stats/CPU needs, and updates tests accordingly.

Data Execution:
- ZipOperator:
  - Convert from bulk to streaming: zip outputs produced incrementally during input ingestion.
  - Generalize _zip to N-ary inputs with per-input buffers and _leftover_blocks handling.
  - Skip empty-row blocks; enqueue output as soon as a zipped block is formed.
  - Add check_still_has_blocks, _zip_blocks (chain zipping), and initialize stats as {"Zip": []}.
  - Assert equal total rows across inputs at completion with clearer error.
Tests:
- test_zip.py: simplify split behavior expectations; remove inversion parameter; ensure block count matches first dataset; add multi-dataset zipping coverage.
- test_dynamic_block_split.py: increase required CPUs to 3 to accommodate streaming zip.

^{Written by Cursor Bugbot for commit 606bbde. This will update automatically on new commits. Configure here.}

Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

…rable dataset size Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

…rator-accept-multiple-input

Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

## Why are these changes needed? Before making zip operator a streaming operator, we make it accept multiple input first. Now Zip operator can be used with ```py >>> import ray >>> ds1 = ray.data.range(5) >>> ds2 = ray.data.range(5) >>> ds3 = ray.data.range(5) >>> ds1.zip(ds2, ds3).take_batch() {'id': array([0, 1, 2, 3, 4]), 'id_1': array([0, 1, 2, 3, 4]), 'id_2': array([0, 1, 2, 3, 4])} >>> ds1.zip(ds2, ds3).take_all() [{'id': 0, 'id_1': 0, 'id_2': 0}, {'id': 1, 'id_1': 1, 'id_2': 1}, {'id': 2, 'id_1': 2, 'id_2': 2}, {'id': 3, 'id_1': 3, 'id_2': 3}, {'id': 4, 'id_1': 4, 'id_2': 4}] >>> ```  ## Related issue number #56504  Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

…rator-streaming

Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

owenowenisme · 2025-09-25T16:19:58Z

python/ray/data/tests/test_zip.py

    )


-@pytest.mark.parametrize(


Removed this test because we don't do invert when streaming.

Could we remove the invert option and retain this test?

Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

## Why are these changes needed? Before making zip operator a streaming operator, we make it accept multiple input first. Now Zip operator can be used with ```py >>> import ray >>> ds1 = ray.data.range(5) >>> ds2 = ray.data.range(5) >>> ds3 = ray.data.range(5) >>> ds1.zip(ds2, ds3).take_batch() {'id': array([0, 1, 2, 3, 4]), 'id_1': array([0, 1, 2, 3, 4]), 'id_2': array([0, 1, 2, 3, 4])} >>> ds1.zip(ds2, ds3).take_all() [{'id': 0, 'id_1': 0, 'id_2': 0}, {'id': 1, 'id_1': 1, 'id_2': 1}, {'id': 2, 'id_1': 2, 'id_2': 2}, {'id': 3, 'id_1': 3, 'id_2': 3}, {'id': 4, 'id_1': 4, 'id_2': 4}] >>> ```  ## Related issue number #56504  Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com> Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>

Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

srinathk10 · 2025-09-29T17:26:05Z

python/ray/data/tests/test_dynamic_block_split.py

    # We need at least 2 CPUs to run a actorpool streaming
-    ray.init(num_cpus=2, object_store_memory=1e9)
+    # We need at least 3 CPUs to run a zip streaming
+    ray.init(num_cpus=4, object_store_memory=1e9)


Did you mean to set to num_cpus=3?

Yeah, thanks for pointing this out.

srinathk10 · 2025-09-29T17:28:11Z

python/ray/data/tests/test_zip.py

    )


-@pytest.mark.parametrize(


Could we remove the invert option and retain this test?

Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

srinathk10

LGTM

alexeykudinkin · 2025-10-02T19:27:29Z

@owenowenisme i think we actually can do more radical simplification of the impl here:

Our goal is to make sure that the blocks produced by the zip a) have the same # of rows and b) hstack the individual blocks from all sequences one by one
To avoid complicated logic of resizing we can actually leverage StreamingRepartition w/ preset target # of rows
Then the only thing zip needs to do is to hstack the blocks

I've discussed this with @srinathk10 and he's gonna help you experiment on this and see it through

owenowenisme · 2025-10-04T01:32:40Z

Update: This PR depends on #57165

## Why are these changes needed? Before making zip operator a streaming operator, we make it accept multiple input first. Now Zip operator can be used with ```py >>> import ray >>> ds1 = ray.data.range(5) >>> ds2 = ray.data.range(5) >>> ds3 = ray.data.range(5) >>> ds1.zip(ds2, ds3).take_batch() {'id': array([0, 1, 2, 3, 4]), 'id_1': array([0, 1, 2, 3, 4]), 'id_2': array([0, 1, 2, 3, 4])} >>> ds1.zip(ds2, ds3).take_all() [{'id': 0, 'id_1': 0, 'id_2': 0}, {'id': 1, 'id_1': 1, 'id_2': 1}, {'id': 2, 'id_1': 2, 'id_2': 2}, {'id': 3, 'id_1': 3, 'id_2': 3}, {'id': 4, 'id_1': 4, 'id_2': 4}] >>> ```  ## Related issue number ray-project#56504  Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com> Signed-off-by: Douglas Strodtman <douglas@anyscale.com>

## Why are these changes needed? Before making zip operator a streaming operator, we make it accept multiple input first. Now Zip operator can be used with ```py >>> import ray >>> ds1 = ray.data.range(5) >>> ds2 = ray.data.range(5) >>> ds3 = ray.data.range(5) >>> ds1.zip(ds2, ds3).take_batch() {'id': array([0, 1, 2, 3, 4]), 'id_1': array([0, 1, 2, 3, 4]), 'id_2': array([0, 1, 2, 3, 4])} >>> ds1.zip(ds2, ds3).take_all() [{'id': 0, 'id_1': 0, 'id_2': 0}, {'id': 1, 'id_1': 1, 'id_2': 1}, {'id': 2, 'id_1': 2, 'id_2': 2}, {'id': 3, 'id_1': 3, 'id_2': 3}, {'id': 4, 'id_1': 4, 'id_2': 4}] >>> ```  ## Related issue number ray-project#56504  Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

github-actions · 2025-10-20T12:25:41Z

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

owenowenisme · 2025-11-17T06:44:06Z

Opening a new one.

## Why are these changes needed? Before making zip operator a streaming operator, we make it accept multiple input first. Now Zip operator can be used with ```py >>> import ray >>> ds1 = ray.data.range(5) >>> ds2 = ray.data.range(5) >>> ds3 = ray.data.range(5) >>> ds1.zip(ds2, ds3).take_batch() {'id': array([0, 1, 2, 3, 4]), 'id_1': array([0, 1, 2, 3, 4]), 'id_2': array([0, 1, 2, 3, 4])} >>> ds1.zip(ds2, ds3).take_all() [{'id': 0, 'id_1': 0, 'id_2': 0}, {'id': 1, 'id_1': 1, 'id_2': 1}, {'id': 2, 'id_1': 2, 'id_2': 2}, {'id': 3, 'id_1': 3, 'id_2': 3}, {'id': 4, 'id_1': 4, 'id_2': 4}] >>> ```  ## Related issue number ray-project#56504  Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

## Why are these changes needed? Before making zip operator a streaming operator, we make it accept multiple input first. Now Zip operator can be used with ```py >>> import ray >>> ds1 = ray.data.range(5) >>> ds2 = ray.data.range(5) >>> ds3 = ray.data.range(5) >>> ds1.zip(ds2, ds3).take_batch() {'id': array([0, 1, 2, 3, 4]), 'id_1': array([0, 1, 2, 3, 4]), 'id_2': array([0, 1, 2, 3, 4])} >>> ds1.zip(ds2, ds3).take_all() [{'id': 0, 'id_1': 0, 'id_2': 0}, {'id': 1, 'id_1': 1, 'id_2': 1}, {'id': 2, 'id_1': 2, 'id_2': 2}, {'id': 3, 'id_1': 3, 'id_2': 3}, {'id': 4, 'id_1': 4, 'id_2': 4}] >>> ```  ## Related issue number ray-project#56504  Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com> Signed-off-by: Future-Outlier <eric901201@gmail.com>

inherit NAry operator first

4fd074c

Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

owenowenisme force-pushed the data/make-zip-operator-streaming branch from 61f995d to 4fd074c Compare September 13, 2025 09:27

owenowenisme and others added 3 commits September 13, 2025 18:23

fix zip operator

13441b7

Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

make zip operator accept list of input

3de48b1

Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

Merge branch 'master' into data/make-zip-operator-streaming

fb2eaff

owenowenisme mentioned this pull request Sep 15, 2025

[Data] Make zip operator accept multiple input #56524

Merged

8 tasks

owenowenisme added 4 commits September 16, 2025 03:00

add assert for input_op length > 2

3e70237

Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

modify test for zip operator to make it parametrize test with configu…

1548a44

…rable dataset size Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

Merge remote-tracking branch 'upstream/master' into data/make-zip-ope…

19dbb20

…rator-accept-multiple-input

update pydo lint

592479d

Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

owenowenisme added 2 commits September 25, 2025 16:14

Merge remote-tracking branch 'upstream/master' into data/make-zip-ope…

001cd27

…rator-streaming

update

c5433fd

Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

owenowenisme force-pushed the data/make-zip-operator-streaming branch from 28f95a2 to c5433fd Compare September 25, 2025 16:15

owenowenisme commented Sep 25, 2025

View reviewed changes

owenowenisme force-pushed the data/make-zip-operator-streaming branch 2 times, most recently from a94c80c to c5433fd Compare September 26, 2025 05:36

increase test ray cluster cpu

d00aed5

Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

owenowenisme marked this pull request as ready for review September 26, 2025 14:51

owenowenisme requested a review from a team as a code owner September 26, 2025 14:51

This comment was marked as outdated.

Sign in to view

fix comment

810da7e

Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

ray-gardener bot added data Ray Data-related issues community-contribution Contributed by the community labels Sep 26, 2025

Add better assertion message and add back todo comment

428f588

Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

This comment was marked as outdated.

Sign in to view

srinathk10 reviewed Sep 29, 2025

View reviewed changes

resolve comment

606bbde

Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

This comment was marked as outdated.

Sign in to view

srinathk10 added the go add ONLY when ready to merge, run all tests label Sep 30, 2025

srinathk10 approved these changes Oct 1, 2025

View reviewed changes

Merge branch 'master' into data/make-zip-operator-streaming

990fa19

github-actions bot added the stale The issue is stale. It will be closed within 7 days unless there are further conversation label Oct 20, 2025

owenowenisme added unstale A PR that has been marked unstale. It will not get marked stale again if this label is on it. and removed stale The issue is stale. It will be closed within 7 days unless there are further conversation labels Oct 20, 2025

owenowenisme closed this Nov 17, 2025

		)


		@pytest.mark.parametrize(

		)


		@pytest.mark.parametrize(

Conversation

owenowenisme commented Sep 13, 2025 • edited by cursor bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why are these changes needed?

Related issue number

Checks

Uh oh!

owenowenisme Sep 25, 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!

srinathk10 Sep 29, 2025

Choose a reason for hiding this comment

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

srinathk10 Sep 29, 2025

Choose a reason for hiding this comment

Uh oh!

owenowenisme Oct 1, 2025

Choose a reason for hiding this comment

Uh oh!

srinathk10 Sep 29, 2025

Choose a reason for hiding this comment

Uh oh!

This comment was marked as outdated.

Uh oh!

srinathk10 left a comment

Choose a reason for hiding this comment

Uh oh!

alexeykudinkin commented Oct 2, 2025

Uh oh!

owenowenisme commented Oct 4, 2025

Uh oh!

github-actions bot commented Oct 20, 2025

Uh oh!

owenowenisme commented Nov 17, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

owenowenisme commented Sep 13, 2025 •

edited by cursor bot

Loading

owenowenisme Sep 25, 2025 •

edited

Loading