[RFC] Multiple inner hits for nested field

[heemin32]

Overview

Following the launch of multi-vector support in k-NN nested fields, search results with inner hits currently return only the single highest-scoring nested document per parent document. However, many users prefer to see the scores of all nested documents within inner hits, allowing them to display all relevant nested documents rather than just the top-scoring one.

Related GitHub Issue

• [FEATURE] inner_hits in nested neural query should return all the chunks #2113
• [FEATURE] add support for more than one kNN query on nested vectors with multiple inner hits and filter #1768

Work Items

• POC for Lucene
• Adding new field in QueryShardContext
• Implementation for faiss
• Implementation for lucene (Maybe after [FEATURE] Improving Lucene Engine Query Performance by reducing number of times a single Lucene k-NN query gets executed #2115 is completed?)

Current state

Let’s try to understand the issue with example. First we are ingesting a document with two nested documents each of which has vector [1, 1, 1], and vector [2, 2, 2].

PUT /my-knn-index-1/_doc/1
{"nested_field":[{"my_vector":[1,1,1]}, {"my_vector":[2,2,2]}]}

Then, we execute the search with inner_hits parameter.

GET /my-knn-index-1/_search
{
  "query": {
    "nested": {
      "path": "nested_field",
      "query": {
        "knn": {
          "nested_field.my_vector": {
            "vector": [1,1,1],
            "k": 2
          }
        }
      },
      "inner_hits": {}
    }
  }
}

The search result returns only single nested document with vector [1, 1, 1]

{
    "_index": "my-knn-index-1",
    "_id": "1",
    "_score": 1.0,
    "_source": {
        "nested_field": [
        {
            "my_vector": [
            1,
            1,
            1
            ]
        },
        {
            "my_vector": [
            2,
            2,
            2
            ]
        }
        ]
    },
    "inner_hits": {
        "nested_field": {
        "hits": {
            "total": {
            "value": 1,
            "relation": "eq"
            },
            "max_score": 1.0,
            "hits": [
            {
                "_index": "my-knn-index-1",
                "_id": "1",
                "_nested": {
                "field": "nested_field",
                "offset": 0
                },
                "_score": 1.0,
                "_source": {
                "my_vector": [
                    1,
                    1,
                    1
                ]
                }
            }
            ]
        }
      }
    }
}

Desired state

We want to have all nested documents inside inner_hits block for each returned parent documents.

"inner_hits": {
    "nested_field": {
    "hits": {
        "total": {
        "value": 2,
        "relation": "eq"
        },
        "max_score": 1.0,
        "hits": [
        {
            "_index": "my-knn-index-1",
            "_id": "1",
            "_nested": {
            "field": "nested_field",
            "offset": 0
            },
            "_score": 1.0,
            "_source": {
            "my_vector": [
                1,
                1,
                1
            ]
            }
        },
       {
            "_index": "my-knn-index-1",
            "_id": "1",
            "_nested": {
            "field": "nested_field",
            "offset": 1
            },
            "_score": 0.25,
            "_source": {
            "my_vector": [
                2,
                2,
                2
            ]
            }
        }
        ]
    }
   }
}

Class diagram

For Faiss

For faiss, we need to migrate the entire query to NativeEngineKnnVectorQuery from KNNQuery. Currently, we use NativeEngineKnnVectorQuery only when the rescoring context exist.

Inside createWeight method of NativeEngineKnnVectorQuery, if it is nested field, we do another exact search with filtering on topK parentIds to retrieve all child documents with its score.

public class NativeEngineKnnVectorQuery extends Query {

    private final KNNQuery knnQuery;

    // Update on existing method
    @Override
    public Weight createWeight(IndexSearcher indexSearcher, ScoreMode scoreMode, float boost) throws IOException {
        final IndexReader reader = indexSearcher.getIndexReader();
        final KNNWeight knnWeight = (KNNWeight) knnQuery.createWeight(indexSearcher, scoreMode, 1);
        List<LeafReaderContext> leafReaderContexts = reader.leaves();
        List<Map<Integer, Float>> perLeafResults;
        RescoreContext rescoreContext = knnQuery.getRescoreContext();
        final int finalK = knnQuery.getK();
        if (rescoreContext == null) {
            perLeafResults = doSearch(indexSearcher, leafReaderContexts, knnWeight, finalK);
        } else {
            boolean isShardLevelRescoringEnabled = KNNSettings.isShardLevelRescoringEnabledForDiskBasedVector(knnQuery.getIndexName());
            int dimension = knnQuery.getQueryVector().length;
            int firstPassK = rescoreContext.getFirstPassK(finalK, isShardLevelRescoringEnabled, dimension);
            perLeafResults = doSearch(indexSearcher, leafReaderContexts, knnWeight, firstPassK);
            if (isShardLevelRescoringEnabled == true) {
                ResultUtil.reduceToTopK(perLeafResults, firstPassK);
            }
    
            StopWatch stopWatch = new StopWatch().start();
            perLeafResults = doRescore(indexSearcher, leafReaderContexts, knnWeight, perLeafResults, finalK);
            long rescoreTime = stopWatch.stop().totalTime().millis();
            log.debug("Rescoring results took {} ms. oversampled k:{}, segments:{}", rescoreTime, firstPassK, leafReaderContexts.size());
        }
        ResultUtil.reduceToTopK(perLeafResults, finalK);
        **// if it is innerHit query
            // get parents doc ids from the topK results
            // do exact search using the parent doc ids as filtered ids
            // return the results of the exact search** 
        TopDocs[] topDocs = new TopDocs[perLeafResults.size()];
        for (int i = 0; i < perLeafResults.size(); i++) {
            topDocs[i] = ResultUtil.resultMapToTopDocs(perLeafResults.get(i), leafReaderContexts.get(i).docBase);
        }
    
        TopDocs topK = TopDocs.merge(knnQuery.getK(), topDocs);
        if (topK.scoreDocs.length == 0) {
            return new MatchNoDocsQuery().createWeight(indexSearcher, scoreMode, boost);
        }
        return createDocAndScoreQuery(reader, topK).createWeight(indexSearcher, scoreMode, boost);
    }
}

For Lucene

For lucene, we will create our own wrapper query around lucene knn query as @navneet1v suggested in this GitHub issue. #2115
The rest of the code logic will be similar to NativeEngineKnnVectorQuery where we do another exact search on topK parentIds to retrieve all child documents with its score.

// This is a new class need to be added
public class LuceneEngineKNNQuery extends Query {
    Query luceneQuery; // this can be KnnFloatVectorQuery, KnnByteVectorQuery, DiversifyingChildrenByteKnnVectorQuery etc.

        @Override
    public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
        // rewrite the lucene query
        Query docAndScoreQuery = searcher.rewrite(luceneQuery);
    
        Weight weight = docAndScoreQuery.createWeight(searcher, scoreMode, boost);
        // get scorer for each segments
        for each leafReaderContext
            List<Scorer> scorers = weight.scorer(leafReaderContext);
        
        // get perLeafResults from scorers
        List<Map<Integer, Float>> perLeafResults;
        perLeafResults = getPerLeafResults(scorers);
        
        // reduce the result to topK
        ResultUtil.reduceToTopK(perLeafResults, firstPassK);
        
        // if it is innerHit query
            // get parents doc ids from the topK results
            // do exact search using the parent doc ids as filtered ids
            // return the results of the exact search
    }
}

