kNN vector rescoring for quantized vectors #116663

carlosdelest · 2024-11-12T13:30:39Z

Adds a new parameter to kNN based searches (rescore_vector.num_candidates_factor) that is used to:

Multiply the number of candidates in the kNN query
Perform an approximate search in the extended candidate set
Perform an exact search over the returned results, and get the top k

This approach uses a FunctionScoreQuery to replace scores using a VectorSimilarity based DoubleValueSource.

The DenseVectorFieldMapper is the one creating the Query - we could have done the rewriting on the KnnVectorQueryBuilder, but given the new code had to access the vectorSimilarity and other internals, it felt more natural to include it in the field mapper.

The parameter has been added to knn section, knn query and knn retriever.

Usage:

GET msmarco-v2-bbq/_search
{
    "query": {
        "knn": {
            "field": "emb",
            "query_vector": [...],
            "k": 10,
            "num_candidates": 100,
            "rescore_vector": {
                "num_candidates_factor": 10.0
            }
        }
    }
}

Documentation will be added as a separate PR once this is merged.

Backlog

Oversample and rescore if its a quantized type - checking indexOptions
Use an object with a parameter on the knn query
Add to the kNN section
Add to the kNN retriever
Implement profiling
Testing

…ased DoubleValueSource

benwtrent

Yes! this is more what I was thinking. Nice use of functionscore!

A couple of things:

we should only oversample & rescore if its a quantized type (you can see this via indexOptions)
I still think the API level thing needs to be an object with a parameter.
We need to make sure this works for knn query and for the top-level KNN object :)

carlosdelest · 2024-11-12T14:33:53Z

Thanks for taking a look @benwtrent ! I was just trying to validate the overall direction. You're perfectly right in your comments, I was not aiming for that yet 🙂

I'll keep iterating on this idea and open up a proper draft and not a PoC.

We need to make sure this works for knn query and for the top-level KNN object :)

Is it OK to tackle that as separate PRs?

benwtrent · 2024-11-12T14:49:04Z

Is it OK to tackle that as separate PRs?

If necessary, for sure. But both hit the same path on the shard level, its all about being able to declare the values at the parser level and propagate them down. If you did separate PRs, you then have separate transport versions, and likely separate capabilities for testing, etc. I am not 100% sure its worth it?

# Conflicts: # server/src/main/java/org/elasticsearch/TransportVersions.java

carlosdelest · 2024-11-18T13:48:24Z

hey @benwtrent ! I'm making progress on this. I've added the rescoring oversample to the knn query, knn section and the knn retriever.

I still have to fix CI, use it just for quantized vectors, and add proper testing - and I'm already at 38 changed files. Do you still think this should go on a single PR, or should I start working on separate PRs for query, section and retriver?

benwtrent · 2024-11-18T13:57:32Z

I still have to fix CI, use it just for quantized vectors, and add proper testing - and I'm already at 38 changed files. Do you still think this should go on a single PR, or should I start working on separate PRs for query, section and retriver?

The main "logic" changes are done down on the mapper. Which all these would use. The rest is serialization & testing. I think all in one is ok.

…r of results from each shard

carlosdelest · 2024-11-20T17:08:46Z

server/src/main/java/org/elasticsearch/search/vectors/KnnRescoreVectorQuery.java

+/**
+ * Wraps a kNN vector query to rescore the results using the non-quantized vectors
+ */
+public class KnnRescoreVectorQuery extends Query implements ProfilingQuery {


@benwtrent I've created this new query to:

Wrap the KNN vector query with a function score for rescoring

Limit the number of results back from each shard when k is specified

LMKWYT!

The idea is good. Once its ready for review, I can give it a thorough look over.

benwtrent

some small things in the new query. I haven't fully reviewed. But I will once its ready

benwtrent · 2024-11-20T17:55:31Z

server/src/main/java/org/elasticsearch/search/vectors/KnnRescoreVectorQuery.java

+        byte[] byteTarget,
+        VectorSimilarityFunction vectorSimilarityFunction,
+        Integer k,
+        Query vectorQuery


the knn query should also be a profiled query

Mmmm, gotcha. There can be multiple levels of wrapping:

KnnRescoreVectorQuery

VectorSimilarityQuery

EsKnnFloatVectorQuery

We'll probably need VectorSimilarityQuery to implement ProfilingQuery as well - right now using a similarity param means that profiling returns 0 vector op counts.

I'll fix that, thanks for the catch 👍

benwtrent · 2024-11-20T17:56:28Z

server/src/main/java/org/elasticsearch/search/vectors/KnnRescoreVectorQuery.java

+
+        if (k == null) {
+            // No need to calculate top k - let the request size limit the results
+            return query;


be sure to account for the inner queries vector ops and ensure that we account for the total vector ops (quantized and non).

benwtrent · 2024-11-20T17:56:53Z

server/src/main/java/org/elasticsearch/search/vectors/KnnRescoreVectorQuery.java

+            scores[i] = scoreDocs[i].score;
+        }
+
+        vectorOpsCount = scoreDocs.length;


this seems like it should be the inner query vector ops + the total docs gathered.

benwtrent · 2024-11-20T17:57:19Z

server/src/main/java/org/elasticsearch/search/vectors/KnnRescoreVectorQuery.java

+/**
+ * Wraps a kNN vector query to rescore the results using the non-quantized vectors
+ */
+public class KnnRescoreVectorQuery extends Query implements ProfilingQuery {


The idea is good. Once its ready for review, I can give it a thorough look over.

carlosdelest · 2024-12-10T14:16:41Z

@benwtrent This is ready for another round of review.

Main changes:

Renaming oversample to num_candidates_factor and make it work with num_candidates
Changed how queries are wrapped to ensure VectorSimilarityQuery wraps the new RescoreKnnVectorQuery, so similarity is applied to the rescored scoring
Updating some YAML tests to use MIP where applicable

LMKWYT!

# Conflicts: # server/src/main/java/org/elasticsearch/TransportVersions.java

benwtrent

Some minor things. but this looks good!

...r/src/main/java/org/elasticsearch/index/mapper/vectors/VectorSimilarityFloatValueSource.java

server/src/main/java/org/elasticsearch/search/vectors/RescoreKnnVectorQuery.java

server/src/main/java/org/elasticsearch/search/vectors/RescoreVectorBuilder.java

benwtrent · 2024-12-10T19:36:52Z

server/src/test/java/org/elasticsearch/search/vectors/RescoreKnnVectorQueryTests.java

+    @ParametersFactory
+    public static Iterable<Object[]> parameters() {
+        List<Object[]> params = new ArrayList<>();
+        params.add(new Object[] { true });
+        params.add(new Object[] { false });
+
+        return params;
+    }


benwtrent · 2024-12-10T19:47:14Z

And please don't forget about adding docs (either in this PR or in a follow up).

Co-authored-by: Benjamin Trent <[email protected]>

elasticsearchmachine · 2024-12-11T08:15:39Z

💔 Backport failed

Status	Branch	Result
❌	8.x	Commit could not be cherrypicked due to conflicts

You can use sqren/backport to manually backport by running backport --upstream elastic/elasticsearch --pr 116663

carlosdelest · 2024-12-11T08:23:42Z

💚 All backports created successfully

Status	Branch	Result
✅	8.x

Questions ?

Please refer to the Backport tool documentation

(cherry picked from commit 5996772) # Conflicts: # server/src/main/java/org/elasticsearch/search/vectors/KnnSearchBuilder.java # x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRankBuilder.java

* kNN vector rescoring for quantized vectors (#116663) (cherry picked from commit 5996772) # Conflicts: # server/src/main/java/org/elasticsearch/search/vectors/KnnSearchBuilder.java # x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRankBuilder.java * FloatVectorValues have a different interface in this Lucene version

(cherry picked from commit 5996772) # Conflicts: # server/src/main/java/org/elasticsearch/search/vectors/KnnSearchBuilder.java # x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRankBuilder.java

(cherry picked from commit 5996772) # Conflicts: # server/src/main/java/org/elasticsearch/search/vectors/KnnSearchBuilder.java # x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRankBuilder.java Co-authored-by: Felix Barnsteiner <[email protected]>

RescoreKnnVectorQuery rewrites to KnnScoreDocQuery, which takes a sorted array of doc ids and corresponding array including scores fo such docs. A binary search is performed on top of the docs array, and such global ids are converted back to segment level ids (subtracting the context docbase) when scoring docs. RescoreKnnVectoryQuery did not sort the array of docs which caused binary search to return non deterministic results, which in turn made us look up wrong docs, something using out of bound ids. One symptom of this was observed in a DFSProfilerIT test failure which triggered a Lucene assertion around doc id being outside of the range of the bitset of live docs. The fix is to simply sort the score docs array before extracting docs ids and scores and providing them to KnnScoreDocQuery upon rewrite. Relates to elastic#116663 Closes elastic#119711

RescoreKnnVectorQuery rewrites to KnnScoreDocQuery, which takes a sorted array of doc ids and corresponding array including scores fo such docs. A binary search is performed on top of the docs array, and such global ids are converted back to segment level ids (subtracting the context docbase) when scoring docs. RescoreKnnVectoryQuery did not sort the array of docs which caused binary search to return non deterministic results, which in turn made us look up wrong docs, something using out of bound ids. One symptom of this was observed in a DFSProfilerIT test failure which triggered a Lucene assertion around doc id being outside of the range of the bitset of live docs. The fix is to simply sort the score docs array before extracting docs ids and scores and providing them to KnnScoreDocQuery upon rewrite. Relates to #116663 Closes #119711

RescoreKnnVectorQuery rewrites to KnnScoreDocQuery, which takes a sorted array of doc ids and corresponding array including scores fo such docs. A binary search is performed on top of the docs array, and such global ids are converted back to segment level ids (subtracting the context docbase) when scoring docs. RescoreKnnVectoryQuery did not sort the array of docs which caused binary search to return non deterministic results, which in turn made us look up wrong docs, something using out of bound ids. One symptom of this was observed in a DFSProfilerIT test failure which triggered a Lucene assertion around doc id being outside of the range of the bitset of live docs. The fix is to simply sort the score docs array before extracting docs ids and scores and providing them to KnnScoreDocQuery upon rewrite. Relates to elastic#116663 Closes elastic#119711

RescoreKnnVectorQuery rewrites to KnnScoreDocQuery, which takes a sorted array of doc ids and corresponding array including scores fo such docs. A binary search is performed on top of the docs array, and such global ids are converted back to segment level ids (subtracting the context docbase) when scoring docs. RescoreKnnVectoryQuery did not sort the array of docs which caused binary search to return non deterministic results, which in turn made us look up wrong docs, something using out of bound ids. One symptom of this was observed in a DFSProfilerIT test failure which triggered a Lucene assertion around doc id being outside of the range of the bitset of live docs. The fix is to simply sort the score docs array before extracting docs ids and scores and providing them to KnnScoreDocQuery upon rewrite. Relates to #116663 Closes #119711

Use a FunctionScoreQuery to replace scores using a VectorSimilarity b…

df06716

…ased DoubleValueSource

elasticsearchmachine added the v9.0.0 label Nov 12, 2024

carlosdelest requested a review from benwtrent November 12, 2024 13:34

carlosdelest mentioned this pull request Nov 12, 2024

PoC - Vector rescoring in kNN #116350

Closed

benwtrent reviewed Nov 12, 2024

View reviewed changes

carlosdelest changed the title ~~PoC - Vector rescoring in kNN, take 2~~ WIP - Vector rescoring Nov 13, 2024

carlosdelest added 8 commits November 13, 2024 17:36

Change API to use "rescore": {"oversample": 1.0}

be76444

Add tests

bd920c5

Fix inference module

91204a1

Merge branch 'main' into feature/knn-vector-rescore-query

6c2c1be

# Conflicts: # server/src/main/java/org/elasticsearch/TransportVersions.java

Fix knn query usage in other modules

b44ec48

Add rescore vector builder to KnnSearchBuilder

2a9e300

Add vector rescore builder to kNN retriever

4497b92

Fix refactoring, spotless

955da1f

carlosdelest added 3 commits November 18, 2024 16:53

Check oversampling is not used for quantized types

ff2c1e9

Minor refactoring to reuse KnnScoreDocQuery

bc1e5c6

Use KnnRescoreVectorQuery to perform rescoring and limiting the numbe…

a7936da

…r of results from each shard

carlosdelest commented Nov 20, 2024

View reviewed changes

benwtrent reviewed Nov 20, 2024

View reviewed changes

carlosdelest added 7 commits November 21, 2024 15:03

Small name refactoring, fix adjusting parameters

f5080a6

Add testing

39e1676

Add tests for RescoreKnnVectorQuery

9946e8d

Spotless

4fbbadd

Add test for knn retriever

0dab8ea

Add tests

257b75d

Parameterize recore knn vector query tests

81384f2

Merge branch 'main' into feature/knn-vector-rescore-query

94963fc

# Conflicts: # server/src/main/java/org/elasticsearch/TransportVersions.java

benwtrent approved these changes Dec 10, 2024

View reviewed changes

Apply suggestions from code review

a256de9

Co-authored-by: Benjamin Trent <[email protected]>

carlosdelest merged commit 5996772 into elastic:main Dec 11, 2024
15 of 16 checks passed

elasticsearchmachine added the backport pending label Dec 11, 2024

carlosdelest mentioned this pull request Dec 11, 2024

[8.x] kNN vector rescoring for quantized vectors (#116663) #118418

Merged

carlosdelest mentioned this pull request Dec 11, 2024

[Docs] kNN vector rescoring for quantized vectors #118425

Merged

carlosdelest mentioned this pull request Dec 18, 2024

Vector rescoring - Simplify code for k == null #118997

Merged

carlosdelest mentioned this pull request Jan 9, 2025

Vector rescoring oversamples k instead of num_candidates #119835

Merged

carlosdelest mentioned this pull request Jan 9, 2025

[8.x] Vector rescoring oversamples k instead of num_candidates #119887

Closed

javanna mentioned this pull request Feb 14, 2025

Knn vector rescoring to sort score docs #122653

Merged

This was referenced Feb 15, 2025

[8.x] Knn vector rescoring to sort score docs (#122653) #122678

Merged

[8.18] Knn vector rescoring to sort score docs (#122653) #122679

Merged

thecoop mentioned this pull request Mar 10, 2025

[CI] CoreWithMultipleProjectsClientYamlTestSuiteIT test {yaml=search.vectors/41_knn_search_bbq_hnsw/Vector rescoring has same scoring as exact search for kNN section} failing #124052

Closed

kNN vector rescoring for quantized vectors #116663

kNN vector rescoring for quantized vectors #116663

Uh oh!

Conversation

carlosdelest commented Nov 12, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Backlog

Uh oh!

benwtrent left a comment

Choose a reason for hiding this comment

Uh oh!

carlosdelest commented Nov 12, 2024

Uh oh!

benwtrent commented Nov 12, 2024

Uh oh!

carlosdelest commented Nov 18, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

benwtrent commented Nov 18, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

benwtrent left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

carlosdelest commented Dec 10, 2024

Uh oh!

benwtrent left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

benwtrent commented Dec 10, 2024

Uh oh!

Uh oh!

elasticsearchmachine commented Dec 11, 2024

💔 Backport failed

Uh oh!

carlosdelest commented Dec 11, 2024

💚 All backports created successfully

Questions ?

Uh oh!

Uh oh!

carlosdelest commented Nov 12, 2024 •

edited

Loading

carlosdelest commented Nov 18, 2024 •

edited

Loading