[Data] AlltoAll OP, Update Data progress bars to use row as the iteration unit

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

@@ -65,6 +65,7 @@ def __init__(
         self._bulk_fn = bulk_fn
         self._next_task_index = 0
         self._num_outputs = num_outputs
+        self._output_rows = 0
         self._sub_progress_bar_names = sub_progress_bar_names
         self._sub_progress_bar_dict = None
         self._input_buffer: List[RefBundle] = []
@@ -79,6 +80,13 @@ def num_outputs_total(self) -> Optional[int]:
             else self.input_dependencies[0].num_outputs_total()
         )
 
+    def num_output_rows_total(self) -> Optional[int]:
+        return (
+            self._output_rows
+            if self._output_rows
+            else self.input_dependencies[0].num_output_rows_total()
+        )
+
     def _add_input_inner(self, refs: RefBundle, input_index: int) -> None:
         assert not self.completed()
         assert input_index == 0, input_index
@@ -99,7 +107,9 @@ def has_next(self) -> bool:
         return len(self._output_buffer) > 0
 
     def _get_next_inner(self) -> RefBundle:
-        return self._output_buffer.pop(0)
+        bundle = self._output_buffer.pop(0)
+        self._output_rows += bundle.num_rows()
+        return bundle
 
     def get_stats(self) -> StatsDict:
         return self._stats
@@ -108,7 +118,7 @@ def get_transformation_fn(self) -> AllToAllTransformFn:
         return self._bulk_fn
 
     def progress_str(self) -> str:
-        return f"{len(self._output_buffer)} output"
+        return f"{self.num_output_rows_total() or 0} rows output"
 
     def initialize_sub_progress_bars(self, position: int) -> int:
         """Initialize all internal sub progress bars, and return the number of bars."""
@@ -117,8 +127,8 @@ def initialize_sub_progress_bars(self, position: int) -> int:
             for name in self._sub_progress_bar_names:
                 bar = ProgressBar(
                     name,
-                    self.num_outputs_total() or 1,
-                    unit="bundle",
+                    self.num_output_rows_total() or 1,
+                    unit="row",
                     position=position,
                 )
                 # NOTE: call `set_description` to trigger the initial print of progress

python/ray/data/_internal/execution/streaming_executor.py

@@ -154,7 +154,9 @@ def get_next(self, output_split_idx: Optional[int] = None) -> RefBundle:
                         output_split_idx
                     )
                     if self._outer._global_info:
-                        self._outer._global_info.update(1, dag.num_outputs_total())
+                        self._outer._global_info.update(
+                            item.num_rows(), dag.num_output_rows_total()
+                        )
                     return item
                 # Needs to be BaseException to catch KeyboardInterrupt. Otherwise we
                 # can leave dangling progress bars by skipping shutdown.

python/ray/data/_internal/progress_bar.py

@@ -134,7 +134,13 @@ def block_until_complete(self, remaining: List[ObjectRef]) -> None:
             done, remaining = ray.wait(
                 remaining, num_returns=len(remaining), fetch_local=False, timeout=0.1
             )
-            self.update(len(done))
+            total_rows_processed = 0
+            for _, result in zip(done, ray.get(done)):
+                num_rows = (
+                    result.num_rows if hasattr(result, "num_rows") else 1
+                )  # Default to 1 if no row count is available
+                total_rows_processed += num_rows
+            self.update(total_rows_processed)
 
             with _canceled_threads_lock:
                 if t in _canceled_threads:
@@ -158,9 +164,15 @@ def fetch_until_complete(self, refs: List[ObjectRef]) -> List[Any]:
             )
             if fetch_local:
                 fetch_local = False
+            total_rows_processed = 0
             for ref, result in zip(done, ray.get(done)):
                 ref_to_result[ref] = result
-            self.update(len(done))
+                num_rows = (
+                    result.num_rows if hasattr(result, "num_rows") else 1
+                )  # Default to 1 if no row count is available
+                total_rows_processed += num_rows
+            # TODO(zhilong): Change the total to total_row when init progress bar
+            self.update(total_rows_processed)
 
             with _canceled_threads_lock:
                 if t in _canceled_threads: