Skip to content

Commit 9c38943

Browse files
davidkyledavidkyle
committed
Rename to TestConfiguration
1 parent 7d2274b commit 9c38943

File tree

4 files changed

+79
-71
lines changed

4 files changed

+79
-71
lines changed

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

Lines changed: 45 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ enum MergePolicyType {
9797
LOG_DOC
9898
}
9999

100-
private static String formatIndexPath(CmdLineArgs args) {
100+
private static String formatIndexPath(TestConfiguration args) {
101101
List<String> suffix = new ArrayList<>();
102102
if (args.indexType() == IndexType.FLAT) {
103103
suffix.add("flat");
@@ -116,7 +116,7 @@ private static String formatIndexPath(CmdLineArgs args) {
116116
return INDEX_DIR + "/" + args.docVectors().get(0).getFileName() + "-" + String.join("-", suffix) + ".index";
117117
}
118118

119-
static Codec createCodec(CmdLineArgs args) {
119+
static Codec createCodec(TestConfiguration args) {
120120
final KnnVectorsFormat format;
121121
int quantizeBits = args.quantizeBits();
122122
if (args.indexType() == IndexType.IVF) {
@@ -233,7 +233,7 @@ public static void main(String[] args) throws Exception {
233233
System.out.println("Where <config-file> is a JSON file containing one or more configurations for the KNN index tester.");
234234
System.out.println("--warnUp is the number of warm iterations ");
235235
System.out.println("An example configuration object: ");
236-
System.out.println(CmdLineArgs.exampleFormatForHelp());
236+
System.out.println(TestConfiguration.exampleFormatForHelp());
237237
return;
238238
}
239239

@@ -244,83 +244,91 @@ public static void main(String[] args) throws Exception {
244244

245245
logger.info("Using configuration file: " + jsonConfigPath);
246246
// Parse the JSON config file to get command line arguments
247-
// This assumes that CmdLineArgs.fromXContent is implemented to parse the JSON file
248-
List<CmdLineArgs> cmdLineArgsList = new ArrayList<>();
247+
// This assumes that the JSON file is the correct format
248+
List<TestConfiguration> testConfigurationList = new ArrayList<>();
249249
try (
250250
InputStream jsonStream = Files.newInputStream(jsonConfigPath);
251251
XContentParser parser = XContentType.JSON.xContent().createParser(XContentParserConfiguration.EMPTY, jsonStream)
252252
) {
253253
// check if the parser is at the start of an object if so, we only have one set of arguments
254254
if (parser.currentToken() == null && parser.nextToken() == XContentParser.Token.START_OBJECT) {
255-
cmdLineArgsList.add(CmdLineArgs.fromXContent(parser));
255+
testConfigurationList.add(TestConfiguration.fromXContent(parser));
256256
} else if (parser.currentToken() == XContentParser.Token.START_ARRAY) {
257257
// if the parser is at the start of an array, we have multiple sets of arguments
258258
while (parser.nextToken() != XContentParser.Token.END_ARRAY) {
259-
cmdLineArgsList.add(CmdLineArgs.fromXContent(parser));
259+
testConfigurationList.add(TestConfiguration.fromXContent(parser));
260260
}
261261
} else {
262262
throw new IllegalArgumentException("Invalid JSON format in config file: " + jsonConfigPath);
263263
}
264264
}
265265
FormattedResults formattedResults = new FormattedResults();
266266

267-
for (CmdLineArgs cmdLineArgs : cmdLineArgsList) {
268-
String indexType = cmdLineArgs.indexType().name().toLowerCase(Locale.ROOT);
269-
Results indexResults = new Results(cmdLineArgs.docVectors().get(0).getFileName().toString(), indexType, cmdLineArgs.numDocs());
270-
Results[] results = new Results[cmdLineArgs.numberOfSearchRuns()];
267+
for (TestConfiguration testConfiguration : testConfigurationList) {
268+
String indexType = testConfiguration.indexType().name().toLowerCase(Locale.ROOT);
269+
Results indexResults = new Results(
270+
testConfiguration.docVectors().get(0).getFileName().toString(),
271+
indexType,
272+
testConfiguration.numDocs()
273+
);
274+
Results[] results = new Results[testConfiguration.numberOfSearchRuns()];
271275
for (int i = 0; i < results.length; i++) {
272-
results[i] = new Results(cmdLineArgs.docVectors().get(0).getFileName().toString(), indexType, cmdLineArgs.numDocs());
276+
results[i] = new Results(
277+
testConfiguration.docVectors().get(0).getFileName().toString(),
278+
indexType,
279+
testConfiguration.numDocs()
280+
);
273281
}
274282
logger.info("Running with Java: " + Runtime.version());
275-
logger.info("Running KNN index tester with arguments: " + cmdLineArgs);
276-
Codec codec = createCodec(cmdLineArgs);
277-
Path indexPath = PathUtils.get(formatIndexPath(cmdLineArgs));
278-
MergePolicy mergePolicy = getMergePolicy(cmdLineArgs);
279-
if (cmdLineArgs.reindex() || cmdLineArgs.forceMerge()) {
283+
logger.info("Running KNN index tester with arguments: " + testConfiguration);
284+
Codec codec = createCodec(testConfiguration);
285+
Path indexPath = PathUtils.get(formatIndexPath(testConfiguration));
286+
MergePolicy mergePolicy = getMergePolicy(testConfiguration);
287+
if (testConfiguration.reindex() || testConfiguration.forceMerge()) {
280288
KnnIndexer knnIndexer = new KnnIndexer(
281-
cmdLineArgs.docVectors(),
289+
testConfiguration.docVectors(),
282290
indexPath,
283291
codec,
284-
cmdLineArgs.indexThreads(),
285-
cmdLineArgs.vectorEncoding(),
286-
cmdLineArgs.dimensions(),
287-
cmdLineArgs.vectorSpace(),
288-
cmdLineArgs.numDocs(),
292+
testConfiguration.indexThreads(),
293+
testConfiguration.vectorEncoding(),
294+
testConfiguration.dimensions(),
295+
testConfiguration.vectorSpace(),
296+
testConfiguration.numDocs(),
289297
mergePolicy,
290-
cmdLineArgs.writerBufferSizeInMb(),
291-
cmdLineArgs.writerMaxBufferedDocs()
298+
testConfiguration.writerBufferSizeInMb(),
299+
testConfiguration.writerMaxBufferedDocs()
292300
);
293-
if (cmdLineArgs.reindex() == false && Files.exists(indexPath) == false) {
301+
if (testConfiguration.reindex() == false && Files.exists(indexPath) == false) {
294302
throw new IllegalArgumentException("Index path does not exist: " + indexPath);
295303
}
296-
if (cmdLineArgs.reindex()) {
304+
if (testConfiguration.reindex()) {
297305
knnIndexer.createIndex(indexResults);
298306
}
299-
if (cmdLineArgs.forceMerge()) {
300-
knnIndexer.forceMerge(indexResults, cmdLineArgs.forceMergeMaxNumSegments());
307+
if (testConfiguration.forceMerge()) {
308+
knnIndexer.forceMerge(indexResults, testConfiguration.forceMergeMaxNumSegments());
301309
}
302310
}
303311
numSegments(indexPath, indexResults);
304-
if (cmdLineArgs.queryVectors() != null && cmdLineArgs.numQueries() > 0) {
312+
if (testConfiguration.queryVectors() != null && testConfiguration.numQueries() > 0) {
305313
if (parsedArgs.warmUpIterations() > 0) {
306314
logger.info("Running the searches for " + parsedArgs.warmUpIterations() + " warm up iterations");
307315
}
308316
// Warm up
309317
for (int warmUpCount = 0; warmUpCount < parsedArgs.warmUpIterations(); warmUpCount++) {
310318
for (int i = 0; i < results.length; i++) {
311319
var ignoreResults = new Results(
312-
cmdLineArgs.docVectors().get(0).getFileName().toString(),
320+
testConfiguration.docVectors().get(0).getFileName().toString(),
313321
indexType,
314-
cmdLineArgs.numDocs()
322+
testConfiguration.numDocs()
315323
);
316-
KnnSearcher knnSearcher = new KnnSearcher(indexPath, cmdLineArgs);
317-
knnSearcher.runSearch(ignoreResults, cmdLineArgs.searchParams().get(i));
324+
KnnSearcher knnSearcher = new KnnSearcher(indexPath, testConfiguration);
325+
knnSearcher.runSearch(ignoreResults, testConfiguration.searchParams().get(i));
318326
}
319327
}
320328

321329
for (int i = 0; i < results.length; i++) {
322-
KnnSearcher knnSearcher = new KnnSearcher(indexPath, cmdLineArgs);
323-
knnSearcher.runSearch(results[i], cmdLineArgs.searchParams().get(i));
330+
KnnSearcher knnSearcher = new KnnSearcher(indexPath, testConfiguration);
331+
knnSearcher.runSearch(results[i], testConfiguration.searchParams().get(i));
324332
}
325333
}
326334
formattedResults.queryResults.addAll(List.of(results));
@@ -329,7 +337,7 @@ public static void main(String[] args) throws Exception {
329337
logger.info("Results: \n" + formattedResults);
330338
}
331339

332-
private static MergePolicy getMergePolicy(CmdLineArgs args) {
340+
private static MergePolicy getMergePolicy(TestConfiguration args) {
333341
MergePolicy mergePolicy = null;
334342
if (args.mergePolicy() != null) {
335343
if (args.mergePolicy() == MergePolicyType.TIERED) {

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

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -107,20 +107,20 @@ class KnnSearcher {
107107
private final VectorSimilarityFunction similarityFunction;
108108
private final VectorEncoding vectorEncoding;
109109

110-
KnnSearcher(Path indexPath, CmdLineArgs cmdLineArgs) {
111-
this.docPath = cmdLineArgs.docVectors();
110+
KnnSearcher(Path indexPath, TestConfiguration testConfiguration) {
111+
this.docPath = testConfiguration.docVectors();
112112
this.indexPath = indexPath;
113-
this.queryPath = cmdLineArgs.queryVectors();
114-
this.numDocs = cmdLineArgs.numDocs();
115-
this.numQueryVectors = cmdLineArgs.numQueries();
116-
this.dim = cmdLineArgs.dimensions();
117-
this.similarityFunction = cmdLineArgs.vectorSpace();
118-
this.vectorEncoding = cmdLineArgs.vectorEncoding();
113+
this.queryPath = testConfiguration.queryVectors();
114+
this.numDocs = testConfiguration.numDocs();
115+
this.numQueryVectors = testConfiguration.numQueries();
116+
this.dim = testConfiguration.dimensions();
117+
this.similarityFunction = testConfiguration.vectorSpace();
118+
this.vectorEncoding = testConfiguration.vectorEncoding();
119119
if (numQueryVectors <= 0) {
120120
throw new IllegalArgumentException("numQueryVectors must be > 0");
121121
}
122-
this.indexType = cmdLineArgs.indexType();
123-
this.randomSeed = cmdLineArgs.seed();
122+
this.indexType = testConfiguration.indexType();
123+
this.randomSeed = testConfiguration.seed();
124124
}
125125

126126
void runSearch(KnnIndexTester.Results finalResults, SearchParameters searchParameters) throws IOException {

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

Lines changed: 18 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -36,15 +36,15 @@ public record SearchParameters(
3636
);
3737

3838
static {
39-
PARSER.declareInt(Builder::setNumCandidates, CmdLineArgs.NUM_CANDIDATES_FIELD);
40-
PARSER.declareInt(Builder::setTopK, CmdLineArgs.K_FIELD);
41-
PARSER.declareDouble(Builder::setVisitPercentage, CmdLineArgs.VISIT_PERCENTAGE_FIELD);
42-
PARSER.declareFloat(Builder::setOverSamplingFactor, CmdLineArgs.OVER_SAMPLING_FACTOR_FIELD);
43-
PARSER.declareInt(Builder::setSearchThreads, CmdLineArgs.SEARCH_THREADS_FIELD);
44-
PARSER.declareInt(Builder::setNumSearchers, CmdLineArgs.NUM_SEARCHERS_FIELD);
45-
PARSER.declareBoolean(Builder::setEarlyTermination, CmdLineArgs.EARLY_TERMINATION_FIELD);
46-
PARSER.declareBoolean(Builder::setFilterCached, CmdLineArgs.FILTER_CACHED);
47-
PARSER.declareFloat(Builder::setFilterSelectivity, CmdLineArgs.FILTER_SELECTIVITY_FIELD);
39+
PARSER.declareInt(Builder::setNumCandidates, TestConfiguration.NUM_CANDIDATES_FIELD);
40+
PARSER.declareInt(Builder::setTopK, TestConfiguration.K_FIELD);
41+
PARSER.declareDouble(Builder::setVisitPercentage, TestConfiguration.VISIT_PERCENTAGE_FIELD);
42+
PARSER.declareFloat(Builder::setOverSamplingFactor, TestConfiguration.OVER_SAMPLING_FACTOR_FIELD);
43+
PARSER.declareInt(Builder::setSearchThreads, TestConfiguration.SEARCH_THREADS_FIELD);
44+
PARSER.declareInt(Builder::setNumSearchers, TestConfiguration.NUM_SEARCHERS_FIELD);
45+
PARSER.declareBoolean(Builder::setEarlyTermination, TestConfiguration.EARLY_TERMINATION_FIELD);
46+
PARSER.declareBoolean(Builder::setFilterCached, TestConfiguration.FILTER_CACHED);
47+
PARSER.declareFloat(Builder::setFilterSelectivity, TestConfiguration.FILTER_SELECTIVITY_FIELD);
4848
}
4949

5050
static SearchParameters.Builder fromXContent(XContentParser parser) {
@@ -156,31 +156,31 @@ public SearchParameters buildWithDefaults(SearchParameters defaults) {
156156
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
157157
builder.startObject();
158158
if (numCandidates != null) {
159-
builder.field(CmdLineArgs.NUM_CANDIDATES_FIELD.getPreferredName(), numCandidates);
159+
builder.field(TestConfiguration.NUM_CANDIDATES_FIELD.getPreferredName(), numCandidates);
160160
}
161161
if (topK != null) {
162-
builder.field(CmdLineArgs.K_FIELD.getPreferredName(), topK);
162+
builder.field(TestConfiguration.K_FIELD.getPreferredName(), topK);
163163
}
164164
if (visitPercentage != null) {
165-
builder.field(CmdLineArgs.VISIT_PERCENTAGE_FIELD.getPreferredName(), visitPercentage);
165+
builder.field(TestConfiguration.VISIT_PERCENTAGE_FIELD.getPreferredName(), visitPercentage);
166166
}
167167
if (overSamplingFactor != null) {
168-
builder.field(CmdLineArgs.OVER_SAMPLING_FACTOR_FIELD.getPreferredName(), overSamplingFactor);
168+
builder.field(TestConfiguration.OVER_SAMPLING_FACTOR_FIELD.getPreferredName(), overSamplingFactor);
169169
}
170170
if (searchThreads != null) {
171-
builder.field(CmdLineArgs.SEARCH_THREADS_FIELD.getPreferredName(), searchThreads);
171+
builder.field(TestConfiguration.SEARCH_THREADS_FIELD.getPreferredName(), searchThreads);
172172
}
173173
if (numSearchers != null) {
174-
builder.field(CmdLineArgs.NUM_SEARCHERS_FIELD.getPreferredName(), numSearchers);
174+
builder.field(TestConfiguration.NUM_SEARCHERS_FIELD.getPreferredName(), numSearchers);
175175
}
176176
if (filterCached != null) {
177-
builder.field(CmdLineArgs.FILTER_CACHED.getPreferredName(), filterCached);
177+
builder.field(TestConfiguration.FILTER_CACHED.getPreferredName(), filterCached);
178178
}
179179
if (filterSelectivity != null) {
180-
builder.field(CmdLineArgs.FILTER_SELECTIVITY_FIELD.getPreferredName(), filterSelectivity);
180+
builder.field(TestConfiguration.FILTER_SELECTIVITY_FIELD.getPreferredName(), filterSelectivity);
181181
}
182182
if (earlyTermination != null) {
183-
builder.field(CmdLineArgs.EARLY_TERMINATION_FIELD.getPreferredName(), earlyTermination);
183+
builder.field(TestConfiguration.EARLY_TERMINATION_FIELD.getPreferredName(), earlyTermination);
184184
}
185185
return builder.endObject();
186186
}

qa/vector/src/main/java/org/elasticsearch/test/knn/CmdLineArgs.java renamed to qa/vector/src/main/java/org/elasticsearch/test/knn/TestConfiguration.java

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@
3131
* Command line arguments for the KNN index tester.
3232
* This class encapsulates all the parameters required to run the KNN index tests.
3333
*/
34-
record CmdLineArgs(
34+
record TestConfiguration(
3535
List<Path> docVectors,
3636
Path queryVectors,
3737
int numDocs,
@@ -97,12 +97,12 @@ record CmdLineArgs(
9797
*/
9898
static final double DEFAULT_WRITER_BUFFER_MB = (JvmInfo.jvmInfo().getMem().getHeapMax().getBytes() / (1024.0 * 1024.0)) * 0.1;
9999

100-
static CmdLineArgs fromXContent(XContentParser parser) throws IOException {
100+
static TestConfiguration fromXContent(XContentParser parser) throws IOException {
101101
Builder builder = PARSER.apply(parser, null);
102102
return builder.build();
103103
}
104104

105-
static final ObjectParser<CmdLineArgs.Builder, Void> PARSER = new ObjectParser<>("cmd_line_args", false, Builder::new);
105+
static final ObjectParser<TestConfiguration.Builder, Void> PARSER = new ObjectParser<>("test_configuration", false, Builder::new);
106106

107107
static {
108108
PARSER.declareStringArray(Builder::setDocVectors, DOC_VECTORS_FIELD);
@@ -144,7 +144,7 @@ public int numberOfSearchRuns() {
144144
}
145145

146146
public static String exampleFormatForHelp() {
147-
var b = new CmdLineArgs.Builder().setDimensions(64)
147+
var b = new TestConfiguration.Builder().setDimensions(64)
148148
.setDocVectors(List.of("/doc/vectors/path"))
149149
.setQueryVectors("/query/vectors/path")
150150
.setSearchParams(
@@ -354,7 +354,7 @@ public Builder setSearchParams(List<SearchParameters.Builder> searchParams) {
354354
return this;
355355
}
356356

357-
public CmdLineArgs build() {
357+
public TestConfiguration build() {
358358
if (docVectors == null) {
359359
throw new IllegalArgumentException("Document vectors path must be provided");
360360
}
@@ -404,7 +404,7 @@ public CmdLineArgs build() {
404404
}
405405
}
406406

407-
return new CmdLineArgs(
407+
return new TestConfiguration(
408408
docVectors,
409409
queryVectors,
410410
numDocs,

0 commit comments

Comments
 (0)