[python][torch] Support distributed iterable dataset sharding

[XiaoHongbo-Hope]

What changed

• Shard streaming Torch datasets across DDP ranks and DataLoader workers when auto_detect_rank=True.
• Keep rank sharding disabled by default because to_torch accepts caller-planned splits.
• Resolve rank from torch.distributed, then torchrun environment variables, and preserve it in spawned workers.
• Reject binding limits with multiple ranks.
• Keep shuffled reads reproducible per epoch, rank, and worker.
• Document uneven inputs and DistributedDataParallel.join().

Reference

The API is inspired by LanceDataset. Unlike Lance, this low-level API accepts external splits, so automatic sharding is opt-in.

Validation

• 43 Torch tests passed.
• CPU torchrun test executes DDP forward/backward with 2 ranks and 2 spawned DataLoader workers.
• Split assignment is complete and non-overlapping.
• Pre-sharded splits remain unchanged by default.
• flake8 and git diff --check passed.

[JingsongLi]

Found two issues in the distributed Torch sharding path; please see the inline comments.

[JingsongLi]

I compared this design with the current PyTorch, Lance, Hugging Face Datasets, WebDataset, and Ray Data implementations. Overall, the direction in this PR looks sound, and I do not see an architectural blocker.

The main reasons are:

• A plain PyTorch IterableDataset must own worker-replica sharding itself (get_worker_info()); DDP does not shard inputs, and DataLoader rejects a custom sampler for an IterableDataset. Therefore, doing process-rank sharding in the dataset and then DataLoader-worker sharding is the correct layer. See the PyTorch DataLoader source and DDP documentation.
• Hugging Face Datasets and WebDataset use the same broad hierarchy: first select the rank/node-local sources, then split those sources among DataLoader workers before opening/decoding them. See HF distributed.py, HF iterable_dataset.py, and WebDataset shardlists.py.
• Keeping auto_detect_rank=False by default is appropriate for Paimon. TableRead.to_torch receives an already-planned split list and cannot know whether the caller has pre-sharded it. Automatically sharding by default could therefore silently apply sharding twice. Systems such as Ray can safely auto-split because the trainer owns the complete dataset plan and all consumers; that is a different ownership boundary. See Ray's dataset-shard API.
• The explicit sharding_rank / sharding_world_size pair is important for DDP subgroups and hybrid parallelism. Resolving these values in the parent and serializing scalar context into spawned workers is also a good design.

One useful comparison is Lance. Lance flattens distributed rank and DataLoader worker id into a global rank and supports fragment- and batch-level sharding. However, based on the current source, its distributed helper only observes a process group initialized in the current process. A spawned DataLoader worker does not retain the parent's initialized process group, while the deprecated explicit rank/world_size branch does not additionally shard among worker replicas. Therefore, as a source-level inference, Lance's iterable path does not transparently cover the full spawn + DDP + num_workers > 0 combination. Paimon's captured parent rank followed by a separate worker split handles this combination more robustly. See LanceDataset, Lance distributed helpers, and Lance samplers.

I would keep the current exact-coverage policy as well: assign every split once, allow uneven rank lengths, and require DDP uneven-input handling. This is preferable to silently duplicating or dropping data. A few non-blocking follow-ups may still be worthwhile:

1. Clarify that set_epoch() changes each rank's local buffer-shuffle order, but does not reshuffle split ownership across ranks as PyTorch DistributedSampler does.
2. Mention that training loops with additional non-DDP collectives (for example, SyncBatchNorm) should use model.join(throw_on_early_termination=True); plain join() can only shadow DDP's own collectives.
3. Consider warning when the number of splits is smaller than world_size * num_workers, since some ranks/workers will be idle.
4. If equal-step training becomes a requirement later, add an explicit policy such as preserve-all/drop-to-equal/repeat-to-equal rather than changing the current behavior implicitly. Similarly, fragment/split versus batch/weighted sharding can be added as an explicit strategy later without coupling it to rank detection.

So my recommendation is to keep the current API boundary and rank -> worker -> reader hierarchy. The comparison with other implementations reinforces this design rather than suggesting that Paimon should copy Lance's defaults.

[XiaoHongbo-Hope]

So my recommendation is to keep the current API boundary and rank -> worker -> reader hierarchy. The comparison with other implementations reinforces this design rather than suggesting that Paimon should copy Lance's defaults.

Agreed, I also went back and forth on Lance’s sharding model before settling on this design.

[JingsongLi]

Reviewed the final head 05cbebc. The previously identified spawn-safety and DDP subgroup sharding issues are fixed, all required checks pass, and the cross-implementation comparison found no remaining architectural blocker.