Skip to content

Commit 037299c

Browse files
pmpailispmpailis
committed
Add telemetry for retrievers (elastic#114109)
1 parent 6665b31 commit 037299c

File tree

19 files changed

+677
-47
lines changed

19 files changed

+677
-47
lines changed

‎docs/changelog/114109.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
pr: 114109
2+
summary: Update cluster stats for retrievers
3+
area: Search
4+
type: enhancement
5+
issues: []

‎docs/reference/cluster/stats.asciidoc‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -762,6 +762,10 @@ Queries are counted once per search request, meaning that if the same query type
762762
(object) Search sections used in selected nodes.
763763
For each section, name and number of times it's been used is reported.
764764

765+
`retrievers`::
766+
(object) Retriever types that were used in selected nodes.
767+
For each retriever, name and number of times it's been used is reported.
768+
765769
=====
766770
767771
`dense_vector`::

‎server/src/internalClusterTest/java/org/elasticsearch/search/retriever/RetrieverTelemetryIT.java‎

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.search.retriever;
11+
12+
import org.elasticsearch.action.admin.cluster.node.capabilities.NodesCapabilitiesRequest;
13+
import org.elasticsearch.action.admin.cluster.node.capabilities.NodesCapabilitiesResponse;
14+
import org.elasticsearch.action.admin.cluster.stats.SearchUsageStats;
15+
import org.elasticsearch.client.Request;
16+
import org.elasticsearch.common.Strings;
17+
import org.elasticsearch.index.query.QueryBuilders;
18+
import org.elasticsearch.rest.RestRequest;
19+
import org.elasticsearch.search.builder.SearchSourceBuilder;
20+
import org.elasticsearch.search.vectors.KnnSearchBuilder;
21+
import org.elasticsearch.search.vectors.KnnVectorQueryBuilder;
22+
import org.elasticsearch.test.ESIntegTestCase;
23+
import org.elasticsearch.xcontent.XContentBuilder;
24+
import org.elasticsearch.xcontent.XContentFactory;
25+
import org.junit.Before;
26+
27+
import java.io.IOException;
28+
import java.util.List;
29+
30+
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
31+
import static org.hamcrest.Matchers.equalTo;
32+
33+
@ESIntegTestCase.ClusterScope(scope = ESIntegTestCase.Scope.TEST, numDataNodes = 0)
34+
public class RetrieverTelemetryIT extends ESIntegTestCase {
35+
36+
private static final String INDEX_NAME = "test_index";
37+
38+
@Override
39+
protected boolean addMockHttpTransport() {
40+
return false; // enable http
41+
}
42+
43+
@Before
44+
public void setup() throws IOException {
45+
XContentBuilder builder = XContentFactory.jsonBuilder()
46+
.startObject()
47+
.startObject("properties")
48+
.startObject("vector")
49+
.field("type", "dense_vector")
50+
.field("dims", 1)
51+
.field("index", true)
52+
.field("similarity", "l2_norm")
53+
.startObject("index_options")
54+
.field("type", "hnsw")
55+
.endObject()
56+
.endObject()
57+
.startObject("text")
58+
.field("type", "text")
59+
.endObject()
60+
.startObject("integer")
61+
.field("type", "integer")
62+
.endObject()
63+
.startObject("topic")
64+
.field("type", "keyword")
65+
.endObject()
66+
.endObject()
67+
.endObject();
68+
69+
assertAcked(prepareCreate(INDEX_NAME).setMapping(builder));
70+
ensureGreen(INDEX_NAME);
71+
}
72+
73+
private void performSearch(SearchSourceBuilder source) throws IOException {
74+
Request request = new Request("GET", INDEX_NAME + "/_search");
75+
request.setJsonEntity(Strings.toString(source));
76+
getRestClient().performRequest(request);
77+
}
78+
79+
public void testTelemetryForRetrievers() throws IOException {
80+
81+
if (false == isRetrieverTelemetryEnabled()) {
82+
return;
83+
}
84+
85+
// search#1 - this will record 1 entry for "retriever" in `sections`, and 1 for "knn" under `retrievers`
86+
{
87+
performSearch(new SearchSourceBuilder().retriever(new KnnRetrieverBuilder("vector", new float[] { 1.0f }, null, 10, 15, null)));
88+
}
89+
90+
// search#2 - this will record 1 entry for "retriever" in `sections`, 1 for "standard" under `retrievers`, and 1 for "range" under
91+
// `queries`
92+
{
93+
performSearch(new SearchSourceBuilder().retriever(new StandardRetrieverBuilder(QueryBuilders.rangeQuery("integer").gte(2))));
94+
}
95+
96+
// search#3 - this will record 1 entry for "retriever" in `sections`, and 1 for "standard" under `retrievers`, and 1 for "knn" under
97+
// `queries`
98+
{
99+
performSearch(
100+
new SearchSourceBuilder().retriever(
101+
new StandardRetrieverBuilder(new KnnVectorQueryBuilder("vector", new float[] { 1.0f }, 10, 15, null))
102+
)
103+
);
104+
}
105+
106+
// search#4 - this will record 1 entry for "retriever" in `sections`, and 1 for "standard" under `retrievers`, and 1 for "term"
107+
// under `queries`
108+
{
109+
performSearch(new SearchSourceBuilder().retriever(new StandardRetrieverBuilder(QueryBuilders.termQuery("topic", "foo"))));
110+
}
111+
112+
// search#5 - t
113+
// his will record 1 entry for "knn" in `sections`
114+
{
115+
performSearch(new SearchSourceBuilder().knnSearch(List.of(new KnnSearchBuilder("vector", new float[] { 1.0f }, 10, 15, null))));
116+
}
117+
118+
// search#6 - this will record 1 entry for "query" in `sections`, and 1 for "match_all" under `queries`
119+
{
120+
performSearch(new SearchSourceBuilder().query(QueryBuilders.matchAllQuery()));
121+
}
122+
123+
// cluster stats
124+
{
125+
SearchUsageStats stats = clusterAdmin().prepareClusterStats().get().getIndicesStats().getSearchUsageStats();
126+
assertEquals(6, stats.getTotalSearchCount());
127+
128+
assertThat(stats.getSectionsUsage().size(), equalTo(3));
129+
assertThat(stats.getSectionsUsage().get("retriever"), equalTo(4L));
130+
assertThat(stats.getSectionsUsage().get("query"), equalTo(1L));
131+
assertThat(stats.getSectionsUsage().get("knn"), equalTo(1L));
132+
133+
assertThat(stats.getRetrieversUsage().size(), equalTo(2));
134+
assertThat(stats.getRetrieversUsage().get("standard"), equalTo(3L));
135+
assertThat(stats.getRetrieversUsage().get("knn"), equalTo(1L));
136+
137+
assertThat(stats.getQueryUsage().size(), equalTo(4));
138+
assertThat(stats.getQueryUsage().get("range"), equalTo(1L));
139+
assertThat(stats.getQueryUsage().get("term"), equalTo(1L));
140+
assertThat(stats.getQueryUsage().get("match_all"), equalTo(1L));
141+
assertThat(stats.getQueryUsage().get("knn"), equalTo(1L));
142+
}
143+
}
144+
145+
private boolean isRetrieverTelemetryEnabled() throws IOException {
146+
NodesCapabilitiesResponse res = clusterAdmin().nodesCapabilities(
147+
new NodesCapabilitiesRequest().method(RestRequest.Method.GET).path("_cluster/stats").capabilities("retrievers-usage-stats")
148+
).actionGet();
149+
return res != null && res.isSupported().orElse(false);
150+
}
151+
}

‎server/src/main/java/org/elasticsearch/TransportVersions.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -238,6 +238,7 @@ static TransportVersion def(int id) {
238238
public static final TransportVersion FAST_REFRESH_RCO = def(8_762_00_0);
239239
public static final TransportVersion TEXT_SIMILARITY_RERANKER_QUERY_REWRITE = def(8_763_00_0);
240240
public static final TransportVersion SIMULATE_INDEX_TEMPLATES_SUBSTITUTIONS = def(8_764_00_0);
241+
public static final TransportVersion RETRIEVERS_TELEMETRY_ADDED = def(8_765_00_0);
241242

242243
/*
243244
* STOP! READ THIS FIRST! No, really,

‎server/src/main/java/org/elasticsearch/action/admin/cluster/stats/SearchUsageStats.java‎

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import java.util.Map;
2323
import java.util.Objects;
2424

25+
import static org.elasticsearch.TransportVersions.RETRIEVERS_TELEMETRY_ADDED;
2526
import static org.elasticsearch.TransportVersions.V_8_12_0;
2627

2728
/**
@@ -34,6 +35,7 @@ public final class SearchUsageStats implements Writeable, ToXContentFragment {
3435
private final Map<String, Long> queries;
3536
private final Map<String, Long> rescorers;
3637
private final Map<String, Long> sections;
38+
private final Map<String, Long> retrievers;
3739

3840
/**
3941
* Creates a new empty stats instance, that will get additional stats added through {@link #add(SearchUsageStats)}
@@ -43,24 +45,33 @@ public SearchUsageStats() {
4345
this.queries = new HashMap<>();
4446
this.sections = new HashMap<>();
4547
this.rescorers = new HashMap<>();
48+
this.retrievers = new HashMap<>();
4649
}
4750

4851
/**
4952
* Creates a new stats instance with the provided info. The expectation is that when a new instance is created using
5053
* this constructor, the provided stats are final and won't be modified further.
5154
*/
52-
public SearchUsageStats(Map<String, Long> queries, Map<String, Long> rescorers, Map<String, Long> sections, long totalSearchCount) {
55+
public SearchUsageStats(
56+
Map<String, Long> queries,
57+
Map<String, Long> rescorers,
58+
Map<String, Long> sections,
59+
Map<String, Long> retrievers,
60+
long totalSearchCount
61+
) {
5362
this.totalSearchCount = totalSearchCount;
5463
this.queries = queries;
5564
this.sections = sections;
5665
this.rescorers = rescorers;
66+
this.retrievers = retrievers;
5767
}
5868

5969
public SearchUsageStats(StreamInput in) throws IOException {
6070
this.queries = in.readMap(StreamInput::readLong);
6171
this.sections = in.readMap(StreamInput::readLong);
6272
this.totalSearchCount = in.readVLong();
6373
this.rescorers = in.getTransportVersion().onOrAfter(V_8_12_0) ? in.readMap(StreamInput::readLong) : Map.of();
74+
this.retrievers = in.getTransportVersion().onOrAfter(RETRIEVERS_TELEMETRY_ADDED) ? in.readMap(StreamInput::readLong) : Map.of();
6475
}
6576

6677
@Override
@@ -72,6 +83,9 @@ public void writeTo(StreamOutput out) throws IOException {
7283
if (out.getTransportVersion().onOrAfter(V_8_12_0)) {
7384
out.writeMap(rescorers, StreamOutput::writeLong);
7485
}
86+
if (out.getTransportVersion().onOrAfter(RETRIEVERS_TELEMETRY_ADDED)) {
87+
out.writeMap(retrievers, StreamOutput::writeLong);
88+
}
7589
}
7690

7791
/**
@@ -81,6 +95,7 @@ public void add(SearchUsageStats stats) {
8195
stats.queries.forEach((query, count) -> queries.merge(query, count, Long::sum));
8296
stats.rescorers.forEach((rescorer, count) -> rescorers.merge(rescorer, count, Long::sum));
8397
stats.sections.forEach((query, count) -> sections.merge(query, count, Long::sum));
98+
stats.retrievers.forEach((query, count) -> retrievers.merge(query, count, Long::sum));
8499
this.totalSearchCount += stats.totalSearchCount;
85100
}
86101

@@ -95,6 +110,8 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
95110
builder.map(rescorers);
96111
builder.field("sections");
97112
builder.map(sections);
113+
builder.field("retrievers");
114+
builder.map(retrievers);
98115
}
99116
builder.endObject();
100117
return builder;
@@ -112,6 +129,10 @@ public Map<String, Long> getSectionsUsage() {
112129
return Collections.unmodifiableMap(sections);
113130
}
114131

132+
public Map<String, Long> getRetrieversUsage() {
133+
return Collections.unmodifiableMap(retrievers);
134+
}
135+
115136
public long getTotalSearchCount() {
116137
return totalSearchCount;
117138
}
@@ -128,12 +149,13 @@ public boolean equals(Object o) {
128149
return totalSearchCount == that.totalSearchCount
129150
&& queries.equals(that.queries)
130151
&& rescorers.equals(that.rescorers)
131-
&& sections.equals(that.sections);
152+
&& sections.equals(that.sections)
153+
&& retrievers.equals(that.retrievers);
132154
}
133155

134156
@Override
135157
public int hashCode() {
136-
return Objects.hash(totalSearchCount, queries, rescorers, sections);
158+
return Objects.hash(totalSearchCount, queries, rescorers, sections, retrievers);
137159
}
138160

139161
@Override

‎server/src/main/java/org/elasticsearch/rest/action/admin/cluster/RestClusterStatsAction.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,8 @@ public class RestClusterStatsAction extends BaseRestHandler {
3232
private static final Set<String> SUPPORTED_CAPABILITIES = Set.of(
3333
"human-readable-total-docs-size",
3434
"verbose-dense-vector-mapping-stats",
35-
"ccs-stats"
35+
"ccs-stats",
36+
"retrievers-usage-stats"
3637
);
3738
private static final Set<String> SUPPORTED_QUERY_PARAMETERS = Set.of("include_remotes", "nodeId", REST_TIMEOUT_PARAM);
3839

‎server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1409,6 +1409,7 @@ private SearchSourceBuilder parseXContent(
14091409
parser,
14101410
new RetrieverParserContext(searchUsage, clusterSupportsFeature)
14111411
);
1412+
searchUsage.trackSectionUsage(RETRIEVER.getPreferredName());
14121413
} else if (QUERY_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
14131414
if (subSearchSourceBuilders.isEmpty() == false) {
14141415
throw new IllegalArgumentException(

‎server/src/main/java/org/elasticsearch/search/retriever/RetrieverBuilder.java‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -62,11 +62,11 @@ protected static void declareBaseParserFields(
6262
String name,
6363
AbstractObjectParser<? extends RetrieverBuilder, RetrieverParserContext> parser
6464
) {
65-
parser.declareObjectArray((r, v) -> r.preFilterQueryBuilders = v, (p, c) -> {
66-
QueryBuilder preFilterQueryBuilder = AbstractQueryBuilder.parseTopLevelQuery(p, c::trackQueryUsage);
67-
c.trackSectionUsage(name + ":" + PRE_FILTER_FIELD.getPreferredName());
68-
return preFilterQueryBuilder;
69-
}, PRE_FILTER_FIELD);
65+
parser.declareObjectArray(
66+
(r, v) -> r.preFilterQueryBuilders = v,
67+
(p, c) -> AbstractQueryBuilder.parseTopLevelQuery(p, c::trackQueryUsage),
68+
PRE_FILTER_FIELD
69+
);
7070
parser.declareString(RetrieverBuilder::retrieverName, NAME_FIELD);
7171
parser.declareFloat(RetrieverBuilder::minScore, MIN_SCORE_FIELD);
7272
}
@@ -138,7 +138,7 @@ protected static RetrieverBuilder parseInnerRetrieverBuilder(XContentParser pars
138138
throw new ParsingException(new XContentLocation(nonfe.getLineNumber(), nonfe.getColumnNumber()), message, nonfe);
139139
}
140140

141-
context.trackSectionUsage(retrieverName);
141+
context.trackRetrieverUsage(retrieverName);
142142

143143
if (parser.currentToken() != XContentParser.Token.END_OBJECT) {
144144
throw new ParsingException(

‎server/src/main/java/org/elasticsearch/search/retriever/RetrieverParserContext.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,10 @@ public void trackRescorerUsage(String name) {
3737
searchUsage.trackRescorerUsage(name);
3838
}
3939

40+
public void trackRetrieverUsage(String name) {
41+
searchUsage.trackRetrieverUsage(name);
42+
}
43+
4044
public boolean clusterSupportsFeature(NodeFeature nodeFeature) {
4145
return clusterSupportsFeature != null && clusterSupportsFeature.test(nodeFeature);
4246
}

‎server/src/main/java/org/elasticsearch/search/retriever/StandardRetrieverBuilder.java‎

Lines changed: 19 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -53,36 +53,28 @@ public final class StandardRetrieverBuilder extends RetrieverBuilder implements
5353
static {
5454
PARSER.declareObject((r, v) -> r.queryBuilder = v, (p, c) -> {
5555
QueryBuilder queryBuilder = AbstractQueryBuilder.parseTopLevelQuery(p, c::trackQueryUsage);
56-
c.trackSectionUsage(NAME + ":" + QUERY_FIELD.getPreferredName());
5756
return queryBuilder;
5857
}, QUERY_FIELD);
5958

60-
PARSER.declareField((r, v) -> r.searchAfterBuilder = v, (p, c) -> {
61-
SearchAfterBuilder searchAfterBuilder = SearchAfterBuilder.fromXContent(p);
62-
c.trackSectionUsage(NAME + ":" + SEARCH_AFTER_FIELD.getPreferredName());
63-
return searchAfterBuilder;
64-
}, SEARCH_AFTER_FIELD, ObjectParser.ValueType.OBJECT_ARRAY);
65-
66-
PARSER.declareField((r, v) -> r.terminateAfter = v, (p, c) -> {
67-
int terminateAfter = p.intValue();
68-
c.trackSectionUsage(NAME + ":" + TERMINATE_AFTER_FIELD.getPreferredName());
69-
return terminateAfter;
70-
}, TERMINATE_AFTER_FIELD, ObjectParser.ValueType.INT);
71-
72-
PARSER.declareField((r, v) -> r.sortBuilders = v, (p, c) -> {
73-
List<SortBuilder<?>> sortBuilders = SortBuilder.fromXContent(p);
74-
c.trackSectionUsage(NAME + ":" + SORT_FIELD.getPreferredName());
75-
return sortBuilders;
76-
}, SORT_FIELD, ObjectParser.ValueType.OBJECT_ARRAY);
77-
78-
PARSER.declareField((r, v) -> r.collapseBuilder = v, (p, c) -> {
79-
CollapseBuilder collapseBuilder = CollapseBuilder.fromXContent(p);
80-
if (collapseBuilder.getField() != null) {
81-
c.trackSectionUsage(COLLAPSE_FIELD.getPreferredName());
82-
}
83-
return collapseBuilder;
84-
}, COLLAPSE_FIELD, ObjectParser.ValueType.OBJECT);
85-
59+
PARSER.declareField(
60+
(r, v) -> r.searchAfterBuilder = v,
61+
(p, c) -> SearchAfterBuilder.fromXContent(p),
62+
SEARCH_AFTER_FIELD,
63+
ObjectParser.ValueType.OBJECT_ARRAY
64+
);
65+
PARSER.declareField((r, v) -> r.terminateAfter = v, (p, c) -> p.intValue(), TERMINATE_AFTER_FIELD, ObjectParser.ValueType.INT);
66+
PARSER.declareField(
67+
(r, v) -> r.sortBuilders = v,
68+
(p, c) -> SortBuilder.fromXContent(p),
69+
SORT_FIELD,
70+
ObjectParser.ValueType.OBJECT_ARRAY
71+
);
72+
PARSER.declareField(
73+
(r, v) -> r.collapseBuilder = v,
74+
(p, c) -> CollapseBuilder.fromXContent(p),
75+
COLLAPSE_FIELD,
76+
ObjectParser.ValueType.OBJECT
77+
);
8678
RetrieverBuilder.declareBaseParserFields(NAME, PARSER);
8779
}
8880

0 commit comments

Comments
 (0)