Knn vector rescoring to sort score docs

[javanna]

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

[elasticsearchmachine]

Pinging @elastic/es-search-relevance (Team:Search Relevance)

[elasticsearchmachine]

Hi @javanna, I've created a changelog YAML for you.

[javanna]

This looks like a pretty bad one. It had nothing to do with profiling although it surfaced from a DFSProfilerIT test failure. I could reproduce it in DFSProfilerIT removing profiling, trimming down to a single shard, and disabling inter-segment concurrency, via the following code using 7E03FF959CA490B6 as seed. I assume we are missing some test coverage, cause bugs like these should be quickly and more easily caught? Was knn vector rescoring providing correct results at all? We were effectively doing binary search against an unsorted array there....

    public void testProfileDfs() throws Exception {
        String textField = "text_field";
        String numericField = "number";
        String vectorField = "vector";
        String indexName = "text-dfs-profile";
        createIndex(indexName, vectorField);
        ensureGreen();

        int numDocs = randomIntBetween(10, 50);
        IndexRequestBuilder[] docs = new IndexRequestBuilder[numDocs];
        for (int i = 0; i < numDocs; i++) {
            docs[i] = prepareIndex(indexName).setId(String.valueOf(i))
                .setSource(
                    textField,
                    English.intToEnglish(i),
                    numericField,
                    i,
                    vectorField,
                    new float[] { randomFloat(), randomFloat(), randomFloat() }
                );
        }
        indexRandom(true, docs);
        refresh();

        float[] queryVector = new float[]{0.2956158f, 0.67327607f, 0.67715585f};
        int k = 10;
        RescoreVectorBuilder rescoreVectorBuilder = new RescoreVectorBuilder(1.8569797f);
        KnnSearchBuilder knnSearchBuilder = new KnnSearchBuilder(
            vectorField,
            queryVector,
            k,
            50,
            rescoreVectorBuilder, //required to fail the test
            null
        );

        assertResponse(
            prepareSearch()
                .setTrackTotalHits(true) //required to fail the test
                .setSearchType(SearchType.DFS_QUERY_THEN_FETCH)
                .setKnnSearch(randomList(2, 5, () -> knnSearchBuilder)),
            response -> {
            }
        );
    }

[benwtrent]

good catch!

[elasticsearchmachine]

💔 Backport failed

Status	Branch	Result
✅	9.0
❌	8.18	Commit could not be cherrypicked due to conflicts
❌	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 122653

[carlosdelest]

Thanks for the fix @javanna . This is a nasty one.

I've opened #122723 as a follow up for adding test coverage.