[Data] Add disable Block Shaping option to BlockOutputBuffer (ray-project#58757)

## Description

In addition to Block shaping by Block Size and Num Rows, add an option
to skip Block Shaping altogether in BlockOutputBuffer.

Signed-off-by: Srinath Krishnamachari <srinath.krishnamachari@anyscale.com>

Author: srinathk10

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

@@ -101,10 +101,6 @@ def __call__(
         results = self._apply_transform(ctx, batches)
         yield from self._post_process(results)
 
-    @abstractmethod
-    def _can_skip_block_sizing(self):
-        pass
-
     @property
     def output_block_size_option(self):
         return self._output_block_size_option
@@ -310,9 +306,6 @@ def _apply_transform(
     def _post_process(self, results: Iterable[MapTransformFnData]) -> Iterable[Block]:
         return self._shape_blocks(results)
 
-    def _can_skip_block_sizing(self):
-        return False
-
     def __repr__(self) -> str:
         return f"RowMapTransformFn({self._row_fn})"
 
@@ -363,12 +356,6 @@ def _apply_transform(
     def _post_process(self, results: Iterable[MapTransformFnData]) -> Iterable[Block]:
         return self._shape_blocks(results)
 
-    def _can_skip_block_sizing(self):
-        return self._output_block_size_option is None and self._batch_format in (
-            BatchFormat.ARROW,
-            BatchFormat.PANDAS,
-        )
-
     def __repr__(self) -> str:
         return f"BatchMapTransformFn({self._batch_fn=}, {self._batch_format=}, {self._batch_size=}, {self._zero_copy_batch=})"
 
@@ -419,9 +406,6 @@ def _post_process(self, results: Iterable[MapTransformFnData]) -> Iterable[Block
 
         return self._shape_blocks(results)
 
-    def _can_skip_block_sizing(self):
-        return self._output_block_size_option is None
-
     def __repr__(self) -> str:
         return (
             f"BlockMapTransformFn({self._block_fn=}, {self._output_block_size_option=})"

python/ray/data/_internal/output_buffer.py

@@ -11,11 +11,13 @@
 class OutputBlockSizeOption:
     target_max_block_size: Optional[int] = None
     target_num_rows_per_block: Optional[int] = None
+    disable_block_shaping: bool = False
 
     def __post_init__(self):
         if (
             self.target_max_block_size is None
             and self.target_num_rows_per_block is None
+            and not self.disable_block_shaping
         ):
             raise ValueError(
                 "Either `target_max_block_size` or `target_num_rows_per_block` "
@@ -27,13 +29,24 @@ def of(
         cls,
         target_max_block_size: Optional[int] = None,
         target_num_rows_per_block: Optional[int] = None,
+        disable_block_shaping: bool = False,
     ) -> Optional["OutputBlockSizeOption"]:
-        if target_max_block_size is None and target_num_rows_per_block is None:
+        if (
+            target_max_block_size is None
+            and target_num_rows_per_block is None
+            and not disable_block_shaping
+        ):
+            # In case
+            #   - Both target_max_block_size and target_num_rows_per_block are None and
+            #   - disable_block_shaping is False
+            #
+            # Buffer won't be yielding incrementally, instead producing just a single block.
             return None
         else:
             return OutputBlockSizeOption(
                 target_max_block_size=target_max_block_size,
                 target_num_rows_per_block=target_num_rows_per_block,
+                disable_block_shaping=disable_block_shaping,
             )
 
 
@@ -93,30 +106,40 @@ def finalize(self) -> None:
         self._finalized = True
 
     def _exceeded_buffer_row_limit(self) -> bool:
+        if self._output_block_size_option.disable_block_shaping:
+            return False
+
         return (
             self._max_num_rows_per_block() is not None
             and self._buffer.num_rows() > self._max_num_rows_per_block()
         )
 
     def _exceeded_buffer_size_limit(self) -> bool:
+        if self._output_block_size_option.disable_block_shaping:
+            return False
+
         return (
             self._max_bytes_per_block() is not None
             and self._buffer.get_estimated_memory_usage() > self._max_bytes_per_block()
         )
 
     def _max_num_rows_per_block(self) -> Optional[int]:
-        return (
-            self._output_block_size_option.target_num_rows_per_block
-            if self._output_block_size_option is not None
-            else None
-        )
+        if self._output_block_size_option is None:
+            return None
+
+        if self._output_block_size_option.disable_block_shaping:
+            return None
+
+        return self._output_block_size_option.target_num_rows_per_block
 
     def _max_bytes_per_block(self) -> Optional[int]:
-        return (
-            self._output_block_size_option.target_max_block_size
-            if self._output_block_size_option is not None
-            else None
-        )
+        if self._output_block_size_option is None:
+            return None
+
+        if self._output_block_size_option.disable_block_shaping:
+            return None
+
+        return self._output_block_size_option.target_max_block_size
 
     def has_next(self) -> bool:
         """Returns true when a complete output block is produced."""
@@ -130,6 +153,9 @@ def has_next(self) -> bool:
             #       is required to align it with semantic of producing 1 block
             #       from 1 block of the input
             return False
+        elif self._output_block_size_option.disable_block_shaping:
+            # When block shaping is disabled, produce blocks immediately
+            return self._buffer.num_rows() > 0
 
         return self._exceeded_buffer_row_limit() or self._exceeded_buffer_size_limit()
 

python/ray/data/tests/test_operators.py

@@ -613,44 +613,34 @@ def _check_batch(block_iter: Iterable[Block], ctx) -> Iterable[Block]:
     assert op.completed()
 
 
-@pytest.mark.parametrize("use_actors", [False, True])
-@pytest.mark.parametrize("preserve_order", [False, True])
-@pytest.mark.parametrize(
-    "target_max_block_size,num_expected_blocks", [(1, 10), (2**20, 1), (None, 1)]
-)
-def test_map_operator_output_unbundling(
+def _run_map_operator_test(
     ray_start_regular_shared,
     use_actors,
     preserve_order,
-    target_max_block_size,
-    num_expected_blocks,
+    transform_fn,
+    output_block_size_option,
+    expected_blocks,
+    test_name="TestMapper",
 ):
-    # Tests that the MapOperator's output queue unbundles the bundles returned from
-    # tasks; this facilitates features such as dynamic block splitting.
-    def noop(block_iter: Iterable[Block], ctx) -> Iterable[Block]:
-        for block in block_iter:
-            yield block
-
+    """Shared test function for MapOperator output unbundling tests."""
     # Create with inputs.
     input_op = InputDataBuffer(
         DataContext.get_current(), make_ref_bundles([[i] for i in range(10)])
     )
     compute_strategy = ActorPoolStrategy() if use_actors else TaskPoolStrategy()
 
     transformer = create_map_transformer_from_block_fn(
-        noop,
-        output_block_size_option=OutputBlockSizeOption.of(
-            target_max_block_size=target_max_block_size,
-        ),
+        transform_fn,
+        output_block_size_option=output_block_size_option,
     )
 
     op = MapOperator.create(
         transformer,
         input_op=input_op,
         data_context=DataContext.get_current(),
-        name="TestMapper",
+        name=test_name,
         compute_strategy=compute_strategy,
-        # Send the everything in a single bundle of 10 blocks.
+        # Send everything in a single bundle of 10 blocks.
         min_rows_per_bundle=10,
     )
 
@@ -670,10 +660,125 @@ def noop(block_iter: Iterable[Block], ctx) -> Iterable[Block]:
     outputs = []
     while op.has_next():
         outputs.append(op.get_next())
-    assert len(outputs) == num_expected_blocks
+    assert len(outputs) == expected_blocks
     assert op.completed()
 
 
+@pytest.mark.parametrize("use_actors", [False, True])
+@pytest.mark.parametrize("preserve_order", [False, True])
+@pytest.mark.parametrize(
+    "target_max_block_size,num_expected_blocks", [(1, 10), (2**20, 1), (None, 1)]
+)
+def test_map_operator_output_unbundling(
+    ray_start_regular_shared,
+    use_actors,
+    preserve_order,
+    target_max_block_size,
+    num_expected_blocks,
+):
+    """Test that MapOperator's output queue unbundles bundles from tasks."""
+
+    def noop(block_iter: Iterable[Block], ctx) -> Iterable[Block]:
+        for block in block_iter:
+            yield block
+
+    _run_map_operator_test(
+        ray_start_regular_shared,
+        use_actors,
+        preserve_order,
+        noop,
+        OutputBlockSizeOption.of(target_max_block_size=target_max_block_size),
+        num_expected_blocks,
+    )
+
+
+@pytest.mark.parametrize("preserve_order", [False, True])
+@pytest.mark.parametrize(
+    "output_block_size_option,expected_blocks",
+    [
+        # Test target_max_block_size
+        (OutputBlockSizeOption.of(target_max_block_size=1), 10),
+        (OutputBlockSizeOption.of(target_max_block_size=2**20), 1),
+        (OutputBlockSizeOption.of(target_max_block_size=None), 1),
+        # Test target_num_rows_per_block
+        (OutputBlockSizeOption.of(target_num_rows_per_block=1), 10),
+        (OutputBlockSizeOption.of(target_num_rows_per_block=5), 2),
+        (OutputBlockSizeOption.of(target_num_rows_per_block=10), 1),
+        (OutputBlockSizeOption.of(target_num_rows_per_block=None), 1),
+        # Test disable_block_shaping
+        (OutputBlockSizeOption.of(disable_block_shaping=True), 10),
+        (OutputBlockSizeOption.of(disable_block_shaping=False), 1),
+        # Test combinations
+        (
+            OutputBlockSizeOption.of(
+                target_max_block_size=1, target_num_rows_per_block=5
+            ),
+            10,
+        ),
+        (
+            OutputBlockSizeOption.of(
+                target_max_block_size=2**20, disable_block_shaping=True
+            ),
+            10,
+        ),
+        (
+            OutputBlockSizeOption.of(
+                target_num_rows_per_block=5, disable_block_shaping=True
+            ),
+            10,
+        ),
+    ],
+)
+def test_map_operator_output_block_size_options(
+    ray_start_regular_shared,
+    preserve_order,
+    output_block_size_option,
+    expected_blocks,
+):
+    """Test MapOperator with various OutputBlockSizeOption configurations."""
+
+    def noop(block_iter: Iterable[Block], ctx) -> Iterable[Block]:
+        for block in block_iter:
+            yield block
+
+    _run_map_operator_test(
+        ray_start_regular_shared,
+        use_actors=False,
+        preserve_order=preserve_order,
+        transform_fn=noop,
+        output_block_size_option=output_block_size_option,
+        expected_blocks=expected_blocks,
+    )
+
+
+@pytest.mark.parametrize("preserve_order", [False, True])
+def test_map_operator_disable_block_shaping_with_batches(
+    ray_start_regular_shared,
+    preserve_order,
+):
+    """Test MapOperator with disable_block_shaping=True using batch operations."""
+
+    def batch_transform(batch_iter, ctx):
+        for batch in batch_iter:
+            # Simple transformation: add 1 to each value
+            if hasattr(batch, "to_pandas"):
+                df = batch.to_pandas()
+                df = df + 1
+                yield df
+            else:
+                yield batch
+
+    _run_map_operator_test(
+        ray_start_regular_shared,
+        use_actors=False,
+        preserve_order=preserve_order,
+        transform_fn=batch_transform,
+        output_block_size_option=OutputBlockSizeOption.of(disable_block_shaping=True),
+        expected_blocks=10,  # With disable_block_shaping=True, we expect 10 blocks
+        test_name="TestBatchMapper",
+    )
+
+
 @pytest.mark.parametrize("use_actors", [False, True])
 def test_map_operator_ray_args(shutdown_only, use_actors):
     ray.shutdown()