Skip to content

Commit 3ed49cd

Browse files
benwtrentbenwtrent
authored
Fix timing bug with DFS profiling (#92421) (#92636)
Introduced in: #90536 Profiling for DFS has had its timing numbers looking weird, additionally, it would trigger some assertion failures because `timer.start()` was called without a `stop()` in between. The key issue was around query `weight` creation. `Weight` creation could be called recursively, thus calling `start` on the timer more than once before calling stop.
1 parent 09f141b commit 3ed49cd

File tree

4 files changed

+174
-49
lines changed

4 files changed

+174
-49
lines changed

‎docs/changelog/92421.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
pr: 92421
2+
summary: Fix timing bug with DFS profiling
3+
area: Search
4+
type: bug
5+
issues: []

‎server/src/internalClusterTest/java/org/elasticsearch/search/profile/dfs/DfsProfilerIT.java‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
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 and the Server Side Public License, v 1; you may not use this file except
5+
* in compliance with, at your election, the Elastic License 2.0 or the Server
6+
* Side Public License, v 1.
7+
*/
8+
9+
package org.elasticsearch.search.profile.dfs;
10+
11+
import org.apache.lucene.tests.util.English;
12+
import org.elasticsearch.action.index.IndexRequestBuilder;
13+
import org.elasticsearch.action.search.SearchResponse;
14+
import org.elasticsearch.action.search.SearchType;
15+
import org.elasticsearch.index.query.QueryBuilder;
16+
import org.elasticsearch.search.profile.ProfileResult;
17+
import org.elasticsearch.search.profile.SearchProfileDfsPhaseResult;
18+
import org.elasticsearch.search.profile.SearchProfileShardResult;
19+
import org.elasticsearch.search.profile.query.CollectorResult;
20+
import org.elasticsearch.search.profile.query.QueryProfileShardResult;
21+
import org.elasticsearch.search.vectors.KnnSearchBuilder;
22+
import org.elasticsearch.test.ESIntegTestCase;
23+
import org.elasticsearch.xcontent.XContentFactory;
24+
25+
import java.io.IOException;
26+
import java.util.List;
27+
import java.util.Map;
28+
29+
import static org.elasticsearch.search.profile.query.RandomQueryGenerator.randomQueryBuilder;
30+
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
31+
import static org.hamcrest.Matchers.emptyOrNullString;
32+
import static org.hamcrest.Matchers.equalTo;
33+
import static org.hamcrest.Matchers.greaterThan;
34+
import static org.hamcrest.Matchers.is;
35+
import static org.hamcrest.Matchers.not;
36+
import static org.hamcrest.Matchers.notNullValue;
37+
38+
public class DfsProfilerIT extends ESIntegTestCase {
39+
40+
private static final int KNN_DIM = 3;
41+
42+
public void testProfileDfs() throws Exception {
43+
String textField = "text_field";
44+
String numericField = "number";
45+
String vectorField = "vector";
46+
String indexName = "text-dfs-profile";
47+
createIndex(indexName, vectorField);
48+
ensureGreen();
49+
50+
int numDocs = randomIntBetween(10, 50);
51+
IndexRequestBuilder[] docs = new IndexRequestBuilder[numDocs];
52+
for (int i = 0; i < numDocs; i++) {
53+
docs[i] = client().prepareIndex(indexName)
54+
.setId(String.valueOf(i))
55+
.setSource(
56+
textField,
57+
English.intToEnglish(i),
58+
numericField,
59+
i,
60+
vectorField,
61+
new float[] { randomFloat(), randomFloat(), randomFloat() }
62+
);
63+
}
64+
indexRandom(true, docs);
65+
refresh();
66+
int iters = between(5, 10);
67+
for (int i = 0; i < iters; i++) {
68+
QueryBuilder q = randomQueryBuilder(List.of(textField), List.of(numericField), numDocs, 3);
69+
logger.info("Query: {}", q);
70+
71+
SearchResponse resp = client().prepareSearch()
72+
.setQuery(q)
73+
.setTrackTotalHits(true)
74+
.setProfile(true)
75+
.setSearchType(SearchType.DFS_QUERY_THEN_FETCH)
76+
.setKnnSearch(new KnnSearchBuilder(vectorField, new float[] { randomFloat(), randomFloat(), randomFloat() }, 5, 50))
77+
.get();
78+
79+
assertNotNull("Profile response element should not be null", resp.getProfileResults());
80+
assertThat("Profile response should not be an empty array", resp.getProfileResults().size(), not(0));
81+
for (Map.Entry<String, SearchProfileShardResult> shard : resp.getProfileResults().entrySet()) {
82+
for (QueryProfileShardResult searchProfiles : shard.getValue().getQueryProfileResults()) {
83+
for (ProfileResult result : searchProfiles.getQueryResults()) {
84+
assertNotNull(result.getQueryName());
85+
assertNotNull(result.getLuceneDescription());
86+
assertThat(result.getTime(), greaterThan(0L));
87+
}
88+
CollectorResult result = searchProfiles.getCollectorResult();
89+
assertThat(result.getName(), is(not(emptyOrNullString())));
90+
assertThat(result.getTime(), greaterThan(0L));
91+
}
92+
SearchProfileDfsPhaseResult searchProfileDfsPhaseResult = shard.getValue().getSearchProfileDfsPhaseResult();
93+
assertThat(searchProfileDfsPhaseResult, is(notNullValue()));
94+
for (ProfileResult result : searchProfileDfsPhaseResult.getQueryProfileShardResult().getQueryResults()) {
95+
assertNotNull(result.getQueryName());
96+
assertNotNull(result.getLuceneDescription());
97+
assertThat(result.getTime(), greaterThan(0L));
98+
}
99+
CollectorResult result = searchProfileDfsPhaseResult.getQueryProfileShardResult().getCollectorResult();
100+
assertThat(result.getName(), is(not(emptyOrNullString())));
101+
assertThat(result.getTime(), greaterThan(0L));
102+
ProfileResult statsResult = searchProfileDfsPhaseResult.getDfsShardResult();
103+
assertThat(statsResult.getQueryName(), equalTo("statistics"));
104+
assertThat(result.getTime(), greaterThan(0L));
105+
}
106+
}
107+
}
108+
109+
private void createIndex(String name, String vectorField) throws IOException {
110+
assertAcked(
111+
prepareCreate(name).setMapping(
112+
XContentFactory.jsonBuilder()
113+
.startObject()
114+
.startObject("properties")
115+
.startObject(vectorField)
116+
.field("type", "dense_vector")
117+
.field("dims", KNN_DIM)
118+
.field("index", true)
119+
.field("similarity", "cosine")
120+
.endObject()
121+
.endObject()
122+
.endObject()
123+
)
124+
);
125+
}
126+
127+
}

‎server/src/main/java/org/elasticsearch/search/dfs/DfsPhase.java‎

Lines changed: 34 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
import org.apache.lucene.search.ScoreMode;
1717
import org.apache.lucene.search.TermStatistics;
1818
import org.apache.lucene.search.TopScoreDocCollector;
19-
import org.apache.lucene.search.Weight;
2019
import org.elasticsearch.index.query.SearchExecutionContext;
2120
import org.elasticsearch.search.builder.SearchSourceBuilder;
2221
import org.elasticsearch.search.internal.SearchContext;
@@ -33,6 +32,7 @@
3332
import java.util.HashMap;
3433
import java.util.List;
3534
import java.util.Map;
35+
import java.util.function.Consumer;
3636

3737
/**
3838
* DFS phase of a search request, used to make scoring 100% accurate by collecting additional info from each shard before the query phase.
@@ -61,28 +61,32 @@ private void collectStatistics(SearchContext context) throws IOException {
6161

6262
Map<String, CollectionStatistics> fieldStatistics = new HashMap<>();
6363
Map<Term, TermStatistics> stats = new HashMap<>();
64+
final Consumer<DfsTimingType> maybeStart = dtt -> {
65+
if (profiler != null) {
66+
profiler.startTimer(dtt);
67+
}
68+
};
69+
final Consumer<DfsTimingType> maybeStop = dtt -> {
70+
if (profiler != null) {
71+
profiler.stopTimer(dtt);
72+
}
73+
};
6474

6575
IndexSearcher searcher = new IndexSearcher(context.searcher().getIndexReader()) {
6676
@Override
6777
public TermStatistics termStatistics(Term term, int docFreq, long totalTermFreq) throws IOException {
6878
if (context.isCancelled()) {
6979
throw new TaskCancelledException("cancelled");
7080
}
71-
72-
if (profiler != null) {
73-
profiler.startTimer(DfsTimingType.TERM_STATISTICS);
74-
}
75-
81+
maybeStart.accept(DfsTimingType.TERM_STATISTICS);
7682
try {
7783
TermStatistics ts = super.termStatistics(term, docFreq, totalTermFreq);
7884
if (ts != null) {
7985
stats.put(term, ts);
8086
}
8187
return ts;
8288
} finally {
83-
if (profiler != null) {
84-
profiler.stopTimer(DfsTimingType.TERM_STATISTICS);
85-
}
89+
maybeStop.accept(DfsTimingType.TERM_STATISTICS);
8690
}
8791
}
8892

@@ -91,51 +95,15 @@ public CollectionStatistics collectionStatistics(String field) throws IOExceptio
9195
if (context.isCancelled()) {
9296
throw new TaskCancelledException("cancelled");
9397
}
94-
95-
if (profiler != null) {
96-
profiler.startTimer(DfsTimingType.COLLECTION_STATISTICS);
97-
}
98-
98+
maybeStart.accept(DfsTimingType.COLLECTION_STATISTICS);
9999
try {
100100
CollectionStatistics cs = super.collectionStatistics(field);
101101
if (cs != null) {
102102
fieldStatistics.put(field, cs);
103103
}
104104
return cs;
105105
} finally {
106-
if (profiler != null) {
107-
profiler.stopTimer(DfsTimingType.COLLECTION_STATISTICS);
108-
}
109-
}
110-
}
111-
112-
@Override
113-
public Weight createWeight(Query query, ScoreMode scoreMode, float boost) throws IOException {
114-
if (profiler != null) {
115-
profiler.startTimer(DfsTimingType.CREATE_WEIGHT);
116-
}
117-
118-
try {
119-
return super.createWeight(query, scoreMode, boost);
120-
} finally {
121-
if (profiler != null) {
122-
profiler.stopTimer(DfsTimingType.CREATE_WEIGHT);
123-
}
124-
}
125-
}
126-
127-
@Override
128-
public Query rewrite(Query original) throws IOException {
129-
if (profiler != null) {
130-
profiler.startTimer(DfsTimingType.REWRITE);
131-
}
132-
133-
try {
134-
return super.rewrite(original);
135-
} finally {
136-
if (profiler != null) {
137-
profiler.stopTimer(DfsTimingType.REWRITE);
138-
}
106+
maybeStop.accept(DfsTimingType.COLLECTION_STATISTICS);
139107
}
140108
}
141109
};
@@ -145,10 +113,27 @@ public Query rewrite(Query original) throws IOException {
145113
}
146114

147115
try {
148-
searcher.createWeight(context.rewrittenQuery(), ScoreMode.COMPLETE, 1);
116+
try {
117+
maybeStart.accept(DfsTimingType.CREATE_WEIGHT);
118+
searcher.createWeight(context.rewrittenQuery(), ScoreMode.COMPLETE, 1);
119+
} finally {
120+
maybeStop.accept(DfsTimingType.CREATE_WEIGHT);
121+
}
149122
for (RescoreContext rescoreContext : context.rescore()) {
150123
for (Query query : rescoreContext.getQueries()) {
151-
searcher.createWeight(searcher.rewrite(query), ScoreMode.COMPLETE, 1);
124+
final Query rewritten;
125+
try {
126+
maybeStart.accept(DfsTimingType.REWRITE);
127+
rewritten = searcher.rewrite(query);
128+
} finally {
129+
maybeStop.accept(DfsTimingType.REWRITE);
130+
}
131+
try {
132+
maybeStart.accept(DfsTimingType.CREATE_WEIGHT);
133+
searcher.createWeight(rewritten, ScoreMode.COMPLETE, 1);
134+
} finally {
135+
maybeStop.accept(DfsTimingType.CREATE_WEIGHT);
136+
}
152137
}
153138
}
154139
} finally {

‎server/src/main/java/org/elasticsearch/search/profile/SearchProfileDfsPhaseResult.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,4 +103,12 @@ public String toString() {
103103
+ queryProfileShardResult
104104
+ '}';
105105
}
106+
107+
public ProfileResult getDfsShardResult() {
108+
return dfsShardResult;
109+
}
110+
111+
public QueryProfileShardResult getQueryProfileShardResult() {
112+
return queryProfileShardResult;
113+
}
106114
}

0 commit comments

Comments
 (0)