Skip to content

Commit c1a4f8a

Browse files
benwtrentbenwtrent
authored
Fix multi-indexing runs for checkVec (#130408)
We now can do multiple query time parameters, but this sort of broke running multiple indexing loads. This corrects multiple indexing & search tests for the KnnIndexTester: Example output. ``` index_name index_type num_docs index_time(ms) force_merge_time(ms) num_segments ----------------------------------- ---------- -------- -------------- -------------------- ------------ corpus-quora-E5-small.fvec.flat ivf 100000 4386 4522 0 cohere-wikipedia-docs-768d.vec ivf 100000 7784 8331 0 corpus-dbpedia-entity-arctic-0.fvec ivf 100000 7513 8969 0 index_name index_type n_probe latency(ms) net_cpu_time(ms) avg_cpu_count QPS recall visited ----------------------------------- ---------- ------- ----------- ---------------- ------------- ------- ------ -------- corpus-quora-E5-small.fvec.flat ivf 10 0.87 0.00 0.00 1149.43 0.88 5511.60 corpus-quora-E5-small.fvec.flat ivf 20 0.73 0.00 0.00 1369.86 0.94 10101.70 corpus-quora-E5-small.fvec.flat ivf 30 0.76 0.00 0.00 1315.79 0.95 14382.15 corpus-quora-E5-small.fvec.flat ivf 40 0.84 0.00 0.00 1190.48 0.96 18554.68 corpus-quora-E5-small.fvec.flat ivf 50 0.99 0.00 0.00 1010.10 0.97 22609.69 cohere-wikipedia-docs-768d.vec ivf 10 0.39 0.00 0.00 2564.10 0.61 2567.00 cohere-wikipedia-docs-768d.vec ivf 20 0.53 0.00 0.00 1886.79 0.78 5017.25 cohere-wikipedia-docs-768d.vec ivf 30 0.66 0.00 0.00 1515.15 0.84 7486.77 cohere-wikipedia-docs-768d.vec ivf 40 0.80 0.00 0.00 1250.00 0.88 10177.80 cohere-wikipedia-docs-768d.vec ivf 50 0.95 0.00 0.00 1052.63 0.90 12744.67 corpus-dbpedia-entity-arctic-0.fvec ivf 10 0.48 0.00 0.00 2083.33 0.62 3302.07 corpus-dbpedia-entity-arctic-0.fvec ivf 20 0.66 0.00 0.00 1515.15 0.74 6333.91 corpus-dbpedia-entity-arctic-0.fvec ivf 30 0.78 0.00 0.00 1282.05 0.81 9545.16 corpus-dbpedia-entity-arctic-0.fvec ivf 40 0.93 0.00 0.00 1075.27 0.85 12786.62 corpus-dbpedia-entity-arctic-0.fvec ivf 50 1.08 0.00 0.00 925.93 0.87 15749.40 ```
1 parent 2144bae commit c1a4f8a

File tree

1 file changed

+42
-35
lines changed

1 file changed

+42
-35
lines changed

‎qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexTester.java‎

Lines changed: 42 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -177,9 +177,11 @@ public static void main(String[] args) throws Exception {
177177
int[] nProbes = cmdLineArgs.indexType().equals(IndexType.IVF) && cmdLineArgs.numQueries() > 0
178178
? cmdLineArgs.nProbes()
179179
: new int[] { 0 };
180+
String indexType = cmdLineArgs.indexType().name().toLowerCase(Locale.ROOT);
181+
Results indexResults = new Results(cmdLineArgs.docVectors().getFileName().toString(), indexType, cmdLineArgs.numDocs());
180182
Results[] results = new Results[nProbes.length];
181183
for (int i = 0; i < nProbes.length; i++) {
182-
results[i] = new Results(cmdLineArgs.indexType().name().toLowerCase(Locale.ROOT), cmdLineArgs.numDocs());
184+
results[i] = new Results(cmdLineArgs.docVectors().getFileName().toString(), indexType, cmdLineArgs.numDocs());
183185
}
184186
logger.info("Running KNN index tester with arguments: " + cmdLineArgs);
185187
Codec codec = createCodec(cmdLineArgs);
@@ -199,12 +201,12 @@ public static void main(String[] args) throws Exception {
199201
throw new IllegalArgumentException("Index path does not exist: " + indexPath);
200202
}
201203
if (cmdLineArgs.reindex()) {
202-
knnIndexer.createIndex(results[0]);
204+
knnIndexer.createIndex(indexResults);
203205
}
204206
if (cmdLineArgs.forceMerge()) {
205-
knnIndexer.forceMerge(results[0]);
207+
knnIndexer.forceMerge(indexResults);
206208
} else {
207-
knnIndexer.numSegments(results[0]);
209+
knnIndexer.numSegments(indexResults);
208210
}
209211
}
210212
if (cmdLineArgs.queryVectors() != null && cmdLineArgs.numQueries() > 0) {
@@ -214,24 +216,27 @@ public static void main(String[] args) throws Exception {
214216
knnSearcher.runSearch(results[i], cmdLineArgs.earlyTermination());
215217
}
216218
}
217-
formattedResults.results.addAll(List.of(results));
219+
formattedResults.queryResults.addAll(List.of(results));
220+
formattedResults.indexResults.add(indexResults);
218221
}
219222
logger.info("Results: \n" + formattedResults);
220223
}
221224

222225
static class FormattedResults {
223-
List<Results> results = new ArrayList<>();
226+
List<Results> indexResults = new ArrayList<>();
227+
List<Results> queryResults = new ArrayList<>();
224228

225229
@Override
226230
public String toString() {
227-
if (results.isEmpty()) {
231+
if (indexResults.isEmpty() && queryResults.isEmpty()) {
228232
return "No results available.";
229233
}
230234

231-
String[] indexingHeaders = { "index_type", "num_docs", "index_time(ms)", "force_merge_time(ms)", "num_segments" };
235+
String[] indexingHeaders = { "index_name", "index_type", "num_docs", "index_time(ms)", "force_merge_time(ms)", "num_segments" };
232236

233237
// Define column headers
234238
String[] searchHeaders = {
239+
"index_name",
235240
"index_type",
236241
"n_probe",
237242
"latency(ms)",
@@ -245,33 +250,34 @@ public String toString() {
245250

246251
StringBuilder sb = new StringBuilder();
247252

248-
Results indexResult = results.get(0); // Assuming all results have the same index type and numDocs
249-
String[] indexData = {
250-
indexResult.indexType,
251-
Integer.toString(indexResult.numDocs),
252-
Long.toString(indexResult.indexTimeMS),
253-
Long.toString(indexResult.forceMergeTimeMS),
254-
Integer.toString(indexResult.numSegments) };
255-
256-
printBlock(sb, indexingHeaders, new String[][] { indexData });
257-
258-
String[][] searchData = new String[results.size()][];
259-
// Format and append each row of data
260-
for (int i = 0; i < results.size(); i++) {
261-
Results result = results.get(i);
262-
searchData[i] = new String[] {
263-
result.indexType,
264-
Integer.toString(result.nProbe),
265-
String.format(Locale.ROOT, "%.2f", result.avgLatency),
266-
String.format(Locale.ROOT, "%.2f", result.netCpuTimeMS),
267-
String.format(Locale.ROOT, "%.2f", result.avgCpuCount),
268-
String.format(Locale.ROOT, "%.2f", result.qps),
269-
String.format(Locale.ROOT, "%.2f", result.avgRecall),
270-
String.format(Locale.ROOT, "%.2f", result.averageVisited) };
271-
253+
String[][] indexResultsArray = new String[indexResults.size()][];
254+
for (int i = 0; i < indexResults.size(); i++) {
255+
Results indexResult = indexResults.get(i);
256+
indexResultsArray[i] = new String[] {
257+
indexResult.indexName,
258+
indexResult.indexType,
259+
Integer.toString(indexResult.numDocs),
260+
Long.toString(indexResult.indexTimeMS),
261+
Long.toString(indexResult.forceMergeTimeMS),
262+
Integer.toString(indexResult.numSegments) };
263+
}
264+
printBlock(sb, indexingHeaders, indexResultsArray);
265+
String[][] queryResultsArray = new String[queryResults.size()][];
266+
for (int i = 0; i < queryResults.size(); i++) {
267+
Results queryResult = queryResults.get(i);
268+
queryResultsArray[i] = new String[] {
269+
queryResult.indexName,
270+
queryResult.indexType,
271+
Integer.toString(queryResult.nProbe),
272+
String.format(Locale.ROOT, "%.2f", queryResult.avgLatency),
273+
String.format(Locale.ROOT, "%.2f", queryResult.netCpuTimeMS),
274+
String.format(Locale.ROOT, "%.2f", queryResult.avgCpuCount),
275+
String.format(Locale.ROOT, "%.2f", queryResult.qps),
276+
String.format(Locale.ROOT, "%.2f", queryResult.avgRecall),
277+
String.format(Locale.ROOT, "%.2f", queryResult.averageVisited) };
272278
}
273279

274-
printBlock(sb, searchHeaders, searchData);
280+
printBlock(sb, searchHeaders, queryResultsArray);
275281

276282
return sb.toString();
277283
}
@@ -331,7 +337,7 @@ private int[] calculateColumnWidths(String[] headers, String[]... data) {
331337
}
332338

333339
static class Results {
334-
final String indexType;
340+
final String indexType, indexName;
335341
final int numDocs;
336342
long indexTimeMS;
337343
long forceMergeTimeMS;
@@ -344,7 +350,8 @@ static class Results {
344350
double netCpuTimeMS;
345351
double avgCpuCount;
346352

347-
Results(String indexType, int numDocs) {
353+
Results(String indexName, String indexType, int numDocs) {
354+
this.indexName = indexName;
348355
this.indexType = indexType;
349356
this.numDocs = numDocs;
350357
}

0 commit comments

Comments
 (0)