fix: reload graph before quizzing

Author: ChenZiHong-Gavin

graphgen/models/storage/__init__.py

@@ -1,3 +1,4 @@
-from .json_storage import JsonKVStorage, JsonListStorage
-from .networkx_storage import NetworkXStorage
+from graphgen.models.storage.graph.networkx_storage import NetworkXStorage
+from graphgen.models.storage.kv.json_storage import JsonKVStorage
+
 from .rocksdb_cache import RocksDBCache

graphgen/models/storage/graph/__init__.py

@@ -170,3 +170,9 @@ def clear(self):
         """
         self._graph.clear()
         logger.info("Graph %s cleared.", self.namespace)
+
+    def reload(self):
+        """
+        Reload the graph from the GraphML file.
+        """
+        self.__post_init__()

graphgen/models/storage/graph/networkx_storage.py

@@ -54,3 +54,7 @@ def upsert(self, data: dict):
     def drop(self):
         if self._data:
             self._data.clear()
+
+    def reload(self):
+        self._data = load_json(self._file_name) or {}
+        logger.info("Reload KV %s with %d data", self.namespace, len(self._data))

graphgen/models/storage/kv/__init__.py

@@ -0,0 +1,79 @@
+import os
+from dataclasses import dataclass
+from typing import Any, Dict, List, Set
+
+# rocksdict is a lightweight C wrapper around RocksDB for Python, pylint may not recognize it
+# pylint: disable=no-name-in-module
+from rocksdict import Rdict
+
+from graphgen.bases.base_storage import BaseKVStorage
+from graphgen.utils import logger
+
+
+@dataclass
+class RocksDBKVStorage(BaseKVStorage):
+    _db: Rdict = None
+    _db_path: str = None
+
+    def __post_init__(self):
+        self._db_path = os.path.join(self.working_dir, f"{self.namespace}.db")
+        self._db = Rdict(self._db_path)
+        logger.info("Load KV (RocksDB) %s at %s", self.namespace, self._db_path)
+
+    @property
+    def data(self):
+        return self._db
+
+    def all_keys(self) -> List[str]:
+        return list(self._db.keys())
+
+    def index_done_callback(self):
+        self._db.flush()
+        logger.info("RocksDB flushed for %s", self.namespace)
+
+    def get_by_id(self, id: str) -> Any:
+        return self._db.get(id, None)
+
+    def get_by_ids(self, ids: List[str], fields: List[str] = None) -> List[Any]:
+        result = []
+        for index in ids:
+            item = self._db.get(index, None)
+            if item is None:
+                result.append(None)
+                continue
+
+            if fields is None:
+                result.append(item)
+            else:
+                result.append({k: v for k, v in item.items() if k in fields})
+        return result
+
+    def get_all(self) -> Dict[str, Dict]:
+        return dict(self._db)
+
+    def filter_keys(self, data: List[str]) -> Set[str]:
+        return {s for s in data if s not in self._db}
+
+    def upsert(self, data: Dict[str, Any]):
+        left_data = {}
+        for k, v in data.items():
+            if k not in self._db:
+                left_data[k] = v
+
+        if left_data:
+            for k, v in left_data.items():
+                self._db[k] = v
+
+            # if left_data is very large, it is recommended to use self._db.write_batch() for optimization
+
+        return left_data
+
+    def drop(self):
+        self._db.close()
+        Rdict.destroy(self._db_path)
+        self._db = Rdict(self._db_path)
+        logger.info("Dropped RocksDB %s", self.namespace)
+
+    def close(self):
+        if self._db:
+            self._db.close()

graphgen/models/storage/kv/json_storage.py

@@ -27,7 +27,7 @@ def __call__(self, batch: pd.DataFrame) -> Iterable[pd.DataFrame]:
         # this operator does not consume any batch data
         # but for compatibility we keep the interface
         _ = batch.to_dict(orient="records")
-
+        self.graph_storage.reload()
         yield from self.quiz()
 
     async def _process_single_quiz(self, item: str) -> dict | None: