Repository: pyg-team/pytorch_geometric
Pull request #8052 (2023-09-18 07:50:41)

`RocksDatabase` implementation

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8052?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#8052](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8052?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (0d72074) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/10464bd3f12fea07c94b39db8e4b71c041787b10?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (10464bd) will **decrease** coverage by `0.73%`.
> Report is 1 commits behind head on master.
> The diff coverage is `55.00%`.

> :exclamation: Current head 0d72074 differs from pull request most recent head 994e5eb. Consider uploading reports for the commit 994e5eb to get more accurate results

```diff
@@            Coverage Diff             @@
##           master    #8052      +/-   ##
==========================================
- Coverage   89.27%   88.54%   -0.73%     
==========================================
  Files         466      466              
  Lines       27686    27716      +30     
==========================================
- Hits        24717    24542     -175     
- Misses       2969     3174     +205     
```


| [Files Changed](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8052?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/database.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8052?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvZGF0YWJhc2UucHk=) | `78.52% <55.00%> (-8.40%)` | :arrow_down: |

... and [32 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8052/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 ffaa29234..f63561942 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,7 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
-- Added a `Database` interface and `SQLite` implementation ([#8028](https://github.com/pyg-team/pytorch_geometric/pull/8028), [#8044](https://github.com/pyg-team/pytorch_geometric/pull/8044), [#8046](https://github.com/pyg-team/pytorch_geometric/pull/8046), [#8051](https://github.com/pyg-team/pytorch_geometric/pull/8051))
+- Added a `Database` interface and `SQLiteDatabase`/`RocksDatabase` implementations ([#8028](https://github.com/pyg-team/pytorch_geometric/pull/8028), [#8044](https://github.com/pyg-team/pytorch_geometric/pull/8044), [#8046](https://github.com/pyg-team/pytorch_geometric/pull/8046), [#8051](https://github.com/pyg-team/pytorch_geometric/pull/8051), [#8052](https://github.com/pyg-team/pytorch_geometric/pull/8052))
 - Added support for weighted/biased sampling in `NeighborLoader`/`LinkNeighborLoader` ([#8038](https://github.com/pyg-team/pytorch_geometric/pull/8038))
 - Added the `MixHopConv` layer and an corresponding example ([#8025](https://github.com/pyg-team/pytorch_geometric/pull/8025))
 - Added the option to pass keyword arguments to the underlying normalization layers within `BasicGNN` and `MLP` ([#8024](https://github.com/pyg-team/pytorch_geometric/pull/8024), [#8033](https://github.com/pyg-team/pytorch_geometric/pull/8033))
diff --git a/torch_geometric/data/database.py b/torch_geometric/data/database.py
index acf5842c6..f8df0cb28 100644
--- a/torch_geometric/data/database.py
+++ b/torch_geometric/data/database.py
@@ -1,9 +1,8 @@
-import io
+import pickle
 from abc import ABC, abstractmethod
 from typing import Any, Iterable, List, Optional, Union
 from uuid import uuid4
 
-import torch
 from torch import Tensor
 from tqdm import tqdm
 
@@ -90,14 +89,12 @@ class Database(ABC):
         if isinstance(data, Tensor):
             data = data.clone()
 
-        buffer = io.BytesIO()
-        torch.save(data, buffer)
-        return buffer.getvalue()
+        return pickle.dumps(data)
 
     @staticmethod
     def deserialize(data: bytes) -> Any:
         r"""Deserializes bytes into the original data."""
-        return torch.load(io.BytesIO(data))
+        return pickle.loads(data)
 
     def slice_to_range(self, indices: slice) -> range:
         start = 0 if indices.start is None else indices.start
@@ -132,7 +129,10 @@ class Database(ABC):
             self.multi_insert(key, value)
 
     def __repr__(self) -> str:
-        return f'{self.__class__.__name__}({len(self)})'
+        try:
+            return f'{self.__class__.__name__}({len(self)})'
+        except NotImplementedError:
+            return f'{self.__class__.__name__}()'
 
 
 class SQLiteDatabase(Database):
@@ -161,10 +161,11 @@ class SQLiteDatabase(Database):
         self._cursor = self._connection.cursor()
 
     def close(self):
-        self._connection.commit()
-        self._connection.close()
-        self._connection = None
-        self._cursor = None
+        if self._connection is not None:
+            self._connection.commit()
+            self._connection.close()
+            self._connection = None
+            self._cursor = None
 
     @property
     def cursor(self) -> Any:
@@ -245,3 +246,55 @@ class SQLiteDatabase(Database):
         query = f'SELECT COUNT(*) FROM {self.name}'
         self.cursor.execute(query)
         return self.cursor.fetchone()[0]
+
+
+class RocksDatabase(Database):
+    def __init__(self, path: str):
+        super().__init__()
+
+        import rocksdict
+
+        self.path = path
+
+        self._db: Optional[rocksdict.Rdict] = None
+
+        self.connect()
+
+    def connect(self):
+        import rocksdict
+        self._db = rocksdict.Rdict(
+            self.path,
+            options=rocksdict.Options(raw_mode=True),
+        )
+
+    def close(self):
+        if self._db is not None:
+            self._db.close()
+            self._db = None
+
+    @property
+    def db(self) -> Any:
+        if self._db is None:
+            raise RuntimeError("No open database connection")
+        return self._db
+
+    @staticmethod
+    def to_key(index: int) -> bytes:
+        return index.to_bytes(8, byteorder='big', signed=True)
+
+    def insert(self, index: int, data: Any):
+        # Ensure that data is not a view of a larger tensor:
+        if isinstance(data, Tensor):
+            data = data.clone()
+
+        self.db[self.to_key(index)] = self.serialize(data)
+
+    def get(self, index: int) -> Any:
+        return self.deserialize(self.db[self.to_key(index)])
+
+    def _multi_get(self, indices: Union[Iterable[int], Tensor]) -> List[Any]:
+        if isinstance(indices, Tensor):
+            indices = indices.tolist()
+        indices = [self.to_key(index) for index in indices]
+        data_list = self.db[indices]
+        return [self.deserialize(data) for data in data_list]


Test patch:
diff --git a/test/data/test_database.py b/test/data/test_database.py
index dcb94addc..0a9a59f18 100644
--- a/test/data/test_database.py
+++ b/test/data/test_database.py
@@ -3,7 +3,7 @@ import os.path as osp
 import pytest
 import torch
 
-from torch_geometric.data.database import SQLiteDatabase
+from torch_geometric.data.database import RocksDatabase, SQLiteDatabase
 from torch_geometric.profile import benchmark
 from torch_geometric.testing import withPackage
 
@@ -35,6 +35,32 @@ def test_sqlite_database(tmp_path, batch_size):
     db.close()
 
 
+@withPackage('rocksdict')
+@pytest.mark.parametrize('batch_size', [None, 1])
+def test_rocks_database(tmp_path, batch_size):
+    path = osp.join(tmp_path, 'rocks.db')
+    db = RocksDatabase(path)
+    assert str(db) == 'RocksDatabase()'
+    with pytest.raises(NotImplementedError):
+        len(db)
+
+    data = torch.randn(5)
+    db.insert(0, data)
+    assert torch.equal(db.get(0), data)
+
+    indices = torch.tensor([1, 2])
+    data_list = torch.randn(2, 5)
+    db.multi_insert(indices, data_list, batch_size=batch_size)
+
+    out_list = db.multi_get(indices, batch_size=batch_size)
+    assert isinstance(out_list, list)
+    assert len(out_list) == 2
+    assert torch.equal(out_list[0], data_list[0])
+    assert torch.equal(out_list[1], data_list[1])
+
+    db.close()
+
+
 @withPackage('sqlite3')
 def test_database_syntactic_sugar(tmp_path):
     path = osp.join(tmp_path, 'sqlite.db')
@@ -72,14 +98,19 @@ if __name__ == '__main__':
     args = parser.parse_args()
 
     data = torch.randn(args.numel, 128)
-
     tmp_dir = tempfile.TemporaryDirectory()
+
     path = osp.join(tmp_dir.name, 'sqlite.db')
-    db = SQLiteDatabase(path, name='test_table')
+    sqlite_db = SQLiteDatabase(path, name='test_table')
+    t = time.perf_counter()
+    sqlite_db.multi_insert(range(args.numel), data, batch_size=100, log=True)
+    print(f'Initialized SQLiteDB in {time.perf_counter() - t:.2f} seconds')
 
+    path = osp.join(tmp_dir.name, 'rocks.db')
+    rocks_db = RocksDatabase(path)
     t = time.perf_counter()
-    db.multi_insert(range(args.numel), data, batch_size=100, log=True)
-    print(f'Initialized DB in {time.perf_counter() - t:.2f} seconds')
+    rocks_db.multi_insert(range(args.numel), data, batch_size=100, log=True)
+    print(f'Initialized RocksDB in {time.perf_counter() - t:.2f} seconds')
 
     def in_memory_get(data):
         index = torch.randint(0, args.numel, (128, ))
@@ -90,9 +121,9 @@ if __name__ == '__main__':
         return db[index]
 
     benchmark(
-        funcs=[in_memory_get, db_get],
-        func_names=['In-Memory', 'SQLite'],
-        args=[(data, ), (db, )],
+        funcs=[in_memory_get, db_get, db_get],
+        func_names=['In-Memory', 'SQLite', 'RocksDB'],
+        args=[(data, ), (sqlite_db, ), (rocks_db, )],
         num_steps=50,
         num_warmups=5,
     )