[python][torch] Support shuffle for Arrow and tensor batch streaming

docs/docs/pypaimon/pytorch.md

@@ -133,12 +133,58 @@ for batch in dataloader:
 converter supports non-null numeric, boolean, and numeric fixed-size-list
 columns. Use `to_tensor_fn` for other types or custom conversion.
 
-Omit `batch_size` to preserve native reader batches. Otherwise, batches are
-combined or sliced to the requested size. Use `DataLoader(batch_size=None)` to
-disable a second batching step. Batch streaming does not support `shuffle=True`.
+Without shuffle, omit `batch_size` to preserve native reader batches. Otherwise,
+batches are combined or sliced to the requested size. Use
+`DataLoader(batch_size=None)` to disable a second batching step.
 Numeric tensors may share read-only Arrow buffers; clone them before in-place
 mutation. Batch formats currently require `prefetch_concurrency=1`.
 
+### Shuffled Batch Streaming
+
+Set `shuffle=True` to mix rows across input batches while keeping values in
+Arrow until the final Tensor conversion. The same options work with
+`batch_format="pyarrow"` and custom `to_tensor_fn` converters:
+
+```python
+dataset = table_read.to_torch(
+    splits,
+    streaming=True,
+    batch_format="torch",
+    batch_size=256,
+    shuffle=True,
+    seed=42,
+    buffer_size=4096,
+    max_buffer_input_splits=4,
+)
+loader = DataLoader(dataset, batch_size=None, num_workers=2)
+
+for epoch in range(10):
+    dataset.set_epoch(epoch)
+    for batch in loader:
+        train(batch["features"], batch["label"])
+```
+
+Each worker retains at most `buffer_size` rows in a rolling shuffle buffer and
+replaces randomly selected slots with rows from incoming Arrow blocks. Input
+blocks contain at most `buffer_size` rows; gathering replacements, output
+batching, and each open format reader use additional memory. This is a row
+bound, not a byte bound. The buffer drains early if combining Arrow blocks
+would overflow a 32-bit offset. Without `batch_size`, output blocks contain at
+most `buffer_size` rows.
+
+`max_buffer_input_splits` bounds the number of interleaved split readers per
+worker; `1` reads splits in order. A binding limit keeps the existing ordered
+read path so it selects the same rows before shuffling. Filters and distributed
+worker/rank sharding also precede shuffle. Each selected row is emitted once.
+
+The shuffle is local to each worker's buffer, not a uniform permutation of the
+whole dataset. The same seed, epoch, input batches and worker/rank configuration
+reproduce the same order; `set_epoch()` also reaches persistent workers. Changing
+reader batching or worker configuration can change the order. Arrow and Tensor
+formats share the same shuffle path, while row-format split interleaving can
+produce a different order. The existing `prefetch_concurrency=1` requirement
+also applies to shuffled batches.
+
 ## Video frame descriptors
 
 For a multimodal frame table, use the higher-level scan API:

paimon-python/pypaimon/read/datasource/torch_dataset.py

@@ -25,6 +25,7 @@
 import random
 import threading
 import warnings
+from collections import deque
 from typing import Any, Callable, Iterator, List, Optional
 
 import pyarrow as pa
@@ -767,6 +768,52 @@ def _default_to_tensor(batch: pa.RecordBatch) -> dict:
     return tensors
 
 
+class _TorchShuffleMixin:
+
+    def _init_shuffle(self, seed, buffer_size, max_buffer_input_splits):
+        self.seed = self._require_int(seed, "seed")
+        self.buffer_size = self._require_positive_int(buffer_size, "buffer_size")
+        self.max_buffer_input_splits = self._require_positive_int(
+            max_buffer_input_splits, "max_buffer_input_splits")
+        self._epoch = _share_epoch_with_torch_workers(0)
+
+    def __setstate__(self, state):
+        self.__dict__ = state
+        self._epoch = _share_epoch_with_torch_workers(self._epoch)
+
+    @property
+    def epoch(self) -> int:
+        return int(self._epoch)
+
+    @epoch.setter
+    def epoch(self, epoch: int) -> None:
+        epoch = self._require_int(epoch, "epoch")
+        self._epoch += epoch - self._epoch
+
+    @staticmethod
+    def _require_int(value: int, name: str) -> int:
+        if not isinstance(value, int):
+            raise ValueError("%s must be an int" % name)
+        return value
+
+    @staticmethod
+    def _require_positive_int(value: int, name: str) -> int:
+        if not isinstance(value, int) or value <= 0:
+            raise ValueError("%s must be a positive int" % name)
+        return value
+
+    def set_epoch(self, epoch: int):
+        self.epoch = epoch
+        return self
+
+    def _shuffle_rng(self, worker_id):
+        rank, world_size = self._distributed_context()
+        rng_seed = self.seed + self.epoch * 1000003 + worker_id
+        if world_size > 1:
+            rng_seed = "%d:%d" % (rng_seed, rank)
+        return random.Random(rng_seed)
+
+
 class TorchBatchIterDataset(_BaseTorchIterDataset):
     """Streaming IterableDataset which yields Arrow or Tensor batches."""
 
@@ -796,16 +843,20 @@ def __iter__(self):
         worker_info = torch.utils.data.get_worker_info()
         splits_to_process = self._worker_splits(worker_info)
         raw_batches = self._arrow_batches_for_splits(splits_to_process)
+        worker_id = worker_info.id if worker_info is not None else 0
+        try:
+            for batch in self._prepare_batches(raw_batches, worker_id):
+                if self.batch_format == "torch":
+                    converter = self.to_tensor_fn or _default_to_tensor
+                    yield converter(batch)
+                else:
+                    yield batch
+        finally:
+            raw_batches.close()
 
-        batches = _sized_record_batches(
-            self._limit_batches(raw_batches), self.batch_size
-        )
-        for batch in batches:
-            if self.batch_format == "torch":
-                converter = self.to_tensor_fn or _default_to_tensor
-                yield converter(batch)
-            else:
-                yield batch
+    def _prepare_batches(self, raw_batches, worker_id):
+        return _sized_record_batches(
+            self._limit_batches(raw_batches), self.batch_size)
 
     def _arrow_batches_for_splits(
         self, splits: List[Split]
@@ -834,7 +885,87 @@ def _limit_batches(
             yield batch
 
 
-class TorchShuffledIterDataset(_BaseTorchIterDataset):
+class TorchShuffledBatchIterDataset(_TorchShuffleMixin, TorchBatchIterDataset):
+    """Shuffle Arrow rows before sizing batches and converting to tensors."""
+
+    def __init__(self, table_read, splits, seed=0, buffer_size=1000,
+                 max_buffer_input_splits=10, **kwargs):
+        super().__init__(table_read, splits, **kwargs)
+        self._init_shuffle(seed, buffer_size, max_buffer_input_splits)
+
+    def _prepare_batches(self, raw_batches, worker_id):
+        shuffled = _shuffle_record_batches(
+            self._limit_batches(raw_batches), self.buffer_size,
+            self._shuffle_rng(worker_id))
+        return _sized_record_batches(shuffled, self.batch_size)
+
+    def _arrow_batches_for_splits(self, splits):
+        # Keep the selected prefix of a binding limit unchanged before shuffle.
+        if self.max_buffer_input_splits == 1 or not self._limit_covers_all_splits():
+            yield from super()._arrow_batches_for_splits(splits)
+            return
+
+        split_iter = iter(splits)
+        active = deque()
+        read_batches = super()._arrow_batches_for_splits
+
+        def add_next():
+            split = next(split_iter, None)
+            if split is not None:
+                active.append(read_batches([split]))
+
+        try:
+            for _ in range(min(self.max_buffer_input_splits, len(splits))):
+                add_next()
+            while active:
+                try:
+                    yield next(active[0])
+                except StopIteration:
+                    active.popleft().close()
+                    add_next()
+                else:
+                    active.rotate(-1)
+        finally:
+            for batches in active:
+                batches.close()
+
+
+def _shuffle_record_batches(batches, buffer_size, rng):
+    # Only row positions enter Python. Payloads, including nested values and
+    # nulls, stay in Arrow. Each input block replaces random reservoir slots.
+    buffer = None
+    for batch in _sized_record_batches(batches, buffer_size):
+        if buffer is None:
+            buffer = batch
+            continue
+        usage = _batch_offset_usage(buffer)
+        if any(usage.get(path, 0) + value > _MAX_ARROW_OFFSET
+               for path, value in _batch_offset_usage(batch).items()):
+            # Drain before concatenation would overflow a 32-bit Arrow offset.
+            yield _shuffled_record_batch(buffer, rng)
+            buffer = batch
+            continue
+        combined = _concat_record_batches([buffer, batch])
+        retained = list(range(buffer.num_rows))
+        selected = []
+        for row in range(batch.num_rows):
+            slot = rng.randrange(buffer.num_rows)
+            selected.append(retained[slot])
+            retained[slot] = buffer.num_rows + row
+        output = combined.take(pa.array(selected, type=pa.int64()))
+        buffer = combined.take(pa.array(retained, type=pa.int64()))
+        yield output
+    if buffer is not None:
+        yield _shuffled_record_batch(buffer, rng)
+
+
+def _shuffled_record_batch(batch, rng):
+    order = list(range(batch.num_rows))
+    rng.shuffle(order)
+    return batch.take(pa.array(order, type=pa.int64()))
+
+
+class TorchShuffledIterDataset(_TorchShuffleMixin, _BaseTorchIterDataset):
     """
     PyTorch IterableDataset with Paimon-controlled streaming shuffle.
 
@@ -861,40 +992,7 @@ def __init__(
             sharding_rank,
             sharding_world_size,
         )
-        self.seed = self._require_int(seed, "seed")
-        self.buffer_size = self._require_positive_int(buffer_size, "buffer_size")
-        self.max_buffer_input_splits = self._require_positive_int(
-            max_buffer_input_splits, "max_buffer_input_splits")
-        self._epoch = _share_epoch_with_torch_workers(0)
-
-    def __setstate__(self, state):
-        self.__dict__ = state
-        self._epoch = _share_epoch_with_torch_workers(self._epoch)
-
-    @property
-    def epoch(self) -> int:
-        return int(self._epoch)
-
-    @epoch.setter
-    def epoch(self, epoch: int) -> None:
-        epoch = self._require_int(epoch, "epoch")
-        self._epoch += epoch - self._epoch
-
-    @staticmethod
-    def _require_int(value: int, name: str) -> int:
-        if not isinstance(value, int):
-            raise ValueError("%s must be an int" % name)
-        return value
-
-    @staticmethod
-    def _require_positive_int(value: int, name: str) -> int:
-        if not isinstance(value, int) or value <= 0:
-            raise ValueError("%s must be a positive int" % name)
-        return value
-
-    def set_epoch(self, epoch: int) -> "TorchShuffledIterDataset":
-        self.epoch = epoch
-        return self
+        self._init_shuffle(seed, buffer_size, max_buffer_input_splits)
 
     def __iter__(self):
         worker_info = torch.utils.data.get_worker_info()
@@ -961,11 +1059,7 @@ def _iter_buffer_shuffled_rows(
         rows: Iterator[dict],
         worker_id: int,
     ) -> Iterator[dict]:
-        rank, world_size = self._distributed_context()
-        rng_seed = self.seed + self.epoch * 1000003 + worker_id
-        if world_size > 1:
-            rng_seed = "%d:%d" % (rng_seed, rank)
-        rng = random.Random(rng_seed)
+        rng = self._shuffle_rng(worker_id)
         buffer = []
         for row in rows:
             if len(buffer) < self.buffer_size:

paimon-python/pypaimon/read/table_read.py

@@ -1207,9 +1207,10 @@ def to_torch(
                 format.
             batch_format: ``"row"``, ``"pyarrow"``, or ``"torch"``. Batch
                 formats require streaming.
-            batch_size: Rows per batch; ``None`` preserves reader batches.
+            batch_size: Rows per batch; ``None`` preserves reader batches when
+                unshuffled, or emits shuffle blocks of at most buffer_size rows.
             to_tensor_fn: Optional RecordBatch converter for Torch batches.
-            shuffle: Whether to shuffle rows; supported only in row format.
+            shuffle: Whether to shuffle rows in a bounded buffer before output.
             auto_detect_rank: Whether streaming reads shard by DDP rank.
             sharding_rank: Explicit rank in the intended DDP process group.
             sharding_world_size: Explicit size of that process group.
@@ -1244,10 +1245,6 @@ def to_torch(
                 raise ValueError(
                     "batch_format=%r requires streaming=True" % batch_format
                 )
-            if shuffle:
-                raise ValueError(
-                    "shuffle=True only supports batch_format='row'"
-                )
             if batch_format == "pyarrow" and to_tensor_fn is not None:
                 raise ValueError("to_tensor_fn requires batch_format='torch'")
             if to_tensor_fn is not None and not callable(to_tensor_fn):
@@ -1263,8 +1260,13 @@ def to_torch(
 
             from pypaimon.read.datasource.torch_dataset import (
                 TorchBatchIterDataset,
+                TorchShuffledBatchIterDataset,
             )
-            return TorchBatchIterDataset(
+            dataset_type = TorchShuffledBatchIterDataset if shuffle else TorchBatchIterDataset
+            shuffle_options = (dict(seed=seed, buffer_size=buffer_size,
+                                    max_buffer_input_splits=max_buffer_input_splits)
+                               if shuffle else {})
+            return dataset_type(
                 self,
                 splits,
                 batch_format=batch_format,
@@ -1273,6 +1275,7 @@ def to_torch(
                 auto_detect_rank=auto_detect_rank,
                 sharding_rank=sharding_rank,
                 sharding_world_size=sharding_world_size,
+                **shuffle_options,
             )
 
         if shuffle: