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.

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
Show file tree
Hide file tree
Changes from 3 commits
Commits
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
13 changes: 13 additions & 0 deletions src/main/java/org/opensearch/knn/index/query/KNNQuery.java
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,17 @@ 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 = null;
if (log.isDebugEnabled()) {
stopWatch = new StopWatch().start();
}

final Weight filterWeight = getFilterWeight(searcher);
if (log.isDebugEnabled() && stopWatch != null) {
stopWatch.stop();
log.debug("Creating filter weight for field [{}] took [{}] nanos", field, stopWatch.totalTime().nanos());
Copy link
Collaborator

Choose a reason for hiding this comment

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

do we need any segment info like total docs which was present? name of the shard etc for better debugging. Because we have a lot of segments then this log statement won't be that helpful

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Added shard and segment info. cat APIs should help trace down total docs and deleted docs

}

if (filterWeight != null) {
return new KNNWeight(this, boost, filterWeight);
}
Expand Down
28 changes: 27 additions & 1 deletion src/main/java/org/opensearch/knn/index/query/KNNWeight.java
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@
import org.apache.lucene.util.BitSetIterator;
import org.apache.lucene.util.Bits;
import org.apache.lucene.util.FixedBitSet;
import org.opensearch.common.Nullable;
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 +131,10 @@ 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 = startStopWatch();
final BitSet filterBitSet = getFilteredDocsBitSet(context);
stopStopWatchAndLog(stopWatch, "FilterBitSet creation");

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 = startStopWatch();
final Map<Integer, Float> docIdsToScoreMap = doANNSearch(context, annFilter, cardinality, k);
stopStopWatchAndLog(annStopWatch, "ANN search");

// 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 All @@ -166,6 +174,14 @@ public PerLeafResult searchLeaf(LeafReaderContext context, int k) throws IOExcep
return new PerLeafResult(filterWeight == null ? null : filterBitSet, docIdsToScoreMap);
}

private void stopStopWatchAndLog(@Nullable final StopWatch stopWatch, final String prefixMessage) {
if (log.isDebugEnabled() && stopWatch != null) {
shatejas marked this conversation as resolved.
Show resolved Hide resolved
stopWatch.stop();
final String logMessage = prefixMessage + ", field: [{}], time in nanos:[{}] ";
log.debug(logMessage, knnQuery.getField(), stopWatch.totalTime().nanos());
}
}

private BitSet getFilteredDocsBitSet(final LeafReaderContext ctx) throws IOException {
if (this.filterWeight == null) {
return new FixedBitSet(0);
Expand Down Expand Up @@ -402,7 +418,10 @@ public Map<Integer, Float> exactSearch(
final LeafReaderContext leafReaderContext,
final ExactSearcher.ExactSearcherContext exactSearcherContext
) throws IOException {
return exactSearcher.searchLeaf(leafReaderContext, exactSearcherContext);
StopWatch stopWatch = startStopWatch();
Map<Integer, Float> exactSearchResults = exactSearcher.searchLeaf(leafReaderContext, exactSearcherContext);
stopStopWatchAndLog(stopWatch, "Exact search");
return exactSearchResults;
}

@Override
Expand Down Expand Up @@ -523,4 +542,11 @@ private boolean isMissingNativeEngineFiles(LeafReaderContext context) {
);
return engineFiles.isEmpty();
}

private StopWatch startStopWatch() {
shatejas marked this conversation as resolved.
Show resolved Hide resolved
if (log.isDebugEnabled()) {
return new StopWatch().start();
}
return null;
}
}
Loading