-
Notifications
You must be signed in to change notification settings - Fork 2.2k
/
test_dbms.py
45 lines (40 loc) · 1.41 KB
/
test_dbms.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
from jina.drivers.index import DBMSIndexDriver
from jina.executors.indexers.dbms import BaseDBMSIndexer
from jina.executors.indexers.dbms.keyvalue import KeyValueDBMSIndexer
from tests import get_documents
def test_dbms_keyvalue(tmpdir, test_metas):
docs = list(get_documents(chunks=False, nr=10, same_content=True))
ids, vecs, meta = zip(
*[
(
doc.id,
doc.embedding,
DBMSIndexDriver._doc_without_embedding(doc).SerializeToString(),
)
for doc in docs
]
)
save_path = None
with KeyValueDBMSIndexer(index_filename='dbms', metas=test_metas) as indexer:
indexer.add(ids, vecs, meta)
assert indexer.size == len(docs)
save_path = indexer.save_abspath
new_docs = list(get_documents(chunks=False, nr=10, same_content=False))
ids, vecs, meta = zip(
*[
(
doc.id,
doc.embedding,
DBMSIndexDriver._doc_without_embedding(doc).SerializeToString(),
)
for doc in new_docs
]
)
# assert contents update
with BaseDBMSIndexer.load(save_path) as indexer:
indexer.update(ids, vecs, meta)
assert indexer.size == len(docs)
# assert contents update
with BaseDBMSIndexer.load(save_path) as indexer:
indexer.delete([d.id for d in docs])
assert indexer.size == 0