perf: Vectorize fuzzy_lookup_embedding with numpy ops#234
Merged
gvanrossum merged 4 commits intomicrosoft:mainfrom Apr 11, 2026
Merged
perf: Vectorize fuzzy_lookup_embedding with numpy ops#234gvanrossum merged 4 commits intomicrosoft:mainfrom
gvanrossum merged 4 commits intomicrosoft:mainfrom
Conversation
Replace Python-level iteration + sort with numpy operations: - No-predicate path: np.flatnonzero for score filtering, np.argpartition for O(n) top-k selection — avoids building ScoredInt for every vector - Predicate path: numpy pre-filters by score, applies predicate only to candidates above threshold - Subset lookup: numpy fancy indexing computes dot products only for subset indices instead of delegating to full-vector scan with predicate
- Add pytest-benchmark to dev dependency group so CI has the benchmark fixture available - Replace hand-rolled StubEmbeddingModel with create_test_embedding_model() to satisfy IEmbeddingModel protocol (fixes pyright)
gvanrossum
approved these changes
Apr 11, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
fuzzy_lookup_embeddingwith numpy vectorized operations (np.flatnonzero,np.argpartition)fuzzy_lookup_embedding_in_subsetto compute dot products only for subset indices instead of scanning all vectors + predicate filterBenchmark (Azure Standard_D2s_v5, 384-dim embeddings)
fuzzy_lookup_embedding(1K vecs)fuzzy_lookup_embedding(10K vecs)fuzzy_lookup_embedding_in_subset(1K of 10K)Why this matters
These functions are called on every
fuzzy_lookup— the core search path for conversation queries. At 10K vectors (a long conversation), the Python iteration path takes 6ms per query. With multiple queries per request (e.g., searching across properties + timestamps + topics), this adds up.The numpy path stays in C for the heavy lifting: score filtering via
np.flatnonzero, O(n) top-k vianp.argpartition, and subset dot products via fancy indexing.ScoredIntobjects are only created for the final top-k results.Test plan
make format check testpasses (470 passed, 12 pre-existing online test failures)tests/benchmarks/test_benchmark_vectorbase.py