Add KNN Profiling Info to Core Profiler - Lucene

CHANGELOG.md

@@ -13,3 +13,4 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 * Asymmetric Distance Computation for binary quantized faiss indices [#2733](https://github.com/opensearch-project/k-NN/pull/2733)
 * [BUGFIX] [Remote Vector Index Build] Don't fall back to CPU on terminal failures [#2773](https://github.com/opensearch-project/k-NN/pull/2773)
 * Add KNN timing info to core profiler [#2785](https://github.com/opensearch-project/k-NN/pull/2785)
+* Add KNN timing info for lucene queries [#2802](https://github.com/opensearch-project/k-NN/pull/2802)

src/main/java/org/opensearch/knn/index/query/KNNQueryFactory.java

@@ -7,14 +7,14 @@
 
 import lombok.NonNull;
 import lombok.extern.log4j.Log4j2;
-import org.apache.lucene.search.KnnByteVectorQuery;
-import org.apache.lucene.search.KnnFloatVectorQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.join.BitSetProducer;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.knn.index.VectorDataType;
 import org.opensearch.knn.index.engine.KNNEngine;
 import org.opensearch.knn.index.query.common.QueryUtils;
+import org.opensearch.knn.index.query.lucene.ProfileKnnByteVectorQuery;
+import org.opensearch.knn.index.query.lucene.ProfileKnnFloatVectorQuery;
 import org.opensearch.knn.index.query.lucenelib.NestedKnnVectorQueryFactory;
 import org.opensearch.knn.index.query.lucene.LuceneEngineKnnVectorQuery;
 import org.opensearch.knn.index.query.nativelib.NativeEngineKnnVectorQuery;
@@ -179,8 +179,8 @@ private static Query getKnnVectorQuery(
         if (parentFilter == null) {
             assert expandNested == false : "expandNested is allowed to be true only for nested fields.";
             return vectorDataType == VectorDataType.FLOAT
-                ? new KnnFloatVectorQuery(fieldName, floatQueryVector, k, filterQuery)
-                : new KnnByteVectorQuery(fieldName, byteQueryVector, k, filterQuery);
+                ? new ProfileKnnFloatVectorQuery(fieldName, floatQueryVector, k, filterQuery)
+                : new ProfileKnnByteVectorQuery(fieldName, byteQueryVector, k, filterQuery);
         }
         // If parentFilter is not null, it is a nested query. Therefore, we delegate creation of query to {@link
         // NestedKnnVectorQueryFactory}

src/main/java/org/opensearch/knn/index/query/common/QueryUtils.java

@@ -20,6 +20,9 @@
 import org.apache.lucene.util.Bits;
 import org.opensearch.knn.index.query.KNNWeight;
 import org.opensearch.knn.index.query.iterators.GroupedNestedDocIdSetIterator;
+import org.opensearch.knn.profile.KNNProfileUtil;
+import org.opensearch.knn.profile.query.KNNQueryTimingType;
+import org.opensearch.search.profile.ContextualProfileBreakdown;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -89,6 +92,35 @@ private int[] findSegmentStarts(final IndexReader reader, final int[] docs) {
         return starts;
     }
 
+    /**
+     * Performs the search in parallel with profiling.
+     *
+     * @param indexSearcher the index searcher
+     * @param leafReaderContexts the leaf reader contexts
+     * @param weight the search weight
+     * @return a list of maps, each mapping document IDs to their scores
+     * @throws IOException
+     */
+    public List<Map<Integer, Float>> doSearch(
+        final IndexSearcher indexSearcher,
+        final List<LeafReaderContext> leafReaderContexts,
+        final Weight weight,
+        ContextualProfileBreakdown profile
+    ) throws IOException {
+        List<Callable<Map<Integer, Float>>> tasks = new ArrayList<>(leafReaderContexts.size());
+        for (LeafReaderContext leafReaderContext : leafReaderContexts) {
+            tasks.add(
+                () -> (Map<Integer, Float>) KNNProfileUtil.profileBreakdown(
+                    profile,
+                    leafReaderContext,
+                    KNNQueryTimingType.ANN_SEARCH,
+                    () -> searchLeaf(leafReaderContext, weight)
+                )
+            );
+        }
+        return indexSearcher.getTaskExecutor().invokeAll(tasks);
+    }
+
     /**
      * Performs the search in parallel.
      *

src/main/java/org/opensearch/knn/index/query/lucene/ProfileDiversifyingChildrenByteKnnVectorQuery.java

@@ -0,0 +1,74 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.index.query.lucene;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.QueryTimeout;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.search.join.BitSetProducer;
+import org.apache.lucene.search.join.DiversifyingChildrenByteKnnVectorQuery;
+import org.apache.lucene.search.knn.KnnCollectorManager;
+import org.apache.lucene.util.Bits;
+import org.opensearch.knn.profile.KNNProfileUtil;
+import org.opensearch.knn.profile.query.KNNQueryTimingType;
+import org.opensearch.search.profile.query.QueryProfiler;
+
+import java.io.IOException;
+
+/**
+ * Wrapper class used for profiling {@link DiversifyingChildrenByteKnnVectorQuery}
+ */
+public class ProfileDiversifyingChildrenByteKnnVectorQuery extends DiversifyingChildrenByteKnnVectorQuery {
+
+    private QueryProfiler profiler;
+
+    public ProfileDiversifyingChildrenByteKnnVectorQuery(
+        String field,
+        byte[] target,
+        Query childFilter,
+        int k,
+        BitSetProducer parentsFilter
+    ) {
+        super(field, target, childFilter, k, parentsFilter);
+    }
+
+    @Override
+    public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+        profiler = KNNProfileUtil.getProfiler(indexSearcher);
+        return super.rewrite(indexSearcher);
+    }
+
+    @Override
+    protected TopDocs approximateSearch(
+        LeafReaderContext context,
+        Bits acceptDocs,
+        int visitedLimit,
+        KnnCollectorManager knnCollectorManager
+    ) throws IOException {
+        return (TopDocs) KNNProfileUtil.profile(profiler, this, context, KNNQueryTimingType.ANN_SEARCH, () -> {
+            try {
+                return super.approximateSearch(context, acceptDocs, visitedLimit, knnCollectorManager);
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        });
+    }
+
+    @Override
+    protected TopDocs exactSearch(LeafReaderContext context, DocIdSetIterator acceptIterator, QueryTimeout queryTimeout)
+        throws IOException {
+        return (TopDocs) KNNProfileUtil.profile(profiler, this, context, KNNQueryTimingType.EXACT_SEARCH, () -> {
+            try {
+                return super.exactSearch(context, acceptIterator, queryTimeout);
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        });
+    }
+}

src/main/java/org/opensearch/knn/index/query/lucene/ProfileDiversifyingChildrenFloatKnnVectorQuery.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.index.query.lucene;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.QueryTimeout;
+import org.apache.lucene.search.*;
+import org.apache.lucene.search.join.BitSetProducer;
+import org.apache.lucene.search.join.DiversifyingChildrenFloatKnnVectorQuery;
+import org.apache.lucene.search.knn.KnnCollectorManager;
+import org.apache.lucene.util.Bits;
+import org.opensearch.knn.profile.KNNProfileUtil;
+import org.opensearch.knn.profile.query.KNNQueryTimingType;
+import org.opensearch.search.profile.query.QueryProfiler;
+
+import java.io.IOException;
+
+/**
+ * Wrapper class used for profiling {@link DiversifyingChildrenFloatKnnVectorQuery}
+ */
+public class ProfileDiversifyingChildrenFloatKnnVectorQuery extends DiversifyingChildrenFloatKnnVectorQuery {
+
+    private QueryProfiler profiler;
+
+    public ProfileDiversifyingChildrenFloatKnnVectorQuery(
+        String field,
+        float[] target,
+        Query childFilter,
+        int k,
+        BitSetProducer parentsFilter
+    ) {
+        super(field, target, childFilter, k, parentsFilter);
+    }
+
+    @Override
+    public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+        profiler = KNNProfileUtil.getProfiler(indexSearcher);
+        return super.rewrite(indexSearcher);
+    }
+
+    @Override
+    protected TopDocs approximateSearch(
+        LeafReaderContext context,
+        Bits acceptDocs,
+        int visitedLimit,
+        KnnCollectorManager knnCollectorManager
+    ) throws IOException {
+        return (TopDocs) KNNProfileUtil.profile(profiler, this, context, KNNQueryTimingType.ANN_SEARCH, () -> {
+            try {
+                return super.approximateSearch(context, acceptDocs, visitedLimit, knnCollectorManager);
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        });
+    }
+
+    @Override
+    protected TopDocs exactSearch(LeafReaderContext context, DocIdSetIterator acceptIterator, QueryTimeout queryTimeout)
+        throws IOException {
+        return (TopDocs) KNNProfileUtil.profile(profiler, this, context, KNNQueryTimingType.EXACT_SEARCH, () -> {
+            try {
+                return super.exactSearch(context, acceptIterator, queryTimeout);
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        });
+    }
+}

src/main/java/org/opensearch/knn/index/query/lucene/ProfileKnnByteVectorQuery.java

@@ -0,0 +1,64 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.index.query.lucene;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.QueryTimeout;
+import org.apache.lucene.search.*;
+import org.apache.lucene.search.knn.KnnCollectorManager;
+import org.apache.lucene.util.Bits;
+import org.opensearch.knn.profile.KNNProfileUtil;
+import org.opensearch.knn.profile.query.KNNQueryTimingType;
+import org.opensearch.search.profile.query.QueryProfiler;
+
+import java.io.IOException;
+
+/**
+ * Wrapper class used for profiling {@link KnnByteVectorQuery}
+ */
+public class ProfileKnnByteVectorQuery extends KnnByteVectorQuery {
+
+    private QueryProfiler profiler;
+
+    public ProfileKnnByteVectorQuery(String field, byte[] target, int k, Query filter) {
+        super(field, target, k, filter);
+    }
+
+    @Override
+    public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+        profiler = KNNProfileUtil.getProfiler(indexSearcher);
+        return super.rewrite(indexSearcher);
+    }
+
+    @Override
+    protected TopDocs approximateSearch(
+        LeafReaderContext context,
+        Bits acceptDocs,
+        int visitedLimit,
+        KnnCollectorManager knnCollectorManager
+    ) throws IOException {
+        return (TopDocs) KNNProfileUtil.profile(profiler, this, context, KNNQueryTimingType.ANN_SEARCH, () -> {
+            try {
+                return super.approximateSearch(context, acceptDocs, visitedLimit, knnCollectorManager);
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        });
+    }
+
+    @Override
+    protected TopDocs exactSearch(LeafReaderContext context, DocIdSetIterator acceptIterator, QueryTimeout queryTimeout)
+        throws IOException {
+        return (TopDocs) KNNProfileUtil.profile(profiler, this, context, KNNQueryTimingType.EXACT_SEARCH, () -> {
+            try {
+                return super.exactSearch(context, acceptIterator, queryTimeout);
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        });
+    }
+
+}

src/main/java/org/opensearch/knn/index/query/lucene/ProfileKnnFloatVectorQuery.java

@@ -0,0 +1,64 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.index.query.lucene;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.QueryTimeout;
+import org.apache.lucene.search.*;
+import org.apache.lucene.search.knn.KnnCollectorManager;
+import org.apache.lucene.util.Bits;
+import org.opensearch.knn.profile.KNNProfileUtil;
+import org.opensearch.knn.profile.query.KNNQueryTimingType;
+import org.opensearch.search.profile.query.QueryProfiler;
+
+import java.io.IOException;
+
+/**
+ * Wrapper class used for profiling {@link KnnFloatVectorQuery}
+ */
+public class ProfileKnnFloatVectorQuery extends KnnFloatVectorQuery {
+
+    private QueryProfiler profiler;
+
+    public ProfileKnnFloatVectorQuery(String field, float[] target, int k, Query filter) {
+        super(field, target, k, filter);
+    }
+
+    @Override
+    public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+        profiler = KNNProfileUtil.getProfiler(indexSearcher);
+        return super.rewrite(indexSearcher);
+    }
+
+    @Override
+    protected TopDocs approximateSearch(
+        LeafReaderContext context,
+        Bits acceptDocs,
+        int visitedLimit,
+        KnnCollectorManager knnCollectorManager
+    ) throws IOException {
+        return (TopDocs) KNNProfileUtil.profile(profiler, this, context, KNNQueryTimingType.ANN_SEARCH, () -> {
+            try {
+                return super.approximateSearch(context, acceptDocs, visitedLimit, knnCollectorManager);
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        });
+    }
+
+    @Override
+    protected TopDocs exactSearch(LeafReaderContext context, DocIdSetIterator acceptIterator, QueryTimeout queryTimeout)
+        throws IOException {
+        return (TopDocs) KNNProfileUtil.profile(profiler, this, context, KNNQueryTimingType.EXACT_SEARCH, () -> {
+            try {
+                return super.exactSearch(context, acceptIterator, queryTimeout);
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        });
+    }
+
+}