Skip to content

Commit 797517e

Browse files
benwtrentbenwtrent
committed
Have top level knn searches tracked in query stats
1 parent 0dbf9f7 commit 797517e

File tree

4 files changed

+257
-14
lines changed

4 files changed

+257
-14
lines changed

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -177,7 +177,7 @@ private static Timer maybeStartTimer(DfsProfiler profiler, DfsTimingType dtt) {
177177
return null;
178178
};
179179

180-
private static void executeKnnVectorQuery(SearchContext context) throws IOException {
180+
static void executeKnnVectorQuery(SearchContext context) throws IOException {
181181
SearchSourceBuilder source = context.request().source();
182182
if (source == null || source.knnSearch().isEmpty()) {
183183
return;
@@ -195,12 +195,18 @@ private static void executeKnnVectorQuery(SearchContext context) throws IOExcept
195195
}
196196
}
197197
List<DfsKnnResults> knnResults = new ArrayList<>(knnVectorQueryBuilders.size());
198+
final long afterQueryTime;
199+
final long beforeQueryTime = System.nanoTime();
200+
var opsListener = context.indexShard().getSearchOperationListener();
201+
opsListener.onPreQueryPhase(context);
198202
for (int i = 0; i < knnSearch.size(); i++) {
199203
String knnField = knnVectorQueryBuilders.get(i).getFieldName();
200204
String knnNestedPath = searchExecutionContext.nestedLookup().getNestedParent(knnField);
201205
Query knnQuery = searchExecutionContext.toQuery(knnVectorQueryBuilders.get(i)).query();
202206
knnResults.add(singleKnnSearch(knnQuery, knnSearch.get(i).k(), context.getProfilers(), context.searcher(), knnNestedPath));
203207
}
208+
afterQueryTime = System.nanoTime();
209+
opsListener.onQueryPhase(context, afterQueryTime - beforeQueryTime);
204210
context.dfsResult().knnResults(knnResults);
205211
}
206212

‎server/src/test/java/org/elasticsearch/index/shard/IndexShardTests.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5608,6 +5608,7 @@ public void testShardExposesWriteLoadStats() throws Exception {
56085608
RetentionLeaseSyncer.EMPTY,
56095609
EMPTY_EVENT_LISTENER,
56105610
fakeClock,
5611+
Collections.emptyList(),
56115612
// Use a listener to advance the fake clock once per indexing operation:
56125613
new IndexingOperationListener() {
56135614
@Override
@@ -5753,6 +5754,7 @@ public void testShardExposesWriteLoadStats_variableRates() throws IOException {
57535754
RetentionLeaseSyncer.EMPTY,
57545755
EMPTY_EVENT_LISTENER,
57555756
fakeClock,
5757+
Collections.emptyList(),
57565758
// Use a listener to advance the fake clock once per indexing operation:
57575759
new IndexingOperationListener() {
57585760
@Override

‎server/src/test/java/org/elasticsearch/search/dfs/DfsPhaseTests.java‎

Lines changed: 122 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,22 +17,44 @@
1717
import org.apache.lucene.search.Query;
1818
import org.apache.lucene.store.Directory;
1919
import org.apache.lucene.tests.index.RandomIndexWriter;
20+
import org.apache.lucene.util.Accountable;
21+
import org.elasticsearch.action.search.SearchShardTask;
22+
import org.elasticsearch.cluster.metadata.IndexMetadata;
23+
import org.elasticsearch.common.settings.Settings;
24+
import org.elasticsearch.index.IndexSettings;
25+
import org.elasticsearch.index.IndexVersion;
26+
import org.elasticsearch.index.cache.bitset.BitsetFilterCache;
27+
import org.elasticsearch.index.mapper.MapperMetrics;
28+
import org.elasticsearch.index.mapper.MappingLookup;
29+
import org.elasticsearch.index.query.ParsedQuery;
30+
import org.elasticsearch.index.query.SearchExecutionContext;
31+
import org.elasticsearch.index.shard.IndexShard;
32+
import org.elasticsearch.index.shard.IndexShardTestCase;
33+
import org.elasticsearch.index.shard.SearchOperationListener;
34+
import org.elasticsearch.index.shard.ShardId;
35+
import org.elasticsearch.search.builder.SearchSourceBuilder;
2036
import org.elasticsearch.search.internal.ContextIndexSearcher;
37+
import org.elasticsearch.search.internal.SearchContext;
2138
import org.elasticsearch.search.profile.Profilers;
2239
import org.elasticsearch.search.profile.SearchProfileDfsPhaseResult;
2340
import org.elasticsearch.search.profile.query.CollectorResult;
2441
import org.elasticsearch.search.profile.query.QueryProfileShardResult;
25-
import org.elasticsearch.test.ESTestCase;
42+
import org.elasticsearch.search.vectors.KnnSearchBuilder;
43+
import org.elasticsearch.test.TestSearchContext;
2644
import org.elasticsearch.threadpool.TestThreadPool;
2745
import org.elasticsearch.threadpool.ThreadPool;
2846
import org.junit.After;
2947
import org.junit.Before;
3048

3149
import java.io.IOException;
50+
import java.util.Collections;
3251
import java.util.List;
3352
import java.util.concurrent.ThreadPoolExecutor;
53+
import java.util.concurrent.atomic.AtomicLong;
3454

35-
public class DfsPhaseTests extends ESTestCase {
55+
import static org.elasticsearch.search.dfs.DfsPhase.executeKnnVectorQuery;
56+
57+
public class DfsPhaseTests extends IndexShardTestCase {
3658

3759
ThreadPoolExecutor threadPoolExecutor;
3860
private TestThreadPool threadPool;
@@ -49,6 +71,104 @@ public void cleanup() {
4971
terminate(threadPool);
5072
}
5173

74+
public void testKnnSearch() throws IOException {
75+
AtomicLong queryCount = new AtomicLong();
76+
AtomicLong queryTime = new AtomicLong();
77+
78+
IndexShard indexShard = newShard(true, List.of(new SearchOperationListener() {
79+
@Override
80+
public void onQueryPhase(SearchContext searchContext, long tookInNanos) {
81+
queryCount.incrementAndGet();
82+
queryTime.addAndGet(tookInNanos);
83+
}
84+
}));
85+
try (Directory dir = newDirectory(); RandomIndexWriter w = new RandomIndexWriter(random(), dir, newIndexWriterConfig())) {
86+
int numDocs = randomIntBetween(900, 1000);
87+
for (int i = 0; i < numDocs; i++) {
88+
Document d = new Document();
89+
d.add(new KnnFloatVectorField("float_vector", new float[] { i, 0, 0 }));
90+
w.addDocument(d);
91+
}
92+
w.flush();
93+
94+
IndexReader reader = w.getReader();
95+
ContextIndexSearcher searcher = new ContextIndexSearcher(
96+
reader,
97+
IndexSearcher.getDefaultSimilarity(),
98+
IndexSearcher.getDefaultQueryCache(),
99+
IndexSearcher.getDefaultQueryCachingPolicy(),
100+
randomBoolean(),
101+
threadPoolExecutor,
102+
threadPoolExecutor.getMaximumPoolSize(),
103+
1
104+
);
105+
IndexSettings indexSettings = new IndexSettings(
106+
IndexMetadata.builder("index")
107+
.settings(Settings.builder().put(IndexMetadata.SETTING_VERSION_CREATED, IndexVersion.current()))
108+
.numberOfShards(1)
109+
.numberOfReplicas(0)
110+
.creationDate(System.currentTimeMillis())
111+
.build(),
112+
Settings.EMPTY
113+
);
114+
BitsetFilterCache bitsetFilterCache = new BitsetFilterCache(indexSettings, new BitsetFilterCache.Listener() {
115+
@Override
116+
public void onCache(ShardId shardId, Accountable accountable) {
117+
118+
}
119+
120+
@Override
121+
public void onRemoval(ShardId shardId, Accountable accountable) {
122+
123+
}
124+
});
125+
SearchExecutionContext searchExecutionContext = new SearchExecutionContext(
126+
0,
127+
0,
128+
indexSettings,
129+
bitsetFilterCache,
130+
null,
131+
null,
132+
MappingLookup.EMPTY,
133+
null,
134+
null,
135+
null,
136+
null,
137+
null,
138+
null,
139+
null,
140+
null,
141+
null,
142+
null,
143+
null,
144+
Collections.emptyMap(),
145+
null,
146+
MapperMetrics.NOOP
147+
);
148+
149+
Query query = new KnnFloatVectorQuery("float_vector", new float[] { 0, 0, 0 }, numDocs, null);
150+
TestSearchContext context = new TestSearchContext(searchExecutionContext, indexShard, searcher) {
151+
@Override
152+
public DfsSearchResult dfsResult() {
153+
return new DfsSearchResult(null, null, null);
154+
}
155+
};
156+
context.request()
157+
.source(
158+
new SearchSourceBuilder().knnSearch(
159+
List.of(new KnnSearchBuilder("float_vector", new float[] { 0, 0, 0 }, numDocs, numDocs, null, null))
160+
)
161+
);
162+
context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap()));
163+
context.parsedQuery(new ParsedQuery(query));
164+
executeKnnVectorQuery(context);
165+
assertTrue(queryCount.get() > 0);
166+
assertTrue(queryTime.get() > 0);
167+
reader.close();
168+
closeShards(indexShard);
169+
}
170+
}
171+
52172
public void testSingleKnnSearch() throws IOException {
53173
try (Directory dir = newDirectory(); RandomIndexWriter w = new RandomIndexWriter(random(), dir, newIndexWriterConfig())) {
54174
int numDocs = randomIntBetween(900, 1000);

0 commit comments

Comments
 (0)