public Scorer scorer()

in src/main/java/org/opensearch/knn/index/KNNWeight.java [84:198]


    public Scorer scorer(LeafReaderContext context) throws IOException {
            SegmentReader reader = (SegmentReader) FilterLeafReader.unwrap(context.reader());
            String directory = ((FSDirectory) FilterDirectory.unwrap(reader.directory())).getDirectory().toString();

            FieldInfo fieldInfo = reader.getFieldInfos().fieldInfo(knnQuery.getField());

            if (fieldInfo == null) {
                logger.debug("[KNN] Field info not found for {}:{}", knnQuery.getField(),
                        reader.getSegmentName());
                return null;
            }

            KNNEngine knnEngine;
            SpaceType spaceType;

            // Check if a modelId exists. If so, the space type and engine will need to be picked up from the model's
            // metadata.
            String modelId = fieldInfo.getAttribute(MODEL_ID);
            if (modelId != null) {
                ModelMetadata modelMetadata = modelDao.getMetadata(modelId);
                if (modelMetadata == null) {
                    throw new RuntimeException("Model \"" + modelId + "\" does not exist.");
                }

                knnEngine = modelMetadata.getKnnEngine();
                spaceType = modelMetadata.getSpaceType();
            } else {
                String engineName = fieldInfo.attributes().getOrDefault(KNN_ENGINE, KNNEngine.NMSLIB.getName());
                knnEngine = KNNEngine.getEngine(engineName);
                String spaceTypeName = fieldInfo.attributes().getOrDefault(SPACE_TYPE, SpaceType.L2.getValue());
                spaceType = SpaceType.getSpace(spaceTypeName);
            }

            /*
             * In case of compound file, extension would be <engine-extension> + c otherwise <engine-extension>
             */
            String engineExtension = reader.getSegmentInfo().info.getUseCompoundFile()
                    ? knnEngine.getExtension() + KNNConstants.COMPOUND_EXTENSION : knnEngine.getExtension();
            String engineSuffix = knnQuery.getField() + engineExtension;
            List<String> engineFiles = reader.getSegmentInfo().files().stream()
                    .filter(fileName -> fileName.endsWith(engineSuffix))
                    .collect(Collectors.toList());

            if(engineFiles.isEmpty()) {
                logger.debug("[KNN] No engine index found for field {} for segment {}",
                        knnQuery.getField(), reader.getSegmentName());
                return null;
            }

            Path indexPath = PathUtils.get(directory, engineFiles.get(0));
            final KNNQueryResult[] results;
            KNNCounter.GRAPH_QUERY_REQUESTS.increment();

            // We need to first get index allocation
            NativeMemoryAllocation indexAllocation = null;

            Map<String, Object> loadParameters = new HashMap<String, Object>() {{
                put(SPACE_TYPE, spaceType.getValue());
            }};

            // For nmslib, we set efSearch from an index setting. This has the advantage of being able to dynamically
            // update this value, which we cannot do at the moment for mapping parameters.
            if (knnEngine.equals(KNNEngine.NMSLIB)) {
                loadParameters.put(KNNConstants.HNSW_ALGO_EF_SEARCH, KNNSettings.getEfSearchParam(knnQuery.getIndexName()));
            }

            try {
                indexAllocation = nativeMemoryCacheManager.get(
                        new NativeMemoryEntryContext.IndexEntryContext(
                                indexPath.toString(),
                                NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance(),
                                loadParameters,
                                knnQuery.getIndexName()
                        ), true);
            } catch (ExecutionException e) {
                GRAPH_QUERY_ERRORS.increment();
                throw new RuntimeException(e);
            }

            // Now that we have the allocation, we need to readLock it
            indexAllocation.readLock();

            try {
                if (indexAllocation.isClosed()) {
                    throw new RuntimeException("Index has already been closed");
                }

                results = JNIService.queryIndex(indexAllocation.getMemoryAddress(), knnQuery.getQueryVector(), knnQuery.getK(), knnEngine.getName());
            } catch (Exception e) {
                GRAPH_QUERY_ERRORS.increment();
                throw new RuntimeException(e);
            } finally {
                indexAllocation.readUnlock();
            }

            /*
             * Scores represent the distance of the documents with respect to given query vector.
             * Lesser the score, the closer the document is to the query vector.
             * Since by default results are retrieved in the descending order of scores, to get the nearest
             * neighbors we are inverting the scores.
             */
            if (results.length == 0) {
                logger.debug("[KNN] Query yielded 0 results");
                return null;
            }

            Map<Integer, Float> scores = Arrays.stream(results).collect(
                    Collectors.toMap(KNNQueryResult::getId, result -> knnEngine.score(result.getScore(), spaceType)));
            int maxDoc = Collections.max(scores.keySet()) + 1;
            DocIdSetBuilder docIdSetBuilder = new DocIdSetBuilder(maxDoc);
            DocIdSetBuilder.BulkAdder setAdder = docIdSetBuilder.grow(maxDoc);
            Arrays.stream(results).forEach(result -> setAdder.add(result.getId()));
            DocIdSetIterator docIdSetIter = docIdSetBuilder.build().iterator();
            return new KNNScorer(this, docIdSetIter, scores, boost);
    }