Skip to content

fix: return doclist of same type as input index in find and findbatched #1470

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
JoanFM merged 5 commits into from
Apr 27, 2023
Merged

fix: return doclist of same type as input index in find and findbatched #1470

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
c75d6e0
fix: return doclist of same type as input index in find and findbatched
Apr 27, 2023
3740d8f
test: add assertion to find_batched test
Apr 27, 2023
a1dbfda
fix: index to int
Apr 27, 2023
af9827f
fix: apply suggestions from code review
Apr 27, 2023
923d450
Merge branch 'main' into fix-doclist-schema-in-find
Apr 27, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions docarray/utils/find.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -213,10 +213,11 @@ class MyDocument(BaseDoc):
for _, (indices_per_query, scores_per_query) in enumerate(
zip(top_indices, top_scores)
):
docs_per_query: DocList = DocList([])
doc_type = cast(Type[BaseDoc], index.doc_type)
docs_per_query: DocList = DocList.__class_getitem__(doc_type)()
for idx in indices_per_query: # workaround until #930 is fixed
docs_per_query.append(index[idx])
batched_docs.append(DocList(docs_per_query))
docs_per_query.append(index[int(idx)])
batched_docs.append(docs_per_query)
scores.append(scores_per_query)
return FindResultBatched(documents=batched_docs, scores=scores)

Expand Down
4 changes: 4 additions & 0 deletions tests/units/util/test_find.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,8 @@ def test_find_torch(random_torch_query, random_torch_index, metric):
)
assert len(top_k) == 7
assert len(scores) == 7
assert top_k.doc_type == random_torch_index.doc_type

if metric.endswith('_dist'):
assert (torch.stack(sorted(scores)) == scores).all()
else:
Expand Down Expand Up @@ -151,6 +153,8 @@ def test_find_batched_torch(random_torch_batch_query, random_torch_index, metric
for top_k, top_scores in zip(documents, scores):
assert len(top_k) == 7
assert len(top_scores) == 7
assert top_k.doc_type == random_torch_index.doc_type

for sc in scores:
if metric.endswith('_dist'):
assert (torch.stack(sorted(sc)) == sc).all()
Expand Down