andras/great-ai
1
0
Fork 0
Code Issues Pull requests 1 Projects Releases 4 Packages Wiki Activity Actions

Improve TinyDB performance for *_ground_truth

Browse source
This commit is contained in:
Andras Schmelczer 2022-08-07 14:12:43 +02:00
parent f382390291
commit 58286f6f1c
No known key found for this signature in database
GPG key ID: 0EA1BC97D0AB076E
1 changed files with 7 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

19
great_ai/persistence/parallel_tinydb_driver.py
View file

@ -63,15 +63,11 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
) )
) )
documents: List[Trace] = [ documents = self._safe_execute(lambda db: db.search(does_match))
Trace.parse_obj(t)
for t in self._safe_execute(lambda db: db.search(does_match))
]
if not documents: if not documents:
return [], 0 return [], 0
df = pd.DataFrame([d.to_flat_dict() for d in documents]) df = pd.DataFrame(documents)
for f in conjunctive_filters: for f in conjunctive_filters:
operator = f.operator.lower() operator = f.operator.lower()
@ -91,10 +87,7 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
count = len(df) count = len(df)
result = df.iloc[skip:] if take is None else df.iloc[skip : skip + take] result = df.iloc[skip:] if take is None else df.iloc[skip : skip + take]
return [ return [Trace.parse_obj(trace) for _, trace in result.iterrows()], count
next(d for d in documents if d.trace_id == trace_id)
for trace_id in result["trace_id"]
], count
def update(self, id: str, new_version: Trace) -> None: def update(self, id: str, new_version: Trace) -> None:
self._safe_execute( self._safe_execute(
@ -105,8 +98,10 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
self._safe_execute(lambda db: db.remove(lambda d: d["trace_id"] == id)) self._safe_execute(lambda db: db.remove(lambda d: d["trace_id"] == id))
def delete_batch(self, ids: List[str]) -> None: def delete_batch(self, ids: List[str]) -> None:
for i in ids: with lock:
self.delete(i) with TinyDB(self.path_to_db) as db:
for id in ids:
db.remove(lambda d: d["trace_id"] == id)
def _safe_execute(self, func: Callable[[TinyDB], Any]) -> Any: def _safe_execute(self, func: Callable[[TinyDB], Any]) -> Any:
with lock: with lock: