Repository: modin-project/modin
Pull request: FEAT-#6139: 🔥 DataLoader interplay.
Added wrappers for modin dataset s.t. it can be used with torch's DataLoader. Closes: #6139 

### Problem

Using a plain `torch`'s `Dataset` to wrap `modin`'s dataframe is..., well, slow. The reason for that is because of `modin`'s relative high latency, coupled with `torch`'s `DataLoader` always extracting element one by one, making the latency very noticeable compared to say, batch access.

### Solution

Custom `ModinDataLoader` that circumvents the `DataFrame -> Dataset -> DataLoader` pattern, s.t. `DataLoader` knows to extract one batch at a time, significantly speeding up the dataloader.
It also works with `torch`'s `Sampler`s, and aims to be a drop in replacement for `DataLoader` that directly wraps a `modin` dataframe.

- [x] first commit message and PR title follow format outlined [here](https://modin.readthedocs.io/en/latest/development/contributing.html#commit-message-formatting)
  > **_NOTE:_**  If you edit the PR title to match this format, you need to add another commit (even if it's empty) or amend your last commit for the CI job that checks the PR title to pick up the new PR title.
- [x] passes `flake8 modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [x] passes `black --check modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [x] signed commit with `git commit -s` <!-- you can amend your commit with a signature via `git commit -amend -s` -->
- [x] Resolves #6139? <!-- issue must be created for each patch -->
- [x] tests added and passing
- [x] module layout described at `docs/development/architecture.rst` is up-to-date <!-- if you have added, renamed or removed files or directories please update the documentation accordingly -->


Problem statement:
FEAT: add interplay with PyTorch DataLoader
Hey devs, thanks for the great work!

As an avid PyTorch user myself, is it possible you could make this work with pytorch's dataloaders? I tried this myself but do not know how to do this efficiently.

Hints from issue discussion:
@RehanSD did you have any ideas for this?

Commit message: FEAT-#6139: Sampler support added.

This makes it more in line with dataloader despite not being a subclass of
🔥 DataLoader.

Signed-off-by: RenChu Wang <patrick1031wang@gmail.com>

Patch to review:
diff --git a/modin/experimental/torch/__init__.py b/modin/experimental/torch/__init__.py
new file mode 100644
index 00000000000..29f35ac6c17
--- /dev/null
+++ b/modin/experimental/torch/__init__.py
@@ -0,0 +1,14 @@
+# Licensed to Modin Development Team under one or more contributor license agreements.
+# See the NOTICE file distributed with this work for additional information regarding
+# copyright ownership.  The Modin Development Team 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.
+
+"""Module supports conversion for torch `DataLoader` interplay."""
diff --git a/modin/experimental/torch/datasets.py b/modin/experimental/torch/datasets.py
new file mode 100644
index 00000000000..2ee924a883f
--- /dev/null
+++ b/modin/experimental/torch/datasets.py
@@ -0,0 +1,94 @@
+# Licensed to Modin Development Team under one or more contributor license agreements.
+# See the NOTICE file distributed with this work for additional information regarding
+# copyright ownership.  The Modin Development Team 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.
+from __future__ import annotations
+
+from typing import Hashable, Iterable, Protocol, Sequence, Type
+
+from pandas import DataFrame
+from torch.utils.data import SequentialSampler
+from modin.pandas import DataFrame as ModinDataFrame
+from abc import abstractmethod
+
+
+class Sampler(Protocol):
+    @abstractmethod
+    def __iter__(self) -> Iterable[int]:
+        ...
+
+
+class ModinDataLoader:
+    "A self explainatory class to convert a DataFrame into a DataLoader that batches rows."
+
+    def __init__(
+        self,
+        df: DataFrame | ModinDataFrame,
+        batch_size: int,
+        features: Sequence[Hashable] = (),
+        sampler: Type[Sampler] | Sampler = SequentialSampler,
+    ) -> None:
+        """
+        Converts a Pandas/Modin DataFrame into a torch DataLoader.
+
+        NOTE: This function should eventually go into modin/utils.py.
+
+        Parameters
+        ----------
+        df : DataFrame
+
+        batch_size : int, default: 1
+
+        features : Sequence[Hashable], default: ()
+            If specified, only these features will be used.
+
+        sampler: Type[Sampler] | Sampler, default: SequentialSampler
+            The sampler to use. By default, iterates over the DataFrame in order.
+
+        Returns
+        -------
+        DataLoader
+            DataLoader object backed by desired data.
+        """
+
+        if features:
+            df = df[features]
+
+        if isinstance(sampler, type):
+            sampler = sampler(df)
+
+        self._df = df
+        self._batch_size = batch_size
+        self._sampler = sampler
+
+    def __len__(self):
+        # Sampler length is always valid.
+        return len(self._sampler)
+
+    def __iter__(self):
+        idx_buffer = []
+
+        for cnt, idx in enumerate(self._sampler):
+            idx_buffer.append(idx)
+
+            if self._end_of_batch(cnt):
+                yield self._df.iloc[idx_buffer].to_numpy()
+                idx_buffer = []
+
+    def _end_of_batch(self, counter: int):
+        return (
+            False
+            or counter % self._batch_size == self._batch_size - 1
+            or counter == len(self._sampler) - 1
+        )
+
+
+PandasDataLoader = ModinDataLoader
diff --git a/modin/experimental/torch/tests/__init__.py b/modin/experimental/torch/tests/__init__.py
new file mode 100644
index 00000000000..cae6413e559
--- /dev/null
+++ b/modin/experimental/torch/tests/__init__.py
@@ -0,0 +1,12 @@
+# Licensed to Modin Development Team under one or more contributor license agreements.
+# See the NOTICE file distributed with this work for additional information regarding
+# copyright ownership.  The Modin Development Team 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.
diff --git a/modin/experimental/torch/tests/test_dataloader.py b/modin/experimental/torch/tests/test_dataloader.py
new file mode 100644
index 00000000000..8d8244ea4e8
--- /dev/null
+++ b/modin/experimental/torch/tests/test_dataloader.py
@@ -0,0 +1,74 @@
+# Licensed to Modin Development Team under one or more contributor license agreements.
+# See the NOTICE file distributed with this work for additional information regarding
+# copyright ownership.  The Modin Development Team 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.
+from __future__ import annotations
+
+from types import ModuleType
+
+import pandas as pd
+from torch.utils.data import RandomSampler
+
+from modin import pandas as mpd
+from modin.experimental.torch.datasets import ModinDataLoader
+
+
+def _load_test_dataframe(pandas: ModuleType):
+    df = pandas.read_csv(
+        "https://raw.githubusercontent.com/ponder-org/ponder-datasets/main/USA_Housing.csv"
+    )
+    return df
+
+
+def _test_torch_dataloader(pandas: ModuleType):
+    df = _load_test_dataframe(pandas)
+    loader = ModinDataLoader(
+        df,
+        batch_size=16,
+        features=[
+            "AVG_AREA_INCOME",
+            "AVG_AREA_HOUSE_AGE",
+            "AVG_AREA_NUM_ROOMS",
+            "AVG_AREA_NUM_BEDROOMS",
+            "POPULATION",
+            "PRICE",
+        ],
+    )
+    for batch in loader:
+        assert batch.shape[0] <= 16, batch.shape
+        assert batch.shape[1] == 6, batch.shape
+
+
+def test_torch_dataloader():
+    _test_torch_dataloader(pd)
+    _test_torch_dataloader(mpd)
+
+
+def test_random():
+    df = pd.read_csv(
+        "https://raw.githubusercontent.com/ponder-org/ponder-datasets/main/USA_Housing.csv"
+    )
+    loader = ModinDataLoader(
+        df,
+        batch_size=16,
+        features=[
+            "AVG_AREA_INCOME",
+            "AVG_AREA_HOUSE_AGE",
+            "AVG_AREA_NUM_ROOMS",
+            "AVG_AREA_NUM_BEDROOMS",
+            "POPULATION",
+            "PRICE",
+        ],
+        sampler=RandomSampler,
+    )
+    for batch in loader:
+        assert batch.shape[0] <= 16, batch.shape
+        assert batch.shape[1] == 6, batch.shape


Code review comments:
diff_hunk: @@ -0,0 +1,74 @@
+# Licensed to Modin Development Team under one or more contributor license agreements.
+# See the NOTICE file distributed with this work for additional information regarding
+# copyright ownership.  The Modin Development Team 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.
+from __future__ import annotations
+
+from types import ModuleType
+
+import pandas as pd
+from torch.utils.data import RandomSampler
+
+from modin import pandas as mpd
path: modin/experimental/torch/tests/test_dataloader.py
text: @user1:
we write the tests using `import modin.pandas as pd` form, like we say in README for the users; could you please also use that form of import so everything is uniform in the codebase?

diff_hunk: @@ -0,0 +1,74 @@
+# Licensed to Modin Development Team under one or more contributor license agreements.
+# See the NOTICE file distributed with this work for additional information regarding
+# copyright ownership.  The Modin Development Team 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.
+from __future__ import annotations
+
+from types import ModuleType
+
+import pandas as pd
+from torch.utils.data import RandomSampler
+
+from modin import pandas as mpd
+from modin.experimental.torch.datasets import ModinDataLoader
+
+
+def _load_test_dataframe(pandas: ModuleType):
+    df = pandas.read_csv(
+        "https://raw.githubusercontent.com/ponder-org/ponder-datasets/main/USA_Housing.csv"
+    )
+    return df
+
+
+def _test_torch_dataloader(pandas: ModuleType):
+    df = _load_test_dataframe(pandas)
+    loader = ModinDataLoader(
+        df,
+        batch_size=16,
+        features=[
+            "AVG_AREA_INCOME",
+            "AVG_AREA_HOUSE_AGE",
+            "AVG_AREA_NUM_ROOMS",
+            "AVG_AREA_NUM_BEDROOMS",
+            "POPULATION",
+            "PRICE",
+        ],
+    )
+    for batch in loader:
+        assert batch.shape[0] <= 16, batch.shape
+        assert batch.shape[1] == 6, batch.shape
+
+
+def test_torch_dataloader():
+    _test_torch_dataloader(pd)
+    _test_torch_dataloader(mpd)
+
+
+def test_random():
+    df = pd.read_csv(
+        "https://raw.githubusercontent.com/ponder-org/ponder-datasets/main/USA_Housing.csv"
+    )
path: modin/experimental/torch/tests/test_dataloader.py
text: @user1:
you already have a function which loads this `.csv`, call it here instead of copy-pasting 🙃

diff_hunk: @@ -0,0 +1,74 @@
+# Licensed to Modin Development Team under one or more contributor license agreements.
+# See the NOTICE file distributed with this work for additional information regarding
+# copyright ownership.  The Modin Development Team 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.
+from __future__ import annotations
+
+from types import ModuleType
+
+import pandas as pd
+from torch.utils.data import RandomSampler
+
+from modin import pandas as mpd
+from modin.experimental.torch.datasets import ModinDataLoader
+
+
+def _load_test_dataframe(pandas: ModuleType):
path: modin/experimental/torch/tests/test_dataloader.py
text: @user1:
I would suggest naming the argument as `lib` or `library`, as we do `import pandas` to use it in our codebase all around, and it looks confusing here

diff_hunk: @@ -0,0 +1,94 @@
+# Licensed to Modin Development Team under one or more contributor license agreements.
+# See the NOTICE file distributed with this work for additional information regarding
+# copyright ownership.  The Modin Development Team 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.
+from __future__ import annotations
+
+from typing import Hashable, Iterable, Protocol, Sequence, Type
+
+from pandas import DataFrame
+from torch.utils.data import SequentialSampler
+from modin.pandas import DataFrame as ModinDataFrame
+from abc import abstractmethod
+
+
+class Sampler(Protocol):
+    @abstractmethod
+    def __iter__(self) -> Iterable[int]:
+        ...
+
+
+class ModinDataLoader:
+    "A self explainatory class to convert a DataFrame into a DataLoader that batches rows."
+
+    def __init__(
+        self,
+        df: DataFrame | ModinDataFrame,
+        batch_size: int,
+        features: Sequence[Hashable] = (),
+        sampler: Type[Sampler] | Sampler = SequentialSampler,
+    ) -> None:
+        """
+        Converts a Pandas/Modin DataFrame into a torch DataLoader.
+
+        NOTE: This function should eventually go into modin/utils.py.
+
+        Parameters
+        ----------
+        df : DataFrame
+
+        batch_size : int, default: 1
+
+        features : Sequence[Hashable], default: ()
+            If specified, only these features will be used.
+
+        sampler: Type[Sampler] | Sampler, default: SequentialSampler
+            The sampler to use. By default, iterates over the DataFrame in order.
+
+        Returns
+        -------
+        DataLoader
+            DataLoader object backed by desired data.
+        """
+
+        if features:
+            df = df[features]
+
+        if isinstance(sampler, type):
+            sampler = sampler(df)
+
+        self._df = df
+        self._batch_size = batch_size
+        self._sampler = sampler
+
+    def __len__(self):
+        # Sampler length is always valid.
+        return len(self._sampler)
+
+    def __iter__(self):
+        idx_buffer = []
+
+        for cnt, idx in enumerate(self._sampler):
+            idx_buffer.append(idx)
+
+            if self._end_of_batch(cnt):
+                yield self._df.iloc[idx_buffer].to_numpy()
+                idx_buffer = []
+
+    def _end_of_batch(self, counter: int):
+        return (
+            False
+            or counter % self._batch_size == self._batch_size - 1
path: modin/experimental/torch/datasets.py
text: @user1:
```suggestion
            counter % self._batch_size == self._batch_size - 1
```

diff_hunk: @@ -0,0 +1,74 @@
+# Licensed to Modin Development Team under one or more contributor license agreements.
+# See the NOTICE file distributed with this work for additional information regarding
+# copyright ownership.  The Modin Development Team 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.
+from __future__ import annotations
+
+from types import ModuleType
+
+import pandas as pd
path: modin/experimental/torch/tests/test_dataloader.py
text: @user1:
we usually do `import pandas` in our codebase

diff_hunk: @@ -0,0 +1,94 @@
+# Licensed to Modin Development Team under one or more contributor license agreements.
+# See the NOTICE file distributed with this work for additional information regarding
+# copyright ownership.  The Modin Development Team 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.
+from __future__ import annotations
+
+from typing import Hashable, Iterable, Protocol, Sequence, Type
+
+from pandas import DataFrame
+from torch.utils.data import SequentialSampler
+from modin.pandas import DataFrame as ModinDataFrame
+from abc import abstractmethod
+
+
+class Sampler(Protocol):
+    @abstractmethod
+    def __iter__(self) -> Iterable[int]:
+        ...
path: modin/experimental/torch/datasets.py
text: @user1:
## Statement has no effect

This statement has no effect.

[Show more details](https://github.com/modin-project/modin/security/code-scanning/1015)