Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adds debug logs for KNNQuery and KNNWeight #2466

Merged
merged 6 commits into from
Jan 30, 2025
Merged

Adds debug logs for KNNQuery and KNNWeight #2466

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
c526a13
Adds debug logs for KNNQuery and KNNWeight
shatejas Jan 29, 2025
2883e80
Adds check to see if log is enabled to start and stop StopWatch
shatejas Jan 29, 2025
38756c3
Addressing comments on the PR
shatejas Jan 29, 2025
d855696
Adds shard and segment info in the logs
shatejas Jan 30, 2025
7b0d3f6
Removes unnecessary segment name param from exact search
shatejas Jan 30, 2025
10821cd
Fixes the build
shatejas Jan 30, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions src/main/java/org/opensearch/knn/index/query/KNNQuery.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import lombok.EqualsAndHashCode;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.log4j.Log4j2;
import org.apache.lucene.search.BooleanClause;
import org.apache.lucene.search.BooleanQuery;
import org.apache.lucene.search.FieldExistsQuery;
Expand All @@ -19,6 +20,7 @@
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.search.Weight;
import org.apache.lucene.search.join.BitSetProducer;
import org.opensearch.common.StopWatch;
import org.opensearch.knn.index.KNNSettings;
import org.opensearch.knn.index.VectorDataType;
import org.opensearch.knn.index.query.rescore.RescoreContext;
Expand All @@ -32,6 +34,7 @@
* Custom KNN query. Query is used for KNNEngine's that create their own custom segment files. These files need to be
* loaded and queried in a custom manner throughout the query path.
*/
@Log4j2
@Getter
@Builder
@AllArgsConstructor
Expand Down Expand Up @@ -168,7 +171,10 @@ public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float bo
if (!KNNSettings.isKNNPluginEnabled()) {
throw new IllegalStateException("KNN plugin is disabled. To enable update knn.plugin.enabled to true");
}
StopWatch stopWatch = new StopWatch().start();
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we wrap this behind log.isDebugEnabled()?
It's not a big overhead but given that it's being used only for debug logging.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I thought about if but wanted to avoid if conditions. Let me add it

final Weight filterWeight = getFilterWeight(searcher);
stopWatch.stop();
log.debug("Creating filter weight for field [{}] took [{}] nanos", field, stopWatch.totalTime().nanos());
if (filterWeight != null) {
return new KNNWeight(this, boost, filterWeight);
}
Expand Down
14 changes: 13 additions & 1 deletion src/main/java/org/opensearch/knn/index/query/KNNWeight.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import org.apache.lucene.util.BitSetIterator;
import org.apache.lucene.util.Bits;
import org.apache.lucene.util.FixedBitSet;
import org.opensearch.common.StopWatch;
import org.opensearch.common.lucene.Lucene;
import org.opensearch.knn.common.FieldInfoExtractor;
import org.opensearch.knn.common.KNNConstants;
Expand Down Expand Up @@ -129,7 +130,11 @@ public Scorer scorer(LeafReaderContext context) throws IOException {
* @return A Map of docId to scores for top k results
*/
public PerLeafResult searchLeaf(LeafReaderContext context, int k) throws IOException {
StopWatch stopWatch = new StopWatch().start();
final BitSet filterBitSet = getFilteredDocsBitSet(context);
stopWatch.stop();
log.debug("Creating filter bitset for field [{}] took [{}] nanos", knnQuery.getField(), stopWatch.totalTime().nanos());

final int maxDoc = context.reader().maxDoc();
int cardinality = filterBitSet.cardinality();
// We don't need to go to JNI layer if no documents are found which satisfy the filters
Expand All @@ -153,7 +158,10 @@ public PerLeafResult searchLeaf(LeafReaderContext context, int k) throws IOExcep
* so that it will not do a bitset look up in bottom search layer.
*/
final BitSet annFilter = (filterWeight != null && cardinality == maxDoc) ? null : filterBitSet;
StopWatch annStopWatch = new StopWatch().start();
final Map<Integer, Float> docIdsToScoreMap = doANNSearch(context, annFilter, cardinality, k);
annStopWatch.stop();
log.debug("ANN search for field [{}] took [{}] nanos", knnQuery.getField(), annStopWatch.totalTime().nanos());

// See whether we have to perform exact search based on approx search results
// This is required if there are no native engine files or if approximate search returned
Expand Down Expand Up @@ -402,7 +410,11 @@ public Map<Integer, Float> exactSearch(
final LeafReaderContext leafReaderContext,
final ExactSearcher.ExactSearcherContext exactSearcherContext
) throws IOException {
return exactSearcher.searchLeaf(leafReaderContext, exactSearcherContext);
StopWatch stopWatch = new StopWatch().start();
Map<Integer, Float> exactSearchResults = exactSearcher.searchLeaf(leafReaderContext, exactSearcherContext);
stopWatch.stop();
log.debug("Exact search for field [{}] took [{}] nanos", knnQuery.getField(), stopWatch.totalTime().nanos());
return exactSearchResults;
}

@Override
Expand Down
Loading