Repository: Project-MONAI/MONAI

Issue:
CSVIterableDataset doesn't work in the same way as CSVDataset
**Is your feature request related to a problem? Please describe.**
Here is the feedback from @holgerroth :

When I am using `CSVIterableDataset`, I get this error
```
Traceback (most recent call last):
  File "run_classification_criteo_monai_mlp.py", line 227, in <module>
    main()
  File "run_classification_criteo_monai_mlp.py", line 143, in main
    f"{step}/{len(train_ds) // train_loader.batch_size}, "
TypeError: object of type 'CSVIterableDataset' has no len()
```
Also, trying to set shuffle=True for the training dataset gives
```
Traceback (most recent call last):
  File "run_classification_criteo_monai_mlp.py", line 227, in <module>
    main()
  File "run_classification_criteo_monai_mlp.py", line 87, in main
    train_loader = torch.utils.data.DataLoader(
  File "/home/hroth/.local/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 224, in __init__
    raise ValueError(
ValueError: DataLoader with IterableDataset: expected unspecified shuffle option, but got shuffle=True
```


Discussion:
Thanks for @holgerroth 's feedback, let me try to enhance the doc-string of `CSVIterableDataset` to make the usage much more clear. And may also enhance the `CSVIterableDataset` with more features.

Thanks.


Patch:
diff --git a/monai/data/__init__.py b/monai/data/__init__.py
--- a/monai/data/__init__.py
+++ b/monai/data/__init__.py
@@ -28,7 +28,7 @@
 from .grid_dataset import GridPatchDataset, PatchDataset, PatchIter
 from .image_dataset import ImageDataset
 from .image_reader import ImageReader, ITKReader, NibabelReader, NumpyReader, PILReader, WSIReader
-from .iterable_dataset import CSVIterableDataset, IterableDataset
+from .iterable_dataset import CSVIterableDataset, IterableBuffer, IterableDataset
 from .nifti_saver import NiftiSaver
 from .nifti_writer import write_nifti
 from .png_saver import PNGSaver
diff --git a/monai/data/iterable_dataset.py b/monai/data/iterable_dataset.py
--- a/monai/data/iterable_dataset.py
+++ b/monai/data/iterable_dataset.py
@@ -9,14 +9,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import math
-from typing import Any, Callable, Dict, Iterable, Optional, Sequence, Union
+from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Union
 
 from torch.utils.data import IterableDataset as _TorchIterableDataset
 from torch.utils.data import get_worker_info
 
 from monai.data.utils import convert_tables_to_dicts
 from monai.transforms import apply_transform
+from monai.transforms.transform import Randomizable
 from monai.utils import ensure_tuple, optional_import
 
 pd, _ = optional_import("pandas")
@@ -25,7 +25,8 @@
 class IterableDataset(_TorchIterableDataset):
     """
     A generic dataset for iterable data source and an optional callable data transform
-    when fetching a data sample.
+    when fetching a data sample. Inherit from PyTorch IterableDataset:
+    https://pytorch.org/docs/stable/data.html?highlight=iterabledataset#torch.utils.data.IterableDataset.
     For example, typical input data can be web data stream which can support multi-process access.
 
     Note that when used with `DataLoader` and `num_workers > 0`, each worker process will have a
@@ -51,10 +52,70 @@ def __iter__(self):
             yield data
 
 
-class CSVIterableDataset(IterableDataset):
+class IterableBuffer(Randomizable, IterableDataset):
+    """
+    Extend the IterableDataset with a buffer and support to randomly pop items.
+
+    Args:
+        data: input data source to load and transform to generate dataset for model.
+        transform: a callable data transform on input data.
+        buffer_size: size of the buffer to store items and randomly pop, default to 512.
+        shuffle: if True, randomly pop a item from the list, otherwise, pop from the beginning,
+            default to False.
+
+    """
+
+    def __init__(self, data, transform=None, buffer_size: int = 512, shuffle: bool = True) -> None:
+        super().__init__(data=data, transform=transform)
+        self.size = buffer_size
+        self.shuffle = shuffle
+        self._idx = 0
+
+    def _rand_pop(self, buffer: List, num_workers: int = 1, id: int = 0):
+        length = len(buffer)
+        for i in range(min(length, num_workers)):
+            if self.shuffle:
+                # randomly select an item for every worker and pop
+                self.randomize(length)
+            # switch random index data and the last index data
+            item, buffer[self._idx] = buffer[self._idx], buffer[-1]
+            buffer.pop()
+            if i == id:
+                if self.transform is not None:
+                    item = apply_transform(self.transform, item)
+                yield item
+
+    def __iter__(self):
+        # pop items for multi-workers
+        info = get_worker_info()
+        num_workers = info.num_workers if info is not None else 1
+        id = info.id if info is not None else 0
+
+        _buffer = []
+        for item in self.data:
+            if len(_buffer) >= self.size:
+                self._rand_pop(_buffer, num_workers=num_workers, id=id)
+            _buffer.append(item)
+        while _buffer:
+            return self._rand_pop(_buffer, num_workers=num_workers, id=id)
+
+    def randomize(self, size: int) -> None:
+        self._idx = self.R.randint(size)
+
+
+class CSVIterableDataset(IterableDataset, Randomizable):
     """
     Iterable dataset to load CSV files and generate dictionary data.
