Skip to content

Commit

Permalink
Add stats for radial search (#1684) (#1701)
Browse files Browse the repository at this point in the history
(cherry picked from commit 9a52b2b)

Signed-off-by: Junqiu Lei <[email protected]>
  • Loading branch information
junqiu-lei authored May 14, 2024
1 parent 46e4051 commit bf8d2a9
Show file tree
Hide file tree
Showing 8 changed files with 245 additions and 40 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
### Features
### Enhancements
* Add KnnCircuitBreakerException and modify exception message [#1688](https://github.com/opensearch-project/k-NN/pull/1688)
* Add stats for radial search [#1684](https://github.com/opensearch-project/k-NN/pull/1684)
### Bug Fixes
* Block commas in model description [#1692](https://github.com/opensearch-project/k-NN/pull/1692)
### Infrastructure
Expand Down
57 changes: 57 additions & 0 deletions src/main/java/org/opensearch/knn/index/VectorQueryType.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

package org.opensearch.knn.index;

import lombok.Getter;
import org.opensearch.knn.common.KNNConstants;
import org.opensearch.knn.plugin.stats.KNNCounter;

@Getter
public enum VectorQueryType {
K(KNNConstants.K) {
@Override
public KNNCounter getQueryStatCounter() {
return KNNCounter.KNN_QUERY_REQUESTS;
}

@Override
public KNNCounter getQueryWithFilterStatCounter() {
return KNNCounter.KNN_QUERY_WITH_FILTER_REQUESTS;
}
},
MIN_SCORE(KNNConstants.MIN_SCORE) {
@Override
public KNNCounter getQueryStatCounter() {
return KNNCounter.MIN_SCORE_QUERY_REQUESTS;
}

@Override
public KNNCounter getQueryWithFilterStatCounter() {
return KNNCounter.MIN_SCORE_QUERY_WITH_FILTER_REQUESTS;
}
},
MAX_DISTANCE(KNNConstants.MAX_DISTANCE) {
@Override
public KNNCounter getQueryStatCounter() {
return KNNCounter.MAX_DISTANCE_QUERY_REQUESTS;
}

@Override
public KNNCounter getQueryWithFilterStatCounter() {
return KNNCounter.MAX_DISTANCE_QUERY_WITH_FILTER_REQUESTS;
}
};

private final String queryTypeName;

VectorQueryType(String queryTypeName) {
this.queryTypeName = queryTypeName;
}

public abstract KNNCounter getQueryStatCounter();

public abstract KNNCounter getQueryWithFilterStatCounter();
}
37 changes: 23 additions & 14 deletions src/main/java/org/opensearch/knn/index/query/KNNQueryBuilder.java
Original file line number Diff line number Diff line change
Expand Up @@ -5,42 +5,43 @@

package org.opensearch.knn.index.query;

import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;
import lombok.extern.log4j.Log4j2;
import org.apache.lucene.search.MatchNoDocsQuery;
import org.opensearch.core.common.Strings;
import org.apache.commons.lang.StringUtils;
import org.apache.lucene.search.MatchNoDocsQuery;
import org.apache.lucene.search.Query;
import org.opensearch.core.ParseField;
import org.opensearch.core.common.ParsingException;
import org.opensearch.core.common.Strings;
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.core.common.io.stream.StreamOutput;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.core.xcontent.XContentParser;
import org.opensearch.index.mapper.MappedFieldType;
import org.opensearch.index.mapper.NumberFieldMapper;
import org.opensearch.index.query.AbstractQueryBuilder;
import org.opensearch.index.query.QueryBuilder;
import org.opensearch.index.query.QueryShardContext;
import org.opensearch.knn.common.KNNConstants;
import org.opensearch.knn.index.KNNMethodContext;
import org.opensearch.knn.index.SpaceType;
import org.opensearch.knn.index.VectorDataType;
import org.opensearch.knn.index.VectorQueryType;
import org.opensearch.knn.index.mapper.KNNVectorFieldMapper;
import org.opensearch.knn.index.util.KNNEngine;
import org.opensearch.knn.indices.ModelDao;
import org.opensearch.knn.indices.ModelMetadata;
import org.opensearch.knn.indices.ModelUtil;
import org.opensearch.knn.plugin.stats.KNNCounter;
import org.opensearch.index.query.AbstractQueryBuilder;
import org.opensearch.index.query.QueryShardContext;

import static org.opensearch.knn.index.IndexUtil.*;
import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;

import static org.opensearch.knn.common.KNNConstants.MAX_DISTANCE;
import static org.opensearch.knn.common.KNNConstants.MIN_SCORE;
import static org.opensearch.knn.common.KNNValidationUtil.validateByteVectorValue;
import static org.opensearch.knn.index.IndexUtil.isClusterOnOrAfterMinRequiredVersion;
import static org.opensearch.knn.index.IndexUtil.minimalRequiredVersionMap;
import static org.opensearch.knn.index.util.KNNEngine.ENGINES_SUPPORTING_RADIAL_SEARCH;

/**
Expand Down Expand Up @@ -246,7 +247,6 @@ public static KNNQueryBuilder fromXContent(XContentParser parser) throws IOExcep
String queryName = null;
String currentFieldName = null;
XContentParser.Token token;
KNNCounter.KNN_QUERY_REQUESTS.increment();
while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
if (token == XContentParser.Token.FIELD_NAME) {
currentFieldName = parser.currentName();
Expand Down Expand Up @@ -283,7 +283,6 @@ public static KNNQueryBuilder fromXContent(XContentParser parser) throws IOExcep
String tokenName = parser.currentName();
if (FILTER_FIELD.getPreferredName().equals(tokenName)) {
log.debug(String.format("Start parsing filter for field [%s]", fieldName));
KNNCounter.KNN_QUERY_WITH_FILTER_REQUESTS.increment();
// Query filters are supported starting from a certain k-NN version only, exact version is defined by
// MINIMAL_SUPPORTED_VERSION_FOR_LUCENE_HNSW_FILTER variable.
// Here we're checking if all cluster nodes has at least that version or higher. This check is required
Expand Down Expand Up @@ -322,7 +321,11 @@ public static KNNQueryBuilder fromXContent(XContentParser parser) throws IOExcep
}
}

validateSingleQueryType(k, maxDistance, minScore);
VectorQueryType vectorQueryType = validateSingleQueryType(k, maxDistance, minScore);
vectorQueryType.getQueryStatCounter().increment();
if (filter != null) {
vectorQueryType.getQueryWithFilterStatCounter().increment();
}

KNNQueryBuilder knnQueryBuilder = new KNNQueryBuilder(fieldName, ObjectsToFloats(vector)).filter(filter)
.boost(boost)
Expand Down Expand Up @@ -580,21 +583,27 @@ public String getWriteableName() {
return NAME;
}

private static void validateSingleQueryType(Integer k, Float distance, Float score) {
private static VectorQueryType validateSingleQueryType(Integer k, Float distance, Float score) {
int countSetFields = 0;
VectorQueryType vectorQueryType = null;

if (k != null && k != 0) {
countSetFields++;
vectorQueryType = VectorQueryType.K;
}
if (distance != null) {
countSetFields++;
vectorQueryType = VectorQueryType.MAX_DISTANCE;
}
if (score != null) {
countSetFields++;
vectorQueryType = VectorQueryType.MIN_SCORE;
}

if (countSetFields != 1) {
throw new IllegalArgumentException(String.format("[%s] requires exactly one of k, distance or score to be set", NAME));
}

return vectorQueryType;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,11 @@ public enum KNNCounter {
SCRIPT_QUERY_ERRORS("script_query_errors"),
TRAINING_REQUESTS("training_requests"),
TRAINING_ERRORS("training_errors"),
KNN_QUERY_WITH_FILTER_REQUESTS("knn_query_with_filter_requests");
KNN_QUERY_WITH_FILTER_REQUESTS("knn_query_with_filter_requests"),
MIN_SCORE_QUERY_REQUESTS("min_score_query_requests"),
MIN_SCORE_QUERY_WITH_FILTER_REQUESTS("min_score_query_with_filter_requests"),
MAX_DISTANCE_QUERY_REQUESTS("max_distance_query_requests"),
MAX_DISTANCE_QUERY_WITH_FILTER_REQUESTS("max_distance_query_with_filter_requests");

private String name;
private AtomicLong count;
Expand Down
20 changes: 20 additions & 0 deletions src/main/java/org/opensearch/knn/plugin/stats/KNNStats.java
Original file line number Diff line number Diff line change
Expand Up @@ -90,12 +90,32 @@ private Map<String, KNNStat<?>> buildStatsMap() {
}

private void addQueryStats(ImmutableMap.Builder<String, KNNStat<?>> builder) {
// KNN Query Stats
builder.put(StatNames.KNN_QUERY_REQUESTS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.KNN_QUERY_REQUESTS)))
.put(
StatNames.KNN_QUERY_WITH_FILTER_REQUESTS.getName(),
new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.KNN_QUERY_WITH_FILTER_REQUESTS))
);

// Min Score Query Stats
builder.put(
StatNames.MIN_SCORE_QUERY_REQUESTS.getName(),
new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.MIN_SCORE_QUERY_REQUESTS))
)
.put(
StatNames.MIN_SCORE_QUERY_WITH_FILTER_REQUESTS.getName(),
new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.MIN_SCORE_QUERY_WITH_FILTER_REQUESTS))
);

// Max Distance Query Stats
builder.put(
StatNames.MAX_DISTANCE_QUERY_REQUESTS.getName(),
new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.MAX_DISTANCE_QUERY_REQUESTS))
)
.put(
StatNames.MAX_DISTANCE_QUERY_WITH_FILTER_REQUESTS.getName(),
new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.MAX_DISTANCE_QUERY_WITH_FILTER_REQUESTS))
);
}

private void addNativeMemoryStats(ImmutableMap.Builder<String, KNNStat<?>> builder) {
Expand Down
6 changes: 5 additions & 1 deletion src/main/java/org/opensearch/knn/plugin/stats/StatNames.java
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,11 @@ public enum StatNames {
KNN_QUERY_WITH_FILTER_REQUESTS(KNNCounter.KNN_QUERY_WITH_FILTER_REQUESTS.getName()),
GRAPH_STATS("graph_stats"),
REFRESH("refresh"),
MERGE("merge");
MERGE("merge"),
MIN_SCORE_QUERY_REQUESTS(KNNCounter.MIN_SCORE_QUERY_REQUESTS.getName()),
MIN_SCORE_QUERY_WITH_FILTER_REQUESTS(KNNCounter.MIN_SCORE_QUERY_WITH_FILTER_REQUESTS.getName()),
MAX_DISTANCE_QUERY_REQUESTS(KNNCounter.MAX_DISTANCE_QUERY_REQUESTS.getName()),
MAX_DISTANCE_QUERY_WITH_FILTER_REQUESTS(KNNCounter.MAX_DISTANCE_QUERY_WITH_FILTER_REQUESTS.getName());

private String name;

Expand Down
30 changes: 30 additions & 0 deletions src/test/java/org/opensearch/knn/index/VectorQueryTypeTests.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
/*
* SPDX-License-Identifier: Apache-2.0
*
* The OpenSearch Contributors require contributions made to
* this file be licensed under the Apache-2.0 license or a
* compatible open source license.
*
* Modifications Copyright OpenSearch Contributors. See
* GitHub history for details.
*/

package org.opensearch.knn.index;

import org.opensearch.knn.KNNTestCase;
import org.opensearch.knn.plugin.stats.KNNCounter;

public class VectorQueryTypeTests extends KNNTestCase {

public void testGetQueryStatCounter() {
assertEquals(KNNCounter.KNN_QUERY_REQUESTS, VectorQueryType.K.getQueryStatCounter());
assertEquals(KNNCounter.MIN_SCORE_QUERY_REQUESTS, VectorQueryType.MIN_SCORE.getQueryStatCounter());
assertEquals(KNNCounter.MAX_DISTANCE_QUERY_REQUESTS, VectorQueryType.MAX_DISTANCE.getQueryStatCounter());
}

public void testGetQueryWithFilterStatCounter() {
assertEquals(KNNCounter.KNN_QUERY_WITH_FILTER_REQUESTS, VectorQueryType.K.getQueryWithFilterStatCounter());
assertEquals(KNNCounter.MIN_SCORE_QUERY_WITH_FILTER_REQUESTS, VectorQueryType.MIN_SCORE.getQueryWithFilterStatCounter());
assertEquals(KNNCounter.MAX_DISTANCE_QUERY_WITH_FILTER_REQUESTS, VectorQueryType.MAX_DISTANCE.getQueryWithFilterStatCounter());
}
}
Loading

0 comments on commit bf8d2a9

Please sign in to comment.