Repository: pyg-team/pytorch_geometric
Pull request #8044 (2023-09-17 08:24:05)

`Database` Updates and Tests

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8044?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#8044](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8044?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (8883c39) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/590879038336d10a141778e091d0b659c1554bcd?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (5908790) will **increase** coverage by `0.22%`.
> The diff coverage is `83.69%`.

```diff
@@            Coverage Diff             @@
##           master    #8044      +/-   ##
==========================================
+ Coverage   88.36%   88.59%   +0.22%     
==========================================
  Files         466      466              
  Lines       27626    27645      +19     
==========================================
+ Hits        24411    24491      +80     
+ Misses       3215     3154      -61     
```


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

: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 5cb3241ae..c5bbaaa8f 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 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 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/docs/source/external/resources.rst b/docs/source/external/resources.rst
index cd62497f9..0443165e3 100644
--- a/docs/source/external/resources.rst
+++ b/docs/source/external/resources.rst
@@ -40,3 +40,5 @@ External Resources
 * Manan Goel: **Recommending Amazon Products using Graph Neural Networks in** :pyg:`null` **PyTorch Geometric** [:wandb:`null` `W&B Report <https://wandb.ai/manan-goel/gnn-recommender/reports/Recommending-Amazon-Products-using-Graph-Neural-Networks-in-PyTorch-Geometric--VmlldzozMTA3MzYw>`__]
 
 * Kùzu: **Remote Backend for** :pyg:`null` **PyTorch Geometric** [:colab:`null` `Colab <https://colab.research.google.com/drive/12fOSqPm1HQTz_m9caRW7E_92vaeD9xq6>`__]
+
+* Aniket Saxena: **Graph Neural Networks-based Node and Graph Classification, and Explanation App using** :pyg:`null` **PyTorch Geometric** [`Website <https://graph-explainability.streamlit.app/>`__, :github:`null` `GitHub <https://github.com/fork123aniket/End-to-End-Node-and-Graph-Classification-and-Explanation-App>`__]
diff --git a/torch_geometric/data/database.py b/torch_geometric/data/database.py
index e1d48d440..b4ac67769 100644
--- a/torch_geometric/data/database.py
+++ b/torch_geometric/data/database.py
@@ -1,123 +1,167 @@
-import abc
-from collections import namedtuple
-from dataclasses import dataclass
 import io
-import itertools
-import json
-from typing import Iterable, Generator, Optional, Any
+from abc import ABC, abstractmethod
+from typing import Any, Iterable, List, Optional, Union
 
-import sqlite3
 import torch
+from torch import Tensor
+
+
+class Database(ABC):
+    r"""Base class for database."""
+    def connect(self):
+        pass
+
+    def close(self):
+        pass
+
+    @abstractmethod
+    def insert(self, index: int, data: Any):
+        raise NotImplementedError
+
+    def multi_insert(
+        self,
+        indices: Union[Iterable[int], Tensor],
+        data_list: Iterable[Any],
+        batch_size: Optional[int] = None,
+    ):
+        if batch_size is None:
+            batch_size = min(len(indices), len(data_list))
+
+        for start in range(0, min(len(indices), len(data_list)), batch_size):
+            self._multi_insert(
+                indices[start:start + batch_size],
+                data_list[start:start + batch_size],
+            )
+
+    def _multi_insert(
+        self,
+        indices: Union[Iterable[int], Tensor],
+        data_list: Iterable[Any],
+    ):
+        if isinstance(indices, Tensor):
+            indices = indices.tolist()
+        for index, data in zip(indices, data_list):
+            self.insert(index, data)
+
+    @abstractmethod
+    def get(self, index: int) -> Any:
+        raise NotImplementedError
+
+    def multi_get(
+        self,
+        indices: Union[Iterable[int], Tensor],
+        batch_size: Optional[int] = None,
+    ) -> List[Any]:
+        if batch_size is None:
+            batch_size = len(indices)
+
+        data_list: List[Any] = []
+        for start in range(0, len(indices), batch_size):
+            chunk_indices = indices[start:start + batch_size]
+            data_list.extend(self._multi_get(chunk_indices))
+        return data_list
+
+    def _multi_get(self, indices: Union[Iterable[int], Tensor]) -> List[Any]:
+        if isinstance(indices, Tensor):
+            indices = indices.tolist()
+        return [self.get(index) for index in indices]
+
+    # Helper functions ########################################################
 
