Repository: pyg-team/pytorch_geometric
Pull request #8046 (2023-09-17 09:47:23)

`SQLiteDatabase`: Better `multi_get` + syntactic sugar

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8046?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#8046](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8046?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (8bb9613) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/3f7a8b9cf8075451d22e77c3e929a124bc831348?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (3f7a8b9) will **increase** coverage by `0.01%`.
> The diff coverage is `90.00%`.

```diff
@@            Coverage Diff             @@
##           master    #8046      +/-   ##
==========================================
+ Coverage   88.59%   88.60%   +0.01%     
==========================================
  Files         466      466              
  Lines       27645    27682      +37     
==========================================
+ Hits        24491    24527      +36     
- Misses       3154     3155       +1     
```


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

... and [1 file with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8046/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 c5bbaaa8f..53e47e037 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))
+- 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 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 b4ac67769..647d3baa3 100644
--- a/torch_geometric/data/database.py
+++ b/torch_geometric/data/database.py
@@ -1,6 +1,7 @@
 import io
 from abc import ABC, abstractmethod
 from typing import Any, Iterable, List, Optional, Union
+from uuid import uuid4
 
 import torch
 from torch import Tensor
@@ -20,14 +21,17 @@ class Database(ABC):
 
     def multi_insert(
         self,
-        indices: Union[Iterable[int], Tensor],
+        indices: Union[Iterable[int], Tensor, slice, range],
         data_list: Iterable[Any],
         batch_size: Optional[int] = None,
     ):
-        if batch_size is None:
-            batch_size = min(len(indices), len(data_list))
+        if isinstance(indices, slice):
+            indices = self.slice_to_range(indices)
+
+        length = min(len(indices), len(data_list))
+        batch_size = length if batch_size is None else batch_size
 
