Repository: pyg-team/pytorch_geometric
Pull request #8066 (2023-09-21 17:54:00)

`OnDiskDataset` class with `Database` integration

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8066?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#8066](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8066?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (1a5ef75) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/5d81cab11c4a29b8975abd46440f3e85c64dc0d4?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (5d81cab) will **decrease** coverage by `0.83%`.
> Report is 5 commits behind head on master.
> The diff coverage is `0.00%`.

```diff
@@            Coverage Diff             @@
##           master    #8066      +/-   ##
==========================================
- Coverage   89.23%   88.40%   -0.83%     
==========================================
  Files         466      467       +1     
  Lines       27812    27848      +36     
==========================================
- Hits        24817    24620     -197     
- Misses       2995     3228     +233     
```


| [Files Changed](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8066?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/data/on\_disk\_dataset.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8066?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvb25fZGlza19kYXRhc2V0LnB5) | `0.00% <0.00%> (ø)` | |

... and [32 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8066/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2999faeff..76e0eafc2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added `OnDiskDataset` interface ([#8066](https://github.com/pyg-team/pytorch_geometric/pull/8066))
 - Added a tutorial for `Node2Vec` and `MetaPath2Vec` usage ([#7938](https://github.com/pyg-team/pytorch_geometric/pull/7938)
 - Added a tutorial for multi-GPU training with pure PyTorch ([#7894](https://github.com/pyg-team/pytorch_geometric/pull/7894)
 - Added `edge_attr` support to `ResGatedGraphConv` ([#8048](https://github.com/pyg-team/pytorch_geometric/pull/8048))
diff --git a/torch_geometric/data/__init__.py b/torch_geometric/data/__init__.py
index 58c144f3d..add215bd3 100644
--- a/torch_geometric/data/__init__.py
+++ b/torch_geometric/data/__init__.py
@@ -9,6 +9,7 @@ from .temporal import TemporalData
 from .database import Database, SQLiteDatabase, RocksDatabase
 from .dataset import Dataset
 from .in_memory_dataset import InMemoryDataset
+from .on_disk_dataset import OnDiskDataset
 from .makedirs import makedirs
 from .download import download_url
 from .extract import extract_tar, extract_zip, extract_bz2, extract_gz
@@ -22,6 +23,7 @@ data_classes = [
     'TemporalData',
     'Dataset',
     'InMemoryDataset',
+    'OnDiskDataset',
 ]
 
 remote_backend_classes = [
diff --git a/torch_geometric/data/dataset.py b/torch_geometric/data/dataset.py
index 2b4b44c30..e62a78a5f 100644
--- a/torch_geometric/data/dataset.py
+++ b/torch_geometric/data/dataset.py
@@ -67,7 +67,7 @@ class Dataset(torch.utils.data.Dataset, ABC):
 
     @abstractmethod
     def len(self) -> int:
-        r"""Returns the number of graphs stored in the dataset."""
+        r"""Returns the number of data objects stored in the dataset."""
         raise NotImplementedError
 
     @abstractmethod
diff --git a/torch_geometric/data/in_memory_dataset.py b/torch_geometric/data/in_memory_dataset.py
index f8b5b1733..81940e4ea 100644
--- a/torch_geometric/data/in_memory_dataset.py
+++ b/torch_geometric/data/in_memory_dataset.py
@@ -27,7 +27,6 @@ from torch_geometric.data.separate import separate
 class InMemoryDataset(Dataset, ABC):
     r"""Dataset base class for creating graph datasets which easily fit
     into CPU memory.
-    Inherits from :class:`torch_geometric.data.Dataset`.
     See `here <https://pytorch-geometric.readthedocs.io/en/latest/tutorial/
     create_dataset.html#creating-in-memory-datasets>`__ for the accompanying
     tutorial.
diff --git a/torch_geometric/data/on_disk_dataset.py b/torch_geometric/data/on_disk_dataset.py
new file mode 100644
index 000000000..b011573f7
--- /dev/null
+++ b/torch_geometric/data/on_disk_dataset.py
@@ -0,0 +1,163 @@
+import os
+from abc import ABC
+from typing import Any, Callable, Iterable, List, Optional, Union
+
+from torch import Tensor
+
+from torch_geometric.data import Database, RocksDatabase, SQLiteDatabase
+from torch_geometric.data.data import BaseData
+from torch_geometric.data.database import Schema
+from torch_geometric.data.dataset import Dataset
+
+
+class OnDiskDataset(Dataset, ABC):
+    r"""Dataset base class for creating large graph datasets which do not
+    easily fit into CPU memory at once by leveraging a :class:`Database`
+    backend for on-disk storage and access of data objects.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        transform (callable, optional): A function/transform that takes in a
+            :class:`~torch_geometric.data.Data` or
+            :class:`~torch_geometric.data.HeteroData` object and returns a
+            transformed version.
+            The data object will be transformed before every access.
+            (default: :obj:`None`)
+        pre_filter (callable, optional): A function that takes in a
+            :class:`~torch_geometric.data.Data` or
+            :class:`~torch_geometric.data.HeteroData` object and returns a
+            boolean value, indicating whether the data object should be
+            included in the final dataset. (default: :obj:`None`)
+        backend (str): The :class:`Database` backend to use
+            (one of :obj:`"sqlite"` or :obj:`"rocksdb"`).
+            (default: :obj:`"sqlite"`)
+        schema (Any or Tuple[Any] or Dict[str, Any], optional): The schema of
+            the input data.
+            Can take :obj:`int`, :obj:`float`, :obj:`str`, :obj:`object`, or a
+            dictionary with :obj:`dtype` and :obj:`size` keys (for specifying
+            tensor data) as input, and can be nested as a tuple or dictionary.
+            Specifying the schema will improve efficiency, since by default the
+            database will use python pickling for serializing and
+            deserializing. If specified to anything different than
+            :obj:`object`, implementations of :class:`OnDiskDataset` need to
+            override :meth:`serialize` and :meth:`deserialize` methods.
+            (default: :obj:`object`)
+        log (bool, optional): Whether to print any console output while
+            downloading and processing the dataset. (default: :obj:`True`)
+    """
+    BACKENDS = {
+        'sqlite': SQLiteDatabase,
+        'rocksdb': RocksDatabase,
+    }
+
+    def __init__(
+        self,
+        root: str,
+        transform: Optional[Callable] = None,
+        *,
+        pre_filter: Optional[Callable] = None,
+        backend: str = 'sqlite',
+        schema: Schema = object,
+        log: bool = True,
+    ):
+        if backend not in self.BACKENDS:
+            raise ValueError(f"Database backend must be one of "
+                             f"{set(self.BACKENDS.keys())} "
+                             f"(got '{backend}')")
+
+        self.backend = backend
+        self.schema = schema
+
+        self._db: Optional[Database] = None
+        self._numel: Optional[int] = None
+
+        super().__init__(root, transform, pre_filter=pre_filter, log=log)
+
+    @property
+    def processed_file_names(self) -> str:
+        return f'{self.backend}.db'
+
+    @property
+    def db(self) -> Database:
+        r"""Returns the underlying :class:`Database`."""
+        if self._db is not None:
+            return self._db
+
+        kwargs = {}
+        cls = self.BACKENDS[self.backend]
+        if issubclass(cls, SQLiteDatabase):
+            kwargs['name'] = self.__class__.__name__
+
+        os.makedirs(self.processed_dir, exist_ok=True)
+        path = self.processed_paths[0]
+        self._db = cls(path=path, schema=self.schema, **kwargs)
+        self._numel = len(self._db)
+        return self._db
+
+    def close(self):
+        r"""Closes the connection to the underlying database."""
+        if self._db is not None:
+            self._db.close()
+
+    def serialize(self, data: BaseData) -> Any:
+        r"""Serializes the :class:`~torch_geometric.data.Data` or
+        :class:`~torch_geometric.data.HeteroData` object into the expected DB
+        schema."""
+        if self.schema == object:
+            return data
+        raise NotImplementedError(f"`{self.__class__.__name__}.serialize()` "
+                                  f"needs to be overridden in case a "
+                                  f"non-default schema was passed")
+
+    def deserialize(self, data: Any) -> BaseData:
+        r"""Deserializes the DB entry into a
+        :class:`~torch_geometric.data.Data` or
+        :class:`~torch_geometric.data.HeteroData` object."""
+        if self.schema == object:
+            return data
+        raise NotImplementedError(f"`{self.__class__.__name__}.deserialize()` "
+                                  f"needs to be overridden in case a "
+                                  f"non-default schema was passed")
+
+    def append(self, data: BaseData):
+        r"""Appends the data object to the dataset."""
+        index = len(self)
+        self.db.insert(index, self.serialize(data))
+        self._numel += 1
+
+    def extend(
+        self,
+        data_list: List[BaseData],
+        batch_size: Optional[int] = None,
+    ):
+        r"""Extends the dataset by a list of data objects."""
+        start = len(self)
+        end = start + len(data_list)
+        data_list = [self.serialize(data) for data in data_list]
+        self.db.multi_insert(range(start, end), data_list, batch_size)
+        self._numel += (end - start)
+
+    def get(self, idx: int) -> BaseData:
+        r"""Gets the data object at index :obj:`idx`."""
+        return self.deserialize(self.db.get(idx))
+
+    def multi_get(
+        self,
+        indices: Union[Iterable[int], Tensor, slice, range],
+        batch_size: Optional[int] = None,
+    ) -> List[BaseData]:
+        r"""Gets a list of data objects from the specified indices."""
+        if len(indices) == 1:
+            data_list = [self.db.get(indices[0])]
+        else:
+            data_list = self.db.multi_get(indices, batch_size)
+
+        return [self.deserialize(data) for data in data_list]
+
+    def len(self) -> int:
+        if self._numel is None:
+            self._numel = len(self.db)
+        return self._numel
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({len(self)})'


Test patch:
diff --git a/test/data/test_on_disk_dataset.py b/test/data/test_on_disk_dataset.py
new file mode 100644
index 000000000..1a19c3213
--- /dev/null
+++ b/test/data/test_on_disk_dataset.py
@@ -0,0 +1,100 @@
+import os.path as osp
+from typing import Any, Dict
+
+import torch
+
+from torch_geometric.data import Data, OnDiskDataset
+from torch_geometric.testing import withPackage
+
+
+@withPackage('sqlite3')
+def test_pickle(tmp_path):
+    dataset = OnDiskDataset(tmp_path)
+    assert len(dataset) == 0
+    assert str(dataset) == 'OnDiskDataset(0)'
+    assert osp.exists(osp.join(tmp_path, 'processed', 'sqlite.db'))
+
+    data_list = [
+        Data(
+            x=torch.randn(5, 8),
+            edge_index=torch.randint(0, 5, (2, 16)),
+            num_nodes=5,
+        ) for _ in range(4)
+    ]
+
+    dataset.append(data_list[0])
+    assert len(dataset) == 1
+
+    dataset.extend(data_list[1:])
+    assert len(dataset) == 4
+
+    out = dataset.get(0)
+    assert torch.equal(out.x, data_list[0].x)
+    assert torch.equal(out.edge_index, data_list[0].edge_index)
+    assert out.num_nodes == data_list[0].num_nodes
+
+    out_list = dataset.multi_get([1, 2, 3])
+    for out, data in zip(out_list, data_list[1:]):
+        assert torch.equal(out.x, data.x)
+        assert torch.equal(out.edge_index, data.edge_index)
+        assert out.num_nodes == data.num_nodes
+
+    dataset.close()
+
+
+@withPackage('sqlite3')
+def test_custom_schema(tmp_path):
+    class CustomSchemaOnDiskDataset(OnDiskDataset):
+        def __init__(self, root: str):
+            schema = {
+                'x': dict(dtype=torch.float, size=(-1, 8)),
+                'edge_index': dict(dtype=torch.long, size=(2, -1)),
+                'num_nodes': int,
+            }
+            self.serialize_count = 0
+            self.deserialize_count = 0
+            super().__init__(root, schema=schema)
+
+        def serialize(self, data: Data) -> Dict[str, Any]:
+            self.serialize_count += 1
+            return data.to_dict()
+
+        def deserialize(self, mapping: Dict[str, Any]) -> Any:
+            self.deserialize_count += 1
+            return Data.from_dict(mapping)
+
+    dataset = CustomSchemaOnDiskDataset(tmp_path)
+    assert len(dataset) == 0
+    assert str(dataset) == 'CustomSchemaOnDiskDataset(0)'
+    assert osp.exists(osp.join(tmp_path, 'processed', 'sqlite.db'))
+
+    data_list = [
+        Data(
+            x=torch.randn(5, 8),
+            edge_index=torch.randint(0, 5, (2, 16)),
+            num_nodes=5,
+        ) for _ in range(4)
+    ]
+
+    dataset.append(data_list[0])
+    assert dataset.serialize_count == 1
+    assert len(dataset) == 1
+
+    dataset.extend(data_list[1:])
+    assert dataset.serialize_count == 4
+    assert len(dataset) == 4
+
+    out = dataset.get(0)
+    assert dataset.deserialize_count == 1
+    assert torch.equal(out.x, data_list[0].x)
+    assert torch.equal(out.edge_index, data_list[0].edge_index)
+    assert out.num_nodes == data_list[0].num_nodes
+
+    out_list = dataset.multi_get([1, 2, 3])
+    assert dataset.deserialize_count == 4
+    for out, data in zip(out_list, data_list[1:]):
+        assert torch.equal(out.x, data.x)
+        assert torch.equal(out.edge_index, data.edge_index)
+        assert out.num_nodes == data.num_nodes
+
+    dataset.close()