Repository: pyg-team/pytorch_geometric
Pull request #8058 (2023-09-19 10:07:12)

`Database` Documentation

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8058?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#8058](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8058?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (854129a) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/8f52944f3bc7fbf5a71370cf75dd2bddd7d9d793?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (8f52944) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #8058   +/-   ##
=======================================
  Coverage   88.53%   88.53%           
=======================================
  Files         466      466           
  Lines       27806    27808    +2     
=======================================
+ Hits        24618    24620    +2     
  Misses       3188     3188           
```


| [Files Changed](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8058?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/8058?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvZGF0YWJhc2UucHk=) | `79.83% <ø> (ø)` | |
| [torch\_geometric/data/\_\_init\_\_.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8058?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvX19pbml0X18ucHk=) | `100.00% <100.00%> (ø)` | |

: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 cad4b1758..e84c025c1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,7 +8,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ### Added
 
 - Added `edge_attr` support to `ResGatedGraphConv` ([#8048](https://github.com/pyg-team/pytorch_geometric/pull/8048))
-- 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), [#8054](https://github.com/pyg-team/pytorch_geometric/pull/8054), [#8057](https://github.com/pyg-team/pytorch_geometric/pull/8057))
+- 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), [#8054](https://github.com/pyg-team/pytorch_geometric/pull/8054), [#8057](https://github.com/pyg-team/pytorch_geometric/pull/8057), [#8058](https://github.com/pyg-team/pytorch_geometric/pull/8058))
 - 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/modules/data.rst b/docs/source/modules/data.rst
index 25699b3c8..a2806801c 100644
--- a/docs/source/modules/data.rst
+++ b/docs/source/modules/data.rst
@@ -31,6 +31,20 @@ Remote Backend Interfaces
      {{ name }}
    {% endfor %}
 
+Databases
+---------
+
+.. currentmodule:: torch_geometric.data
+
+.. autosummary::
+   :nosignatures:
+   :toctree: ../generated
+   :template: autosummary/inherited_class.rst
+
+   {% for name in torch_geometric.data.database_classes %}
+     {{ name }}
+   {% endfor %}
+
 PyTorch Lightning Wrappers
 --------------------------
 
diff --git a/torch_geometric/data/__init__.py b/torch_geometric/data/__init__.py
index a4787b185..58c144f3d 100644
--- a/torch_geometric/data/__init__.py
+++ b/torch_geometric/data/__init__.py
@@ -6,6 +6,7 @@ from .data import Data
 from .hetero_data import HeteroData
 from .batch import Batch
 from .temporal import TemporalData
+from .database import Database, SQLiteDatabase, RocksDatabase
 from .dataset import Dataset
 from .in_memory_dataset import InMemoryDataset
 from .makedirs import makedirs
@@ -30,6 +31,12 @@ remote_backend_classes = [
     'EdgeAttr',
 ]
 
+database_classes = [
+    'Database',
+    'SQLiteDatabase',
+    'RocksDatabase',
+]
+
 helper_functions = [
     'makedirs',
     'download_url',
diff --git a/torch_geometric/data/database.py b/torch_geometric/data/database.py
index c180eac12..9843749bc 100644
--- a/torch_geometric/data/database.py
+++ b/torch_geometric/data/database.py
@@ -35,7 +35,52 @@ Schema = Union[Any, Dict[str, Any], Tuple[Any], List[Any]]
 
 
 class Database(ABC):
-    r"""Base class for inserting and retrieving data from a database."""
+    r"""Base class for inserting and retrieving data from a database.
+    A database acts as a persisted, out-of-memory and index-based key/value
+    store for tensor and custom data:
+
+    .. code-block:: python
+
+        db = Database()
+        db[0] = Data(x=torch.randn(5, 16), y=0, z='id_0')
+        print(db[0])
+        >>> Data(x=[5, 16], y=0, z='id_0')
+
+    To improve efficiency, it is recommended to specify the underlying
+    :obj:`schema` of the data:
+
+    .. code-block:: python
+
+        db = Database(schema={  # Custom schema:
+            # Tensor information can be specified through a dictionary:
+            'x': dict(dtype=torch.float, size=(-1, 16)),
+            'y': int,
+            'z': str,
+        })
+        db[0] = dict(x=torch.randn(5, 16), y=0, z='id_0')
+        print(db[0])
+        >>> {'x': torch.tensor(...), 'y': 0, 'z': 'id_0'}
+
+    In addition, databases support batch-wise insert and get, and support
+    syntactic sugar known from indexing Python lists, *e.g.*:
+
+    .. code-block:: python
+
+        db = Database()
+        db[2:5] = torch.randn(3, 16)
+        print(db[torch.tensor([2, 3])])
+        >>> [torch.tensor(...), torch.tensor(...)]
+
+    Args:
+        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. (default: :obj:`object`)
+    """
     def __init__(self, schema: Schema = object):
         schema = maybe_cast_to_tensor_info(schema)
         schema = self._to_dict(schema)
@@ -47,13 +92,23 @@ class Database(ABC):
         self.schema: Dict[Union[str, int], Any] = schema
 
     def connect(self):
+        r"""Connects to the database.
+        Databases will automatically connect on instantiation.
+        """
         pass
 
     def close(self):
+        r"""Closes the connection to the database."""
         pass
 
     @abstractmethod
     def insert(self, index: int, data: Any):
+        r"""Inserts data at the specified index.
+
+        Args:
+            index (int): The index at which to insert.
+            data (Any): The object to insert.
+        """
         raise NotImplementedError
 
     def multi_insert(
@@ -63,6 +118,18 @@ class Database(ABC):
         batch_size: Optional[int] = None,
         log: bool = False,
     ):
+        r"""Inserts a chunk of data at the specified indices.
+
+        Args:
+            indices (List[int] or torch.Tensor or range): The indices at which
+                to insert.
+            data_list (List[Any]): The objects to insert.
+            batch_size (int, optional): If specified, will insert the data to
+                the database in batches of size :obj:`batch_size`.
+                (default: :obj:`None`)
+            log (bool, optional): If set to :obj:`True`, will log progress to
+                the console. (default: :obj:`False`)
+        """
         if isinstance(indices, slice):
             indices = self.slice_to_range(indices)
 
@@ -93,6 +160,11 @@ class Database(ABC):
 
     @abstractmethod
     def get(self, index: int) -> Any:
+        r"""Gets data from the specified index.
+
+        Args:
+            index (int): The index to query.
+        """
         raise NotImplementedError
 
     def multi_get(
@@ -100,6 +172,14 @@ class Database(ABC):
         indices: Union[Iterable[int], Tensor, slice, range],
         batch_size: Optional[int] = None,
     ) -> List[Any]:
+        r"""Gets a chunk of data from the specified indices.
+
+        Args:
+            indices (List[int] or torch.Tensor or range): The indices to query.
+            batch_size (int, optional): If specified, will request the data
+                from the database in batches of size :obj:`batch_size`.
+                (default: :obj:`None`)
+        """
         if isinstance(indices, slice):
             indices = self.slice_to_range(indices)
 
@@ -168,6 +248,23 @@ class Database(ABC):
 
 
 class SQLiteDatabase(Database):
+    r"""An index-based key/value database based on :obj:`sqlite3`.
+
+    .. note::
+        This database implementation requires the :obj:`sqlite3` package.
+
+    Args:
+        path (str): The path to where the database should be saved.
+        name (str): The name of the table to save the data to.
+        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. (default: :obj:`object`)
+    """
     def __init__(self, path: str, name: str, schema: Schema = object):
         super().__init__(schema)
 
@@ -366,6 +463,26 @@ class SQLiteDatabase(Database):
 
 
 class RocksDatabase(Database):
+    r"""An index-based key/value database based on :obj:`RocksDB`.
+
+    .. note::
+        This database implementation requires the :obj:`rocksdict` package.
+
+    .. warning::
+        :class:`RocksDatabase` is currently less optimized than
+        :class:`SQLiteDatabase`.
+
+    Args:
+        path (str): The path to where the database should be saved.
+        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. (default: :obj:`object`)
+    """
     def __init__(self, path: str, schema: Schema = object):
         super().__init__(schema)
 


Test patch:
diff --git a/test/data/test_database.py b/test/data/test_database.py
index 7d3b99a25..5d7d055fb 100644
--- a/test/data/test_database.py
+++ b/test/data/test_database.py
@@ -3,12 +3,8 @@ import os.path as osp
 import pytest
 import torch
 
-from torch_geometric.data import Data
-from torch_geometric.data.database import (
-    RocksDatabase,
-    SQLiteDatabase,
-    TensorInfo,
-)
+from torch_geometric.data import Data, RocksDatabase, SQLiteDatabase
+from torch_geometric.data.database import TensorInfo
 from torch_geometric.profile import benchmark
 from torch_geometric.testing import has_package, withPackage