update

Author: rusty1s

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 @@ def serialize(data: Any) -> bytes:
         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
@@ -265,7 +262,10 @@ def __init__(self, path: str):
 
     def connect(self):
         import rocksdict
-        self._db = rocksdict.Rdict(self.path)
+        self._db = rocksdict.Rdict(
+            self.path,
+            options=rocksdict.Options(raw_mode=True),
+        )
 
     def close(self):
         if self._db is not None:
@@ -278,17 +278,23 @@ def db(self) -> Any:
             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[index] = data
+        self.db[self.to_key(index)] = self.serialize(data)
 
     def get(self, index: int) -> Any:
-        return self.db[self.to_key(index)]
+        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()
-        return self.db[indices]
+        indices = [self.to_key(index) for index in indices]
+        data_list = self.db[indices]
+        return [self.deserialize(data) for data in data_list]