-from torch_geometric.data.data import Data
-from torch_geometric.typing import OptTensor
-
-
-class GraphLabel:
-
-    def __init__(self, id: str, *args, **kwargs):
-        self.id = id
-        for key, value in kwargs.items():
-            setattr(self, key, value)
-
-@dataclass
-class GraphRow:
-    id: str
-    data: Optional[dict[str, Optional[bytes]]]
-
-
-def chunk(seq: Iterable, chunk_size: int) -> Generator[list, Any, None]:
-    "chunk data into chunks of chunk_size"
-    it = iter(seq)
-    while True:
-        batch = list(itertools.islice(it, chunk_size))
-        if not batch:
-            return
-        yield batch
-
-
-def namedtuple_factory(cursor, row):
-    """util function to create a namedtuple Row foe db results"""
-    fields = [column[0] for column in cursor.description]
-    cls = namedtuple("Row", fields)
-    return cls._make(row)
-
-
-class Database(abc.ABC):
-
-    def __init__(self, credentials, *args, **kwargs):
-        self.connection = self._get_connection(credentials)
-
-    @abc.abstractmethod
-    def _initialize(self):
-        """initialize the database in some way if needed"""
-        raise NotImplementedError()
-
-    @abc.abstractmethod
-    def insert(self, labels: Iterable[GraphLabel], values: Iterable[Data], batch_size=10000) -> list[str]:
-        """insert data into a database"""
-        raise NotImplementedError()
-
-    @abc.abstractmethod
-    def serialize_data(self, data: Data) -> GraphRow:
-        """serialize the data"""
-        raise NotImplementedError()
-
-    @abc.abstractmethod
-    def get(self, key: str) -> Data:
-        """get data by key"""
-        raise NotImplementedError()
-
-    def multi_get(self, keys: list[str]) -> list[Data]:
-        """get multiple keys"""
-        return [self.get(key) for key in keys]
-
-
-class SQLiteDatabase(abc.ABC):
-
-    def __init__(self, credentials, table='pyg_database') -> None:
-        self.table = table
-        super().__init__(credentials)
-
-    def _initialize(self):
-        create = """CREATE TABLE ? (id TEXT, data, TEXT)"""
-        self.cursor.execute(create, self.table)
-
-    def insert(self, labels: Iterable[GraphLabel], values: Iterable[Data], batch_size=10000) -> list[GraphRow]:
-        for chunk_data in chunk(zip(labels, values), batch_size):
-            serialized = [self.serialize_data(label, value) for label, value in chunk_data]
-            query = f"""
-            INSERT INTO {self.table} (id, data) 
-              VALUES (?, ?)"""
-            self.cursor.executemany(query, [(row['id'], json.dumps(row['data'])) for row in serialized])
-
-    def get(self, label: GraphLabel):
-        query = f"""SELECT * FROM {self.table} where id = ?"""
-        self.cursor.execute(query, (label.id))
-        return self.cursor.fetchone()
-    
-    def multi_get(self, labels: Iterable[GraphLabel], batch_size=999):
-        for chunk_data in chunk(labels, batch_size):
-            query = f"SELECT * FROM {self.table} WHERE id IN ({','.join('?' * len(chunk_data))})"
-            self.cursor.execute(query, (label.id for label in chunk_data))
-
-    def serialize_data(self, label: GraphLabel, data: Data) -> GraphRow:
-        row_dict = {k: self._serialize_tensor(v) if isinstance(v, OptTensor) else v for k, v in vars(data).items()}
-        return GraphRow(
-            **vars(label),
-            **row_dict
-        )
+    @staticmethod
+    def serialize(data: Any) -> bytes:
+        r"""Serializes :obj:`data` into bytes."""
+        buffer = io.BytesIO()
+        torch.save(data, buffer)
+        return buffer.getvalue()
 
     @staticmethod