-    It can be helpful when loading extremely big CSV files that can't read into memory directly.
+    It is particularly useful when data come from a stream, inherits from PyTorch IterableDataset:
+    https://pytorch.org/docs/stable/data.html?highlight=iterabledataset#torch.utils.data.IterableDataset.
+
+    It also can be helpful when loading extremely big CSV files that can't read into memory directly,
+    just treat the big CSV file as stream input, call `reset()` of `CSVIterableDataset` for every epoch.
+    Note that as a stream input, it can't get the length of dataset.
+
+    To effectively shuffle the data in the big dataset, users can set a big buffer to continuously store
+    the loaded chunks, then randomly pick data from the buffer for following tasks.
+
     To accelerate the loading process, it can support multi-processing based on PyTorch DataLoader workers,
     every process executes transforms on part of every loaded chunk.
     Note: the order of output data may not match data source in multi-processing mode.
@@ -70,10 +131,11 @@ class CSVIterableDataset(IterableDataset):
         ]
 
     Args:
-        filename: the filename of expected CSV file to load. if providing a list
-            of filenames, it will load all the files and join tables.
+        filename: the filename of CSV file to load. it can be a str, URL, path object or file-like object.
+            if providing a list of filenames, it will load all the files and join tables.
         chunksize: rows of a chunk when loading iterable data from CSV files, default to 1000. more details:
             https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.read_csv.html.
+        buffer_size: size of the buffer to store the loaded chunks, if None, set to `2 x chunksize`.
         col_names: names of the expected columns to load. if None, load all the columns.
         col_types: `type` and `default value` to convert the loaded columns, if None, use original data.
             it should be a dictionary, every item maps to an expected column, the `key` is the column
@@ -93,6 +155,7 @@ class CSVIterableDataset(IterableDataset):
             be the new column name, the `value` is the names of columns to combine. for example:
             `col_groups={"ehr": [f"ehr_{i}" for i in range(10)], "meta": ["meta_1", "meta_2"]}`
         transform: transform to apply on the loaded items of a dictionary data.
+        shuffle: whether to shuffle all the data in the buffer every time a new chunk loaded.
         kwargs: additional arguments for `pandas.merge()` API to join tables.
 
     """
@@ -101,19 +164,23 @@ def __init__(
         self,
         filename: Union[str, Sequence[str]],
         chunksize: int = 1000,
+        buffer_size: Optional[int] = None,
         col_names: Optional[Sequence[str]] = None,
         col_types: Optional[Dict[str, Optional[Dict[str, Any]]]] = None,
         col_groups: Optional[Dict[str, Sequence[str]]] = None,
         transform: Optional[Callable] = None,
+        shuffle: bool = False,
         **kwargs,
     ):
         self.files = ensure_tuple(filename)
         self.chunksize = chunksize
-        self.iters = self.reset()
+        self.buffer_size = 2 * chunksize if buffer_size is None else buffer_size
         self.col_names = col_names
         self.col_types = col_types
         self.col_groups = col_groups
+        self.shuffle = shuffle
         self.kwargs = kwargs
+        self.iters = self.reset()
         super().__init__(data=None, transform=transform)  # type: ignore
 
     def reset(self, filename: Optional[Union[str, Sequence[str]]] = None):
@@ -123,20 +190,19 @@ def reset(self, filename: Optional[Union[str, Sequence[str]]] = None):
         self.iters = [pd.read_csv(f, chunksize=self.chunksize) for f in self.files]
         return self.iters
 
-    def __iter__(self):
+    def _flattened(self):
         for chunks in zip(*self.iters):
-            self.data = convert_tables_to_dicts(
+            yield from convert_tables_to_dicts(
                 dfs=chunks,
                 col_names=self.col_names,
                 col_types=self.col_types,
                 col_groups=self.col_groups,
                 **self.kwargs,
             )
-            info = get_worker_info()
-            if info is not None:
-                length = len(self.data)
-                per_worker = int(math.ceil(length / float(info.num_workers)))
-                start = info.id * per_worker
-                self.data = self.data[start : min(start + per_worker, length)]
-
-            return super().__iter__()
+
+    def __iter__(self):
+        buffer = IterableBuffer(
+            data=self._flattened(), transform=self.transform, buffer_size=self.buffer_size, shuffle=self.shuffle
+        )
+        buffer.set_random_state(state=self.R)
+        yield from buffer


Test patch:
diff --git a/tests/test_csv_iterable_dataset.py b/tests/test_csv_iterable_dataset.py
--- a/tests/test_csv_iterable_dataset.py
+++ b/tests/test_csv_iterable_dataset.py
@@ -146,16 +146,19 @@ def prepare_csv_file(data, filepath):
 
             # test transform
             dataset = CSVIterableDataset(
+                chunksize=2,
+                buffer_size=4,
                 filename=[filepath1, filepath2, filepath3],
                 col_groups={"ehr": [f"ehr_{i}" for i in range(5)]},
                 transform=ToNumpyd(keys="ehr"),
+                shuffle=True,
             )
+            dataset.set_random_state(123)
             expected = [
+                [3.7725, 4.2118, 4.6353, 5.298, 9.5451],
                 [2.0078, 2.2902, 2.0549, 3.0196, 3.8078],
-                [6.8392, 6.4745, 5.8627, 5.1922, 5.2745],
-                [3.7725, 4.2118, 4.6353, 5.2980, 9.5451],
-                [3.3333, 3.2353, 3.4000, 3.1647, 3.0863],
                 [6.4275, 6.2549, 5.9765, 6.2627, 7.7176],
+                [6.8392, 6.4745, 5.8627, 5.1922, 5.2745],
             ]
             for item, exp in zip(dataset, expected):
                 self.assertTrue(isinstance(item["ehr"], np.ndarray))