adjusting vector reclal testing

Author: benwtrent

test/external-modules/vector/build.gradle renamed to qa/vector/build.gradle

@@ -26,10 +26,16 @@ tasks.register("checkVec", JavaExec) {
   mainClass.set("org.elasticsearch.test.knn.KnnIndexTester")
   // TODO REGISTER ARGS
   args = [
-    "--docVectors=<path>",
+    "--docVectors=/some/path",
     "--numDocs=100000",
-    "--reindex",
-    "--queryVectors=<path>",
+    //"--reindex",
+    //"--forceMerge",
+    "--queryVectors=/some/other/path",
     "--dimensions=768",
+    "--overSamplingFactor=3",
+    "--nProbe=5"
   ]
+  if (buildParams.getRuntimeJavaVersion().map { it.majorVersion.toInteger() }.get() >= 21) {
+    jvmArgs '--add-modules=jdk.incubator.vector', '--enable-native-access=ALL-UNNAMED', '-Djava.util.concurrent.ForkJoinPool.common.parallelism=8'
+  }
 }

test/external-modules/vector/src/main/java/module-info.java renamed to qa/vector/src/main/java/module-info.java

@@ -135,7 +135,7 @@ static CmdLineArgs parse(String[] args) {
         int hnswEfConstruction = 100;
         int searchThreads = 1;
         int indexThreads = 1;
-        boolean reindex = true;
+        boolean reindex = false;
         boolean forceMerge = false;
         VectorSimilarityFunction vectorSpace = VectorSimilarityFunction.EUCLIDEAN;
         // 32 means no quantization
@@ -205,10 +205,10 @@ public static void main(String[] args) throws Exception {
         StringBuilder resultHeaders = new StringBuilder();
         StringBuilder resultValues = new StringBuilder();
         // indicate params used for index creation
-        resultHeaders.append("index_type,");
-        resultValues.append(cmdLineArgs.indexType).append(",");
-        resultHeaders.append("num_docs,");
-        resultValues.append(cmdLineArgs.numDocs).append(",");
+        resultHeaders.append("index_type\t");
+        resultValues.append(cmdLineArgs.indexType).append("\t");
+        resultHeaders.append("num_docs\t");
+        resultValues.append(cmdLineArgs.numDocs).append("\t");
         if (cmdLineArgs.reindex || cmdLineArgs.forceMerge) {
             KnnIndexer knnIndexer = new KnnIndexer(
                 cmdLineArgs.docVectors,
@@ -230,28 +230,27 @@ public static void main(String[] args) throws Exception {
             }
         }
         if (indexCreationTimeMS > 0) {
-            resultHeaders.append("index_time(ms)").append(",");
-            resultValues.append(indexCreationTimeMS).append(",");
+            resultHeaders.append("index_time(ms)").append("\t");
+            resultValues.append(indexCreationTimeMS).append("\t");
         }
         if (forceMergeTimeMS > 0) {
-            resultHeaders.append("force_merge_time(ms)").append(",");
-            resultValues.append(forceMergeTimeMS).append(",");
+            resultHeaders.append("force_merge_time(ms)").append("\t");
+            resultValues.append(forceMergeTimeMS).append("\t");
         }
-        resultHeaders.append("num_segments,");
-        resultValues.append(numSegments).append(",");
+        resultHeaders.append("num_segments\t");
+        resultValues.append(numSegments).append("\t");
 
         if (cmdLineArgs.queryVectors != null) {
             KnnSearcher knnSearcher = new KnnSearcher(indexPath, cmdLineArgs);
             KnnSearcher.SearcherResults results = knnSearcher.runSearch();
-            resultHeaders.append("latency(ms),");
-            resultValues.append(results.avgLatency()).append(",");
-            resultHeaders.append("qps,");
-            resultValues.append(results.qps()).append(",");
-            resultHeaders.append("recall,");
-            resultValues.append(results.avgRecall()).append(",");
-            resultHeaders.append("visited,");
-            resultValues.append(results.averageVisited()).append(",");
-            return;
+            resultHeaders.append("latency(ms)\t");
+            resultValues.append(results.avgLatency()).append("\t");
+            resultHeaders.append("qps\t");
+            resultValues.append(results.qps()).append("\t");
+            resultHeaders.append("recall\t");
+            resultValues.append(results.avgRecall()).append("\t");
+            resultHeaders.append("visited\t");
+            resultValues.append(results.averageVisited()).append("\t");
         }
         System.out.println(resultHeaders);
         System.out.println(resultValues);

test/external-modules/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexTester.java renamed to qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexTester.java

@@ -213,7 +213,7 @@ private int[][] getOrCalculateExactNN() throws IOException {
             36
         );
         String nnFileName = "nn-" + hash + ".bin";
-        Path nnPath = Paths.get(nnFileName);
+        Path nnPath = Paths.get("target/" + nnFileName);
         if (Files.exists(nnPath)) {
             logger.info("read pre-cached exact match vectors from cache file \"" + nnPath + "\"");
             return readExactNN(nnPath);
@@ -352,7 +352,7 @@ private int[][] computeExactNN(Path queryPath) throws IOException {
                 for (int i = 0; i < numQueryVectors; i++) {
                     float[] queryVector = new float[dim];
                     queryReader.next(queryVector);
-                    tasks.add(new ComputeNNFloatTask(i, queryVector, result, reader, similarityFunction));
+                    tasks.add(new ComputeNNFloatTask(i, topK, queryVector, result, reader, similarityFunction));
                 }
                 ForkJoinPool.commonPool().invokeAll(tasks);
             }
@@ -384,19 +384,27 @@ static class ComputeNNFloatTask implements Callable<Void> {
         private final int[][] result;
         private final IndexReader reader;
         private final VectorSimilarityFunction similarityFunction;
+        private final int topK;
 
-        ComputeNNFloatTask(int queryOrd, float[] query, int[][] result, IndexReader reader, VectorSimilarityFunction similarityFunction) {
+        ComputeNNFloatTask(
+            int queryOrd,
+            int topK,
+            float[] query,
+            int[][] result,
+            IndexReader reader,
+            VectorSimilarityFunction similarityFunction
+        ) {
             this.queryOrd = queryOrd;
             this.query = query;
             this.result = result;
             this.reader = reader;
             this.similarityFunction = similarityFunction;
+            this.topK = topK;
         }
 
         @Override
         public Void call() {
             IndexSearcher searcher = new IndexSearcher(reader);
-            int topK = result[0].length;
             try {
                 var queryVector = new ConstKnnFloatValueSource(query);
                 var docVectors = new FloatKnnVectorFieldSource(VECTOR_FIELD);

test/external-modules/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexer.java renamed to qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexer.java

@@ -172,4 +172,4 @@ if (extraProjects.exists()) {
   }
 }
 
-include 'test:external-modules:vector'
+include 'qa:vector'

test/external-modules/vector/src/main/java/org/elasticsearch/test/knn/KnnSearcher.java renamed to qa/vector/src/main/java/org/elasticsearch/test/knn/KnnSearcher.java

@@ -8,14 +8,11 @@
  */
 
 subprojects {
-  if (project.name != 'vector') {
-
     apply plugin: 'elasticsearch.base-internal-es-plugin'
 
     esplugin {
       name = it.name
       licenseFile = layout.settingsDirectory.file('licenses/AGPL-3.0+SSPL-1.0+ELASTIC-LICENSE-2.0.txt').asFile
       noticeFile = layout.settingsDirectory.file('NOTICE.txt').asFile
     }
-  }
 }