Separating nested field query with innerHit query

Currently, KNNQueryBuilder.doToQuery(QueryShardContext) is called twice when innerHit is requested: one for nested field query and another for inner hit query. We want to add new field in QueryShardContext to distinguish if this is for nested field query or inner hit query so that we can retrieve all child documents only for inner hit query. The reason of not retrieving all child documents for nested field query is to avoid a degradation on latency. This is a two way door where we can support retrieving all child documents for nested field query if we want to support different scoring mode other than max. #1743

[navneet1v]

We want to add new field in QueryShardContext to distinguish if this is for nested field query or inner hit query so that we can retrieve all child documents only for inner hit query.

@heemin32 do we have details around this what will be the parameter and what will be the shape of that parameter?

[jmazanec15]

@heemin32 So for quantization (assuming no re-scoring) would this mean that some of the results have approximated scores and some have full precision scores? Or would they all have full precision scores? If the latter is the case, then would we need to update the overall docs score?

[heemin32]

We want to add new field in QueryShardContext to distinguish if this is for nested field query or inner hit query so that we can retrieve all child documents only for inner hit query.

@heemin32 do we have details around this what will be the parameter and what will be the shape of that parameter?

Something like private boolean isInnerHit;

[heemin32]

We want to add new field in QueryShardContext to distinguish if this is for nested field query or inner hit query so that we can retrieve all child documents only for inner hit query.

@heemin32 do we have details around this what will be the parameter and what will be the shape of that parameter?

@heemin32 So for quantization (assuming no re-scoring) would this mean that some of the results have approximated scores and some have full precision scores? Or would they all have full precision scores? If the latter is the case, then would we need to update the overall docs score?

They all will have full precision scores. I was going to accept the score difference to avoid latency degradation but we can discuss about it. If we don't want the score to be different between those two, we should do this retrieval for both nested field query and inner hit query.