chore: import upstream snapshot with attribution

python/ray/data/_internal/datasource/torch_datasource.py

@@ -0,0 +1,82 @@
+from typing import TYPE_CHECKING, Optional
+
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data.block import BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+
+if TYPE_CHECKING:
+    import torch
+
+    from ray.data.context import DataContext
+
+
+TORCH_DATASOURCE_READER_BATCH_SIZE = 32
+
+
+class TorchDatasource(Datasource):
+    """Torch datasource, for reading from `Torch
+    datasets <https://pytorch.org/docs/stable/data.html/>`_.
+    This datasource implements a streaming read using a single read task.
+    """
+
+    def __init__(
+        self,
+        dataset: "torch.utils.data.Dataset",
+    ):
+        self._dataset = dataset
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+        data_context: Optional["DataContext"] = None,
+    ):
+        assert parallelism == 1
+
+        meta = BlockMetadata(
+            # Note: avoid len(self._dataset) because it will trigger
+            # iterating through IterableDataset, which can cause OOM.
+            num_rows=None,
+            size_bytes=None,
+            input_files=None,
+            exec_stats=None,
+        )
+        read_task = ReadTask(
+            lambda subset=self._dataset: _read_subset(
+                subset,
+            ),
+            metadata=meta,
+            per_task_row_limit=per_task_row_limit,
+        )
+
+        return [read_task]
+
+    def estimate_inmemory_data_size(self):
+        return None
+
+
+def _read_subset(subset: "torch.utils.data.Subset"):
+    batch = []
+
+    # Get items from dataset based on its type
+    if hasattr(subset, "__iter__"):
+        # IterableDataset: Use the iterator directly
+        items = subset
+    else:
+        # Map-style dataset: Respect __len__
+        items = (subset[i] for i in range(len(subset)))
+
+    # Process items in batches
+    for item in items:
+        batch.append(item)
+        if len(batch) == TORCH_DATASOURCE_READER_BATCH_SIZE:
+            builder = DelegatingBlockBuilder()
+            builder.add_batch({"item": batch})
+            yield builder.build()
+            batch.clear()
+
+    # Handle any remaining items
+    if len(batch) > 0:
+        builder = DelegatingBlockBuilder()
+        builder.add_batch({"item": batch})
+        yield builder.build()