-    def _serialize_tensor(t: OptTensor) -> bytes:
-        """convert a tensor into bytes"""
-        buff = io.BytesIO()
-        torch.save(t, buff)
-        return buff.getvalue()
-
-    def _get_connection(self, credentials):
-        """a method to get the db cursor to executor SQL"""
-        con = sqlite3.connect(credentials)
-        cursor = con.cursor()
-        cursor.row_factory = namedtuple_factory
-        return cursor
+    def deserialize(data: bytes) -> Any:
+        r"""Deserializes bytes into the original data."""
+        return torch.load(io.BytesIO(data))
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}()'
+
+
+class SQLiteDatabase(Database):
+    def __init__(self, path: str, name: str):
+        super().__init__()
+
+        import sqlite3
+
+        self.path = path
+        self.name = name
+
+        self._connection: Optional[sqlite3.Connection] = None
+        self._cursor: Optional[sqlite3.Cursor] = None
+
+        self.connect()
+
+        query = (f'CREATE TABLE IF NOT EXISTS {self.name} (\n'
+                 f'  id INTEGER PRIMARY KEY,\n'
+                 f'  data BLOB NOT NULL\n'
+                 f')')
+        self.cursor.execute(query)
+
+    def connect(self):
+        import sqlite3
+        self._connection = sqlite3.connect(self.path)
+        self._cursor = self._connection.cursor()
+
+    def close(self):
+        self._connection.commit()
+        self._connection.close()
+        self._connection = None
+        self._cursor = None
+
+    @property
+    def cursor(self) -> Any:
+        if self._cursor is None:
+            raise RuntimeError("No open database connection")
+        return self._cursor
+
+    def insert(self, index: int, data: Any):
+        query = f'INSERT INTO {self.name} (id, data) VALUES (?, ?)'
+        self.cursor.execute(query, (index, self.serialize(data)))
+
+    def _multi_insert(
+        self,
+        indices: Union[Iterable[int], Tensor],
+        data_list: Iterable[Any],
+    ):
+        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 (?, ?)'
+        self.cursor.executemany(query, zip(indices, data_list))
+
+    def get(self, index: int) -> Any:
+        query = f'SELECT data FROM {self.name} WHERE id = ?'
+        self.cursor.execute(query, (index, ))
+        return self.deserialize(self.cursor.fetchone()[0])
+
+    def multi_get(
+        self,
+        indices: Union[Iterable[int], Tensor],
+        batch_size: Optional[int] = None,
+    ) -> List[Any]:
+        if 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)
+
+        if batch_size is None:
+            data_list = self.cursor.fetchall()
+        else:
+            data_list: List[Any] = []
+            while True:
+                chunk_list = self.cursor.fetchmany(size=batch_size)
+                if len(chunk_list) == 0:
+                    break
+                data_list.extend(chunk_list)
+
+        return [self.deserialize(data[0]) for data in data_list]
diff --git a/torch_geometric/explain/algorithm/graphmask_explainer.py b/torch_geometric/explain/algorithm/graphmask_explainer.py
index 27ace22ae..2385ae320 100644
--- a/torch_geometric/explain/algorithm/graphmask_explainer.py
+++ b/torch_geometric/explain/algorithm/graphmask_explainer.py
@@ -50,6 +50,10 @@ class GraphMaskExplainer(ExplainerAlgorithm):
         <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
         /explain/graphmask_explainer.py>`_.
 
+        A working real-time example of :class:`GraphMaskExplainer` in the form
+        of a deployed app can be accessed `here
+        <https://graph-explainability.streamlit.app/>`_.
+
     Args:
         num_layers (int): The number of layers to use.
         epochs (int, optional): The number of epochs to train.


Test patch:
diff --git a/test/data/test_database.py b/test/data/test_database.py
new file mode 100644
index 000000000..cbd4379df
--- /dev/null
+++ b/test/data/test_database.py
@@ -0,0 +1,31 @@
+import os.path as osp
+
+import pytest
+import torch
+
+from torch_geometric.data.database import SQLiteDatabase
+from torch_geometric.testing import withPackage
+
+
+@withPackage('sqlite3')
+@pytest.mark.parametrize('batch_size', [None, 1])
+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()'
+
+    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()