[python][torch] Support distributed iterable dataset sharding

docs/docs/pypaimon/pytorch.md

@@ -60,6 +60,22 @@ when it is false, it will read the full amount of data into memory.
 **`prefetch_concurrency`** (default: 1): In streaming row mode, controls
 reader threads per DataLoader worker. It has no effect in non-streaming mode.
 
+### Distributed Sharding
+
+Streaming reads shard splits across DDP ranks and DataLoader workers:
+
+```python
+dataset = table_read.to_torch(
+    splits,
+    streaming=True,
+)
+dataloader = DataLoader(dataset, batch_size=32, num_workers=2)
+```
+
+PyPaimon checks `torch.distributed`, then `RANK` and `WORLD_SIZE`. Detection is
+enabled by default; set `auto_detect_rank=False` to disable rank sharding.
+A limit that may truncate the input is rejected when multiple ranks are active.
+
 ### Batch Streaming
 
 For batch-oriented training, make the streaming dataset yield batches directly:

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

@@ -18,6 +18,7 @@
 """
 Module to read a Paimon table into PyTorch Dataset.
 """
+import os
 import queue
 import random
 import threading
@@ -40,6 +41,59 @@ def _share_epoch_with_torch_workers(value):
     return torch.tensor(value, dtype=torch.long).share_memory_()
 
 
+def _validate_distributed_context(rank: int, world_size: int):
+    if isinstance(rank, bool) or not isinstance(rank, int):
+        raise ValueError("rank must be an int")
+    if isinstance(world_size, bool) or not isinstance(world_size, int):
+        raise ValueError("world_size must be an int")
+    if world_size <= 0:
+        raise ValueError("world_size must be greater than 0")
+    if rank < 0 or rank >= world_size:
+        raise ValueError("rank must satisfy 0 <= rank < world_size")
+    return rank, world_size
+
+
+def _resolve_distributed_context(auto_detect_rank: bool):
+    if not isinstance(auto_detect_rank, bool):
+        raise ValueError("auto_detect_rank must be a bool")
+    if not auto_detect_rank:
+        return 0, 1
+
+    distributed = getattr(torch, "distributed", None)
+    if (
+        distributed is not None
+        and distributed.is_available()
+        and distributed.is_initialized()
+    ):
+        rank = distributed.get_rank()
+        world_size = distributed.get_world_size()
+        return _validate_distributed_context(rank, world_size)
+
+    env_rank = os.environ.get("RANK")
+    env_world_size = os.environ.get("WORLD_SIZE")
+    if env_rank is not None or env_world_size is not None:
+        if env_rank is None or env_world_size is None:
+            raise ValueError(
+                "RANK and WORLD_SIZE environment variables must be set together"
+            )
+        try:
+            rank, world_size = int(env_rank), int(env_world_size)
+        except ValueError:
+            raise ValueError(
+                "RANK and WORLD_SIZE environment variables must be integers"
+            )
+        return _validate_distributed_context(rank, world_size)
+
+    return 0, 1
+
+
+def _balanced_slice(values: List[Any], shard_id: int, shard_count: int):
+    base_size, remainder = divmod(len(values), shard_count)
+    start = shard_id * base_size + min(shard_id, remainder)
+    size = base_size + (1 if shard_id < remainder else 0)
+    return values[start:start + size]
+
+
 class TorchDataset(Dataset):
     """
     PyTorch Dataset implementation for reading Paimon table data.
@@ -92,10 +146,33 @@ class _BaseTorchIterDataset(IterableDataset):
     Shared helpers for streaming PyTorch datasets backed by Paimon splits.
     """
 
-    def __init__(self, table_read: TableRead, splits: List[Split]):
+    def __init__(
+        self,
+        table_read: TableRead,
+        splits: List[Split],
+        auto_detect_rank: bool = True,
+    ):
         self.table_read = table_read
         self.splits = splits
         self.field_names = [field.name for field in table_read.read_type]