-        for start in range(0, min(len(indices), len(data_list)), batch_size):
+        for start in range(0, length, batch_size):
             self._multi_insert(
                 indices[start:start + batch_size],
                 data_list[start:start + batch_size],
@@ -35,7 +39,7 @@ class Database(ABC):
 
     def _multi_insert(
         self,
-        indices: Union[Iterable[int], Tensor],
+        indices: Union[Iterable[int], Tensor, range],
         data_list: Iterable[Any],
     ):
         if isinstance(indices, Tensor):
@@ -49,14 +53,17 @@ class Database(ABC):
 
     def multi_get(
         self,
-        indices: Union[Iterable[int], Tensor],
+        indices: Union[Iterable[int], Tensor, slice, range],
         batch_size: Optional[int] = None,
     ) -> List[Any]:
-        if batch_size is None:
-            batch_size = len(indices)
+        if isinstance(indices, slice):
+            indices = self.slice_to_range(indices)
+
+        length = len(indices)
+        batch_size = length if batch_size is None else batch_size
 
         data_list: List[Any] = []
-        for start in range(0, len(indices), batch_size):
+        for start in range(0, length, batch_size):
             chunk_indices = indices[start:start + batch_size]
             data_list.extend(self._multi_get(chunk_indices))
         return data_list
@@ -80,6 +87,38 @@ class Database(ABC):
         r"""Deserializes bytes into the original data."""
         return torch.load(io.BytesIO(data))
 
+    def slice_to_range(self, indices: slice) -> range:
+        start = 0 if indices.start is None else indices.start
+        stop = len(self) if indices.stop is None else indices.stop
+        step = 1 if indices.step is None else indices.step
+
+        return range(start, stop, step)
+
+    # Python built-ins ########################################################
+
+    def __len__(self) -> int:
+        raise NotImplementedError
+
+    def __getitem__(
+        self,
+        key: Union[int, Iterable[int], Tensor, slice, range],
+    ) -> Union[Any, List[Any]]:
+
+        if isinstance(key, int):
+            return self.get(key)
+        else:
+            return self.multi_get(key)
+
+    def __setitem__(
+        self,
+        key: Union[int, Iterable[int], Tensor, slice, range],
+        value: Union[Any, Iterable[Any]],
+    ):
+        if isinstance(key, int):
+            self.insert(key, value)
+        else:
+            self.multi_insert(key, value)
+
     def __repr__(self) -> str:
         return f'{self.__class__.__name__}()'
 
@@ -127,7 +166,7 @@ class SQLiteDatabase(Database):
 
     def _multi_insert(
         self,
-        indices: Union[Iterable[int], Tensor],
+        indices: Union[Iterable[int], Tensor, range],
         data_list: Iterable[Any],
     ):
         if isinstance(indices, Tensor):
@@ -144,15 +183,35 @@ class SQLiteDatabase(Database):
 
     def multi_get(
         self,
-        indices: Union[Iterable[int], Tensor],
+        indices: Union[Iterable[int], Tensor, slice, range],
         batch_size: Optional[int] = None,
     ) -> List[Any]:
-        if isinstance(indices, Tensor):
+
+        if isinstance(indices, slice):
+            indices = self.slice_to_range(indices)
+        elif isinstance(indices, Tensor):
             indices = indices.tolist()
 
-        query = (f'SELECT data FROM {self.name} '
-                 f'WHERE id IN ({", ".join("?" * len(indices))})')
-        self.cursor.execute(query, indices)
+        # We first create a temporary ID table to then perform an INNER JOIN.
+        # This avoids having a long IN clause and guarantees sorted outputs:
+        join_table_name = f'{self.name}__join__{uuid4().hex}'
+        query = (f'CREATE TABLE {join_table_name} (\n'
+                 f'  id INTEGER,\n'
+                 f'  row_id INTEGER\n'
+                 f')')
+        self.cursor.execute(query)
+
+        query = f'INSERT INTO {join_table_name} (id, row_id) VALUES (?, ?)'
+        self.cursor.executemany(query, zip(indices, range(len(indices))))
+
+        query = f'SELECT * FROM {join_table_name}'
+        self.cursor.execute(query)
+
+        query = (f'SELECT {self.name}.data '
+                 f'FROM {self.name} INNER JOIN {join_table_name} '
+                 f'ON {self.name}.id = {join_table_name}.id '
+                 f'ORDER BY {join_table_name}.row_id')
+        self.cursor.execute(query)
 
         if batch_size is None:
             data_list = self.cursor.fetchall()
@@ -164,4 +223,12 @@ class SQLiteDatabase(Database):
                     break
                 data_list.extend(chunk_list)
 
+        query = f'DROP TABLE {join_table_name}'
+        self.cursor.execute(query)
+
         return [self.deserialize(data[0]) for data in data_list]
+
+    def __len__(self) -> int:
+        query = f'SELECT COUNT(*) FROM {self.name}'
+        self.cursor.execute(query)
+        return self.cursor.fetchone()[0]


Test patch:
diff --git a/test/data/test_database.py b/test/data/test_database.py
index cbd4379df..dd5586bc6 100644
--- a/test/data/test_database.py
+++ b/test/data/test_database.py
@@ -13,14 +13,17 @@ 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 len(db) == 0
 
     data = torch.randn(5)
     db.insert(0, data)
+    assert len(db) == 1
     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)
+    assert len(db) == 3
 
     out_list = db.multi_get(indices, batch_size=batch_size)
     assert isinstance(out_list, list)
@@ -29,3 +32,29 @@ def test_sqlite_database(tmp_path, batch_size):
     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')
+    db = SQLiteDatabase(path, name='test_table')
+
+    data = torch.randn(5, 16)
+    db[0] = data[0]
+    db[1:3] = data[1:3]
+    db[torch.tensor([3, 4])] = data[torch.tensor([3, 4])]
+    assert len(db) == 5
+
+    assert torch.equal(db[0], data[0])
+    assert torch.equal(torch.stack(db[:3], dim=0), data[:3])
+    assert torch.equal(torch.stack(db[3:], dim=0), data[3:])
+    assert torch.equal(torch.stack(db[1::2], dim=0), data[1::2])
+    assert torch.equal(torch.stack(db[[4, 3]], dim=0), data[[4, 3]])
+    assert torch.equal(
+        torch.stack(db[torch.tensor([4, 3])], dim=0),
+        data[torch.tensor([4, 3])],
+    )
+    assert torch.equal(
+        torch.stack(db[torch.tensor([4, 4])], dim=0),
+        data[torch.tensor([4, 4])],
+    )