Streaming executor fixes #4 (ray-project#32882)

edoakes · Mar 1, 2023 · 2875357 · 2875357
1 parent 0c02ae8
commit 2875357
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/python/ray/data/tests/test_bulk_executor.py b/python/ray/data/tests/test_bulk_executor.py
@@ -125,6 +125,7 @@ def test_actor_strategy(ray_start_10_cpus_shared):
 
 def test_new_execution_backend_invocation(ray_start_10_cpus_shared):
     DatasetContext.get_current().new_execution_backend = True
+    DatasetContext.get_current().use_streaming_executor = False
     # Read-only: will use legacy executor for now.
     ds = ray.data.range(10)
     assert ds.take_all() == list(range(10))

diff --git a/python/ray/data/tests/test_dataset.py b/python/ray/data/tests/test_dataset.py
@@ -382,12 +382,13 @@ def test_zip_different_num_blocks_split_smallest(
         [{str(i): i for i in range(num_cols1, num_cols1 + num_cols2)}] * n,
         parallelism=num_blocks2,
     )
-    ds = ds1.zip(ds2)
+    ds = ds1.zip(ds2).fully_executed()
+    num_blocks = ds._plan._snapshot_blocks.executed_num_blocks()
     assert ds.take() == [{str(i): i for i in range(num_cols1 + num_cols2)}] * n
     if should_invert:
-        assert ds.num_blocks() == num_blocks2
+        assert num_blocks == num_blocks2
     else:
-        assert ds.num_blocks() == num_blocks1
+        assert num_blocks == num_blocks1
 
 
 def test_zip_pandas(ray_start_regular_shared):