Skip to content

Commit 3b0922a

Browse files
author
Joan Fontanals Martinez
committed
Merge branch 'get-current-count-hnswlib' of https://github.com/docarray/docarray into refactor-hnswlib-performance
2 parents fa9b9ce + a91e23e commit 3b0922a

File tree

2 files changed

+14
-14
lines changed

2 files changed

+14
-14
lines changed

‎tests/index/hnswlib/test_index_get_del.py‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -211,7 +211,7 @@ def test_get_single(ten_simple_docs, ten_flat_docs, ten_nested_docs, tmp_path):
211211
for d in ten_simple_docs:
212212
id_ = d.id
213213
assert index[id_].id == id_
214-
assert np.all(index[id_].tens == d.tens)
214+
assert np.allclose(index[id_].tens == d.tens)
215215

216216
# flat
217217
index = HnswDocumentIndex[FlatDoc](work_dir=str(flat_path))
@@ -221,8 +221,8 @@ def test_get_single(ten_simple_docs, ten_flat_docs, ten_nested_docs, tmp_path):
221221
for d in ten_flat_docs:
222222
id_ = d.id
223223
assert index[id_].id == id_
224-
assert np.all(index[id_].tens_one == d.tens_one)
225-
assert np.all(index[id_].tens_two == d.tens_two)
224+
assert np.allclose(index[id_].tens_one == d.tens_one)
225+
assert np.allclose(index[id_].tens_two == d.tens_two)
226226

227227
# nested
228228
index = HnswDocumentIndex[NestedDoc](work_dir=str(nested_path))
@@ -233,7 +233,7 @@ def test_get_single(ten_simple_docs, ten_flat_docs, ten_nested_docs, tmp_path):
233233
id_ = d.id
234234
assert index[id_].id == id_
235235
assert index[id_].d.id == d.d.id
236-
assert np.all(index[id_].d.tens == d.d.tens)
236+
assert np.allclose(index[id_].d.tens == d.d.tens)
237237

238238

239239
def test_get_multiple(ten_simple_docs, ten_flat_docs, ten_nested_docs, tmp_path):
@@ -252,7 +252,7 @@ def test_get_multiple(ten_simple_docs, ten_flat_docs, ten_nested_docs, tmp_path)
252252
retrieved_docs = index[ids_to_get]
253253
for id_, d_in, d_out in zip(ids_to_get, docs_to_get, retrieved_docs):
254254
assert d_out.id == id_
255-
assert np.all(d_out.tens == d_in.tens)
255+
assert np.allclose(d_out.tens == d_in.tens)
256256

257257
# flat
258258
index = HnswDocumentIndex[FlatDoc](work_dir=str(flat_path))
@@ -264,8 +264,8 @@ def test_get_multiple(ten_simple_docs, ten_flat_docs, ten_nested_docs, tmp_path)
264264
retrieved_docs = index[ids_to_get]
265265
for id_, d_in, d_out in zip(ids_to_get, docs_to_get, retrieved_docs):
266266
assert d_out.id == id_
267-
assert np.all(d_out.tens_one == d_in.tens_one)
268-
assert np.all(d_out.tens_two == d_in.tens_two)
267+
assert np.allclose(d_out.tens_one == d_in.tens_one)
268+
assert np.allclose(d_out.tens_two == d_in.tens_two)
269269

270270
# nested
271271
index = HnswDocumentIndex[NestedDoc](work_dir=str(nested_path))
@@ -278,7 +278,7 @@ def test_get_multiple(ten_simple_docs, ten_flat_docs, ten_nested_docs, tmp_path)
278278
for id_, d_in, d_out in zip(ids_to_get, docs_to_get, retrieved_docs):
279279
assert d_out.id == id_
280280
assert d_out.d.id == d_in.d.id
281-
assert np.all(d_out.d.tens == d_in.d.tens)
281+
assert np.allclose(d_out.d.tens == d_in.d.tens)
282282

283283

284284
def test_get_key_error(ten_simple_docs, ten_flat_docs, ten_nested_docs, tmp_path):
@@ -303,7 +303,7 @@ def test_del_single(ten_simple_docs, ten_flat_docs, ten_nested_docs, tmp_path):
303303
index[id_]
304304
else:
305305
assert index[id_].id == id_
306-
assert np.all(index[id_].tens == d.tens)
306+
assert np.allclose(index[id_].tens == d.tens)
307307
# delete again
308308
del index[ten_simple_docs[3].id]
309309
assert index.num_docs() == 8
@@ -314,7 +314,7 @@ def test_del_single(ten_simple_docs, ten_flat_docs, ten_nested_docs, tmp_path):
314314
index[id_]
315315
else:
316316
assert index[id_].id == id_
317-
assert np.all(index[id_].tens == d.tens)
317+
assert np.allclose(index[id_].tens == d.tens)
318318

319319

320320
def test_del_multiple(ten_simple_docs, ten_flat_docs, ten_nested_docs, tmp_path):
@@ -333,7 +333,7 @@ def test_del_multiple(ten_simple_docs, ten_flat_docs, ten_nested_docs, tmp_path)
333333
index[doc.id]
334334
else:
335335
assert index[doc.id].id == doc.id
336-
assert np.all(index[doc.id].tens == doc.tens)
336+
assert np.allclose(index[doc.id].tens == doc.tens)
337337

338338

339339
def test_del_key_error(ten_simple_docs, ten_flat_docs, ten_nested_docs, tmp_path):

‎tests/index/hnswlib/test_persist_data.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ def test_persist_and_restore(tmp_path):
2222
query = SimpleDoc(tens=np.random.random((10,)))
2323

2424
# create index
25-
index = HnswDocumentIndex[SimpleDoc](work_dir=str(tmp_path))
25+
_ = HnswDocumentIndex[SimpleDoc](work_dir=str(tmp_path))
2626

2727
# load existing index file
2828
index = HnswDocumentIndex[SimpleDoc](work_dir=str(tmp_path))
@@ -38,7 +38,7 @@ def test_persist_and_restore(tmp_path):
3838
find_results_after = index.find(query, search_field='tens', limit=5)
3939
for doc_before, doc_after in zip(find_results_before[0], find_results_after[0]):
4040
assert doc_before.id == doc_after.id
41-
assert (doc_before.tens == doc_after.tens).all()
41+
assert np.allclose(doc_before.tens == doc_after.tens)
4242

4343
# add new data
4444
index.index([SimpleDoc(tens=np.random.random((10,))) for _ in range(5)])
@@ -70,7 +70,7 @@ def test_persist_and_restore_nested(tmp_path):
7070
find_results_after = index.find(query, search_field='d__tens', limit=5)
7171
for doc_before, doc_after in zip(find_results_before[0], find_results_after[0]):
7272
assert doc_before.id == doc_after.id
73-
assert (doc_before.tens == doc_after.tens).all()
73+
assert np.allclose(doc_before.tens == doc_after.tens)
7474

7575
# delete and restore
7676
index.index(

0 commit comments

Comments
 (0)