[Data] Deprecate num_blocks parameter from Dataset.random_shuffle (…

…#41111) This PR is to deprecate `num_blocks` parameter from `random_shuffle()` This is to avoid surfacing block concept from our API. Signed-off-by: Cheng Su <[email protected]>
ray-project · Nov 14, 2023 · ce55056 · ce55056
1 parent b441fdd
commit ce55056
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/python/ray/data/_internal/logical/operators/all_to_all_operator.py b/python/ray/data/_internal/logical/operators/all_to_all_operator.py
@@ -60,13 +60,11 @@ def __init__(
  input_op: LogicalOperator,
  name: str = "RandomShuffle",
  seed: Optional[int] = None,
- num_outputs: Optional[int] = None,
  ray_remote_args: Optional[Dict[str, Any]] = None,
  ):
  super().__init__(
  name,
  input_op,
- num_outputs=num_outputs,
  sub_progress_bar_names=[
  ExchangeTaskSpec.MAP_SUB_PROGRESS_BAR_NAME,
  ExchangeTaskSpec.REDUCE_SUB_PROGRESS_BAR_NAME,

diff --git a/python/ray/data/dataset.py b/python/ray/data/dataset.py
@@ -1068,13 +1068,19 @@ def random_shuffle(
  Args:
  seed: Fix the random seed to use, otherwise one is chosen
  based on system randomness.
- num_blocks: The number of output blocks after the shuffle, or ``None``
- to retain the number of blocks.
 
  Returns:
  The shuffled :class:`Dataset`.
  """ # noqa: E501
 
+ if num_blocks is not None:
+ warnings.warn(
+ "`num_blocks` parameter is deprecated in Ray 2.9. random_shuffle() "
+ "does not support to change the number of output blocks. Use "
+ "repartition() instead.", # noqa: E501
+ DeprecationWarning,
+ )
+
  plan = self._plan.with_stage(
  RandomShuffleStage(seed, num_blocks, ray_remote_args)
  )
@@ -1084,7 +1090,6 @@ def random_shuffle(
  op = RandomShuffle(
  logical_plan.dag,
  seed=seed,
- num_outputs=num_blocks,
  ray_remote_args=ray_remote_args,
  )
  logical_plan = LogicalPlan(op)