+        self.auto_detect_rank = auto_detect_rank
+        self.rank, self.world_size = _resolve_distributed_context(auto_detect_rank)
+        self._context_pid = os.getpid()
+
+    def _distributed_context(self):
+        rank, world_size = _resolve_distributed_context(
+            self.auto_detect_rank
+        )
+        current_pid = os.getpid()
+        if (
+            current_pid != self._context_pid
+            and world_size == 1
+            and self.world_size > 1
+        ):
+            return self.rank, self.world_size
+        self.rank, self.world_size = rank, world_size
+        self._context_pid = current_pid
+        return rank, world_size
 
     def _row_to_dict(self, offset_row) -> dict:
         row_dict = {}
@@ -136,30 +213,25 @@ def _limit_covers_all_splits(self) -> bool:
         return True
 
     def _worker_splits(self, worker_info) -> List[Split]:
-        if worker_info is None:
-            return self.splits
+        rank, world_size = self._distributed_context()
+        worker_id = worker_info.id if worker_info is not None else 0
+        num_workers = worker_info.num_workers if worker_info is not None else 1
 
-        # DataLoader workers cannot share a limit budget that may truncate.
+        if self.table_read.limit == 0:
+            return []
         if (
             self.table_read.limit is not None
             and not self._limit_covers_all_splits()
         ):
-            return self.splits if worker_info.id == 0 else []
-
-        worker_id = worker_info.id
-        num_workers = worker_info.num_workers
-        total_splits = len(self.splits)
-        splits_per_worker = total_splits // num_workers
-        remainder = total_splits % num_workers
-
-        if worker_id < remainder:
-            start_idx = worker_id * (splits_per_worker + 1)
-            end_idx = start_idx + splits_per_worker + 1
-        else:
-            start_idx = worker_id * splits_per_worker + remainder
-            end_idx = start_idx + splits_per_worker
+            if world_size > 1:
+                raise ValueError(
+                    "limit is not supported with distributed Torch sharding"
+                )
+            # A binding limit cannot be shared safely.
+            return self.splits if worker_id == 0 else []
 
-        return self.splits[start_idx:end_idx]
+        rank_splits = _balanced_slice(self.splits, rank, world_size)
+        return _balanced_slice(rank_splits, worker_id, num_workers)
 
 
 class TorchIterDataset(_BaseTorchIterDataset):
@@ -179,7 +251,13 @@ class TorchIterDataset(_BaseTorchIterDataset):
     _PREFETCH_GET_TIMEOUT_SEC = 300.0
     _PREFETCH_JOIN_TIMEOUT_SEC = 5.0
 
-    def __init__(self, table_read: TableRead, splits: List[Split], prefetch_concurrency: int = 1):
+    def __init__(
+        self,
+        table_read: TableRead,
+        splits: List[Split],
+        prefetch_concurrency: int = 1,
+        auto_detect_rank: bool = True,
+    ):
         """
         Initialize TorchIterDataset.
 
@@ -190,7 +268,7 @@ def __init__(self, table_read: TableRead, splits: List[Split], prefetch_concurre
                 this worker (default 1). When > 1, splits are partitioned across
                 threads to increase read throughput.
         """
-        super().__init__(table_read, splits)
+        super().__init__(table_read, splits, auto_detect_rank)
         self.prefetch_concurrency = max(1, int(prefetch_concurrency))
 
     def __iter__(self):
@@ -393,8 +471,9 @@ def __init__(
         batch_format: str,
         batch_size: Optional[int],
         to_tensor_fn: Optional[Callable[[pa.RecordBatch], Any]] = None,
+        auto_detect_rank: bool = True,
     ):
-        super().__init__(table_read, splits)
+        super().__init__(table_read, splits, auto_detect_rank)
         self.batch_format = batch_format
         self.batch_size = batch_size
         self.to_tensor_fn = to_tensor_fn
@@ -457,8 +536,9 @@ def __init__(
         seed: int = 0,
         buffer_size: int = 1000,
         max_buffer_input_splits: int = 10,
+        auto_detect_rank: bool = True,
     ):
-        super().__init__(table_read, splits)
+        super().__init__(table_read, splits, auto_detect_rank)
         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(
@@ -559,7 +639,11 @@ def _iter_buffer_shuffled_rows(
         rows: Iterator[dict],
         worker_id: int,
     ) -> Iterator[dict]:
-        rng = random.Random(self.seed + self.epoch * 1000003 + 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)
+        rng = random.Random(rng_seed)
         buffer = []
         for row in rows:
             if len(buffer) < self.buffer_size:

paimon-python/pypaimon/read/table_read.py

@@ -661,6 +661,7 @@ def to_torch(
         seed: int = 0,
         buffer_size: int = 1000,
         max_buffer_input_splits: int = 10,
+        auto_detect_rank: bool = True,
     ) -> "torch.utils.data.Dataset":
         """Wrap Paimon table data in a PyTorch Dataset.
 
@@ -674,6 +675,7 @@ def to_torch(
             batch_size: Rows per batch; ``None`` preserves reader batches.
             to_tensor_fn: Optional RecordBatch converter for Torch batches.
             shuffle: Whether to shuffle rows; supported only in row format.
+            auto_detect_rank: Whether streaming reads detect the DDP context.
         """
         valid_batch_formats = {"row", "pyarrow", "torch"}
         if batch_format not in valid_batch_formats:
@@ -725,6 +727,7 @@ def to_torch(
                 batch_format=batch_format,
                 batch_size=batch_size,
                 to_tensor_fn=to_tensor_fn,
+                auto_detect_rank=auto_detect_rank,
             )
 
         if shuffle:
@@ -739,12 +742,18 @@ def to_torch(
                 seed=seed,
                 buffer_size=buffer_size,
                 max_buffer_input_splits=max_buffer_input_splits,
+                auto_detect_rank=auto_detect_rank,
             )
             return dataset
 
         if streaming:
             from pypaimon.read.datasource.torch_dataset import TorchIterDataset
-            dataset = TorchIterDataset(self, splits, prefetch_concurrency)
+            dataset = TorchIterDataset(
+                self,
+                splits,
+                prefetch_concurrency,
+                auto_detect_rank=auto_detect_rank,
+            )
             return dataset
         else:
             from pypaimon.read.datasource.torch_dataset import TorchDataset

paimon-python/pypaimon/tests/torch_distributed_sharding_worker.py

@@ -0,0 +1,75 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import json
+import os
+import sys
+from types import SimpleNamespace
+
+import torch
+from torch.utils.data import DataLoader
+
+from pypaimon.read.datasource.torch_dataset import TorchIterDataset
+
+
+class _OffsetRow:
+    def __init__(self, values):
+        self._values = values
+
+    def get_field(self, index):
+        return self._values[index]
+
+
+class _TableRead:
+    limit = None
+    read_type = [
+        SimpleNamespace(name="split_id"),
+        SimpleNamespace(name="rank"),
+        SimpleNamespace(name="worker"),
+    ]
+
+    def to_iterator(self, splits):
+        worker_info = torch.utils.data.get_worker_info()
+        worker_id = worker_info.id if worker_info is not None else 0
+        rank = int(os.environ["RANK"])
+        for split_id in splits:
+            yield _OffsetRow([split_id, rank, worker_id])
+
+
+def main():
+    output_dir = sys.argv[1]
+    torch.distributed.init_process_group("gloo")
+    rank = torch.distributed.get_rank()
+    try:
+        dataset = TorchIterDataset(
+            _TableRead(),
+            list(range(11)),
+        )
+        rows = list(DataLoader(dataset, batch_size=None, num_workers=2))
+        with open(
+            os.path.join(output_dir, "rank-%d.json" % rank),
+            "w",
+            encoding="utf-8",
+        ) as result_file:
+            json.dump(rows, result_file)
+        torch.distributed.barrier()
+    finally:
+        torch.distributed.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()