fix

python/ray/data/dataset.py

@@ -2658,9 +2658,7 @@ def write_datasource(
  # WriteResult (one element per write task). Otherwise, an error will
  # be raised. The Datasource can handle execution outcomes with the
  # on_write_complete() and on_write_failed().
- def transform(
- blocks: Iterable[Block], ctx, fn
- ) -> Iterable[Block]:
+ def transform(blocks: Iterable[Block], ctx, fn) -> Iterable[Block]:
  return [[datasource.write(blocks, ctx, **write_args)]]
 
  plan = self._plan.with_stage(
@@ -2675,7 +2673,7 @@ def transform(
  try:
  self._write_ds = Dataset(plan, self._epoch, self._lazy).fully_executed()
  datasource.on_write_complete(
- self._write_ds._plan.execute().get_blocks()
+ ray.get(self._write_ds._plan.execute().get_blocks())
  )
  except Exception as e:
  datasource.on_write_failed([], e)

python/ray/data/datasource/datasource.py

@@ -362,8 +362,8 @@ def write(
  ray.get(tasks)
  return "ok"
 
- def on_write_complete(self, write_results: List[ObjectRef[WriteResult]]) -> None:
- assert all(ray.get(w) == ["ok"] for w in write_results), write_results
+ def on_write_complete(self, write_results: List[WriteResult]) -> None:
+ assert all(w == ["ok"] for w in write_results), write_results
  self.num_ok += 1
 
  def on_write_failed(

python/ray/data/datasource/file_based_datasource.py

@@ -17,10 +17,8 @@
 
 from ray.data._internal.arrow_block import ArrowRow
 from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
-from ray.data._internal.block_list import BlockMetadata
 from ray.data._internal.execution.interfaces import TaskContext
 from ray.data._internal.output_buffer import BlockOutputBuffer
-from ray.data._internal.remote_fn import cached_remote_fn
 from ray.data._internal.util import _check_pyarrow_version, _resolve_custom_scheme
 from ray.data.block import Block, BlockAccessor
 from ray.data.context import DatasetContext
@@ -36,7 +34,7 @@
 )
 
 from ray.types import ObjectRef
-from ray.util.annotations import Deprecated, DeveloperAPI, PublicAPI
+from ray.util.annotations import DeveloperAPI, PublicAPI
 from ray._private.utils import _add_creatable_buckets_param_if_s3_uri
 
 if TYPE_CHECKING:

python/ray/data/datasource/mongo_datasource.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Dict, List, Optional, TYPE_CHECKING
+from typing import Dict, List, Optional, TYPE_CHECKING
 
 from ray.data.datasource.datasource import Datasource, Reader, ReadTask, WriteResult
 from ray.data.block import (
@@ -8,10 +8,8 @@
  BlockMetadata,
 )
 from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
-from ray.data._internal.remote_fn import cached_remote_fn
 from ray.data._internal.execution.interfaces import TaskContext
-from ray.types import ObjectRef
-from ray.util.annotations import Deprecated, PublicAPI
+from ray.util.annotations import PublicAPI
 from typing import Iterable
 
 if TYPE_CHECKING:

python/ray/data/tests/test_dataset.py

@@ -5516,6 +5516,7 @@ def test_ragged_tensors(ray_start_regular_shared):
  ArrowVariableShapedTensorType(dtype=new_type, ndim=3),
  ]
 
+
 class LoggerWarningCalled(Exception):
  """Custom exception used in test_warning_execute_with_no_cpu() and
  test_nowarning_execute_with_cpu(). Raised when the `logger.warning` method

python/ray/data/tests/test_dataset_formats.py

@@ -267,7 +267,7 @@ def write(b):
  return "ok"
 
  def on_write_complete(self, write_results: List[WriteResult]) -> None:
- assert all(ray.get(w) == ["ok"] for w in write_results), write_results
+ assert all(w == ["ok"] for w in write_results), write_results
  self.num_ok += 1
 
  def on_write_failed(