Repository: pyg-team/pytorch_geometric
Pull request #8051 (2023-09-18 06:14:47)

`SQLiteDatabase` benchmark

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8051?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#8051](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8051?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (e706ebb) 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.68%`.
> The diff coverage is `77.77%`.

```diff
@@            Coverage Diff             @@
##           master    #8051      +/-   ##
==========================================
- Coverage   89.27%   88.59%   -0.68%     
==========================================
  Files         466      466              
  Lines       27686    27689       +3     
==========================================
- Hits        24717    24532     -185     
- Misses       2969     3157     +188     
```


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

... and [31 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8051/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 53e47e037..ffaa29234 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))
+- 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 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 647d3baa3..acf5842c6 100644
--- a/torch_geometric/data/database.py
+++ b/torch_geometric/data/database.py
@@ -5,6 +5,7 @@ from uuid import uuid4
 
 import torch
 from torch import Tensor
+from tqdm import tqdm
 
 
 class Database(ABC):
@@ -24,6 +25,7 @@ class Database(ABC):
         indices: Union[Iterable[int], Tensor, slice, range],
         data_list: Iterable[Any],
         batch_size: Optional[int] = None,
+        log: bool = False,
     ):
         if isinstance(indices, slice):
             indices = self.slice_to_range(indices)
@@ -31,7 +33,13 @@ class Database(ABC):
         length = min(len(indices), len(data_list))
         batch_size = length if batch_size is None else batch_size
 
-        for start in range(0, length, batch_size):
+        if log and length > batch_size:
+            desc = f'Insert {length} entries'
+            offsets = tqdm(range(0, length, batch_size), desc=desc)
+        else:
+            offsets = range(0, length, batch_size)
+
+        for start in offsets:
             self._multi_insert(
                 indices[start:start + batch_size],
                 data_list[start:start + batch_size],
@@ -78,6 +86,10 @@ class Database(ABC):
     @staticmethod
     def serialize(data: Any) -> bytes:
         r"""Serializes :obj:`data` into bytes."""
+        # Ensure that data is not a view of a larger tensor:
+        if isinstance(data, Tensor):
+            data = data.clone()
+
         buffer = io.BytesIO()
         torch.save(data, buffer)
         return buffer.getvalue()
@@ -120,7 +132,7 @@ class Database(ABC):
             self.multi_insert(key, value)
 
     def __repr__(self) -> str:
-        return f'{self.__class__.__name__}()'
+        return f'{self.__class__.__name__}({len(self)})'
 
 
 class SQLiteDatabase(Database):
@@ -171,6 +183,7 @@ class SQLiteDatabase(Database):
     ):
         if isinstance(indices, Tensor):
             indices = indices.tolist()
+
         data_list = [self.serialize(data) for data in data_list]
 
         query = f'INSERT INTO {self.name} (id, data) VALUES (?, ?)'


Test patch:
diff --git a/test/data/test_database.py b/test/data/test_database.py
index dd5586bc6..dcb94addc 100644
--- a/test/data/test_database.py
+++ b/test/data/test_database.py
@@ -4,6 +4,7 @@ import pytest
 import torch
 
 from torch_geometric.data.database import SQLiteDatabase
+from torch_geometric.profile import benchmark
 from torch_geometric.testing import withPackage
 
 
@@ -12,7 +13,7 @@ from torch_geometric.testing import withPackage
 def test_sqlite_database(tmp_path, batch_size):
     path = osp.join(tmp_path, 'sqlite.db')
     db = SQLiteDatabase(path, name='test_table')
-    assert str(db) == 'SQLiteDatabase()'
+    assert str(db) == 'SQLiteDatabase(0)'
     assert len(db) == 0
 
     data = torch.randn(5)
@@ -58,3 +59,42 @@ def test_database_syntactic_sugar(tmp_path):
         torch.stack(db[torch.tensor([4, 4])], dim=0),
         data[torch.tensor([4, 4])],
     )
+
+
+if __name__ == '__main__':
+    import argparse
+    import tempfile
+    import time
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--numel', type=int, default=100_000)
+    parser.add_argument('--batch_size', type=int, default=256)
+    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')
+
+    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')
+
+    def in_memory_get(data):
+        index = torch.randint(0, args.numel, (128, ))
+        return data[index]
+
+    def db_get(db):
+        index = torch.randint(0, args.numel, (128, ))
+        return db[index]
+
+    benchmark(
+        funcs=[in_memory_get, db_get],
+        func_names=['In-Memory', 'SQLite'],
+        args=[(data, ), (db, )],
+        num_steps=50,
+        num_warmups=5,
+    )
+
+    tmp_dir.cleanup()