kNN vector rescoring for quantized vectors

[carlosdelest]

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

[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]

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]

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?

[carlosdelest]

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]

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.

[carlosdelest]

@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!

[benwtrent]

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]

the knn query should also be a profiled query

[carlosdelest]

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]

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

[benwtrent]

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

[benwtrent]

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

[carlosdelest]

@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!

[benwtrent]

Some minor things. but this looks good!

[benwtrent]

🤯

[benwtrent]

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

[elasticsearchmachine]

💔 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]

💚 All backports created successfully

Status	Branch	Result
✅	8.x

Questions ?

Please refer to the Backport tool documentation