iter

Author: benwtrent

qa/vector/build.gradle

@@ -15,6 +15,10 @@ tasks.named("dependencyLicenses").configure {
   mapping from: /lucene-.*/, to: 'lucene'
 }
 
+tasks.named('forbiddenApisMain').configure {
+  enabled = false
+}
+
 dependencies {
   api "org.apache.lucene:lucene-core:${versions.lucene}"
   api "org.apache.lucene:lucene-queries:${versions.lucene}"

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

@@ -12,6 +12,7 @@
 import org.apache.lucene.index.VectorEncoding;
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.core.PathUtils;
 import org.elasticsearch.xcontent.ObjectParser;
 import org.elasticsearch.xcontent.ParseField;
 import org.elasticsearch.xcontent.ToXContentObject;
@@ -158,12 +159,12 @@ static class Builder {
         private int dimensions;
 
         public Builder setDocVectors(String docVectors) {
-            this.docVectors = Path.of(docVectors);
+            this.docVectors = PathUtils.get(docVectors);
             return this;
         }
 
         public Builder setQueryVectors(String queryVectors) {
-            this.queryVectors = Path.of(queryVectors);
+            this.queryVectors = PathUtils.get(queryVectors);
             return this;
         }
 

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

@@ -15,6 +15,7 @@
 import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.logging.LogConfigurator;
+import org.elasticsearch.core.PathUtils;
 import org.elasticsearch.index.codec.vectors.ES813Int8FlatVectorFormat;
 import org.elasticsearch.index.codec.vectors.ES814HnswScalarQuantizedVectorsFormat;
 import org.elasticsearch.index.codec.vectors.IVFVectorsFormat;
@@ -133,8 +134,8 @@ public static void main(String[] args) throws Exception {
         }
         String jsonConfig = args[0];
         // Parse command line arguments
-        Path jsonConfigPath = Path.of(jsonConfig);
-        if (jsonConfigPath.toFile().exists() == false) {
+        Path jsonConfigPath = PathUtils.get(jsonConfig);
+        if (Files.exists(jsonConfigPath) == false) {
             throw new IllegalArgumentException("JSON config file does not exist: " + jsonConfigPath);
         }
         // Parse the JSON config file to get command line arguments
@@ -161,7 +162,7 @@ public static void main(String[] args) throws Exception {
             Results result = new Results(cmdLineArgs.indexType().name().toLowerCase(Locale.ROOT), cmdLineArgs.numDocs());
             System.out.println("Running KNN index tester with arguments: " + cmdLineArgs);
             Codec codec = createCodec(cmdLineArgs);
-            Path indexPath = Path.of(formatIndexPath(cmdLineArgs));
+            Path indexPath = PathUtils.get(formatIndexPath(cmdLineArgs));
             if (cmdLineArgs.reindex() || cmdLineArgs.forceMerge()) {
                 KnnIndexer knnIndexer = new KnnIndexer(
                     cmdLineArgs.docVectors(),
@@ -253,7 +254,7 @@ private String formatRow(String[] values, int[] widths) {
             for (int i = 0; i < values.length; i++) {
                 // Left-align text column (index_type), right-align numeric columns
                 String format = (i == 0) ? "%-" + widths[i] + "s" : "%" + widths[i] + "s";
-                row.append(String.format(format, values[i]));
+                row.append(Strings.format(format, values[i]));
 
                 // Add separation between columns
                 if (i < values.length - 1) {

qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexer.java

@@ -35,12 +35,14 @@
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.store.FSDirectory;
 import org.apache.lucene.util.PrintStreamInfoStream;
+import org.elasticsearch.common.io.Channels;
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 import java.nio.channels.FileChannel;
+import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.List;
@@ -122,11 +124,7 @@ public boolean isEnabled(String component) {
             similarityFunction
         );
 
-        if (indexPath.toFile().mkdirs()) {
-            logger.info("Created index directory: %s", indexPath);
-        } else {
-            logger.info("Index directory already exists: %s", indexPath);
-        }
+        Files.createDirectory(indexPath);
 
         long start = System.nanoTime();
         try (
@@ -149,7 +147,7 @@ public boolean isEnabled(String component) {
             );
 
             VectorReader inReader = VectorReader.create(in, dim, vectorEncoding);
-            try (ExecutorService exec = Executors.newFixedThreadPool(numIndexThreads)) {
+            try (ExecutorService exec = Executors.newFixedThreadPool(numIndexThreads, r -> new Thread(r, "KnnIndexer-Thread"))) {
                 AtomicInteger numDocsIndexed = new AtomicInteger();
                 List<Future<?>> threads = new ArrayList<>();
                 for (int i = 0; i < numIndexThreads; i++) {
@@ -294,12 +292,12 @@ void reset() throws IOException {
         }
 
         private void readNext() throws IOException {
-            int bytesRead = this.input.read(bytes);
+            int bytesRead = Channels.readFromFileChannel(this.input, this.input.position(), bytes);
             if (bytesRead < bytes.capacity()) {
                 // wrap around back to the start of the file if we hit the end:
                 logger.warn("VectorReader hit EOF when reading " + this.input + "; now wrapping around to start of file again");
                 this.input.position(0);
-                bytesRead = this.input.read(bytes);
+                bytesRead = Channels.readFromFileChannel(this.input, this.input.position(), bytes);
                 if (bytesRead < bytes.capacity()) {
                     throw new IllegalStateException(
                         "vector file " + input + " doesn't even have enough bytes for a single vector?  got bytesRead=" + bytesRead

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

@@ -40,6 +40,7 @@
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.store.FSDirectory;
 import org.apache.lucene.store.MMapDirectory;
+import org.elasticsearch.core.PathUtils;
 import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper;
 import org.elasticsearch.search.profile.query.QueryProfiler;
 import org.elasticsearch.search.vectors.ESKnnByteVectorQuery;
@@ -58,7 +59,6 @@
 import java.nio.channels.FileChannel;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import java.nio.file.Paths;
 import java.nio.file.attribute.FileTime;
 import java.util.ArrayList;
 import java.util.HashSet;
@@ -119,8 +119,10 @@ void runSearch(KnnIndexTester.Results finalResults) throws IOException {
         TopDocs[] results = new TopDocs[numQueryVectors];
         int[][] resultIds = new int[numQueryVectors][];
         long elapsed, totalCpuTimeMS, totalVisited = 0;
-        ExecutorService executorService = Executors.newFixedThreadPool(searchThreads);
-        try (FileChannel input = FileChannel.open(queryPath)) {
+        try (
+            FileChannel input = FileChannel.open(queryPath);
+            ExecutorService executorService = Executors.newFixedThreadPool(searchThreads, r -> new Thread(r, "KnnSearcher-Thread"))
+        ) {
             long queryPathSizeInBytes = input.size();
             logger.info(
                 "queryPath size: "
@@ -177,13 +179,10 @@ void runSearch(KnnIndexTester.Results finalResults) throws IOException {
                     );
                 }
             }
-        } finally {
-            executorService.shutdown();
         }
         logger.info("checking results");
         int[][] nn = getOrCalculateExactNN();
-        float recall = checkResults(resultIds, nn, topK);
-        finalResults.avgRecall = recall;
+        finalResults.avgRecall = checkResults(resultIds, nn, topK);
         finalResults.qps = (1000f * numQueryVectors) / elapsed;
         finalResults.avgLatency = (float) elapsed / numQueryVectors;
         finalResults.averageVisited = (double) totalVisited / numQueryVectors;
@@ -206,7 +205,7 @@ private int[][] getOrCalculateExactNN() throws IOException {
             36
         );
         String nnFileName = "nn-" + hash + ".bin";
-        Path nnPath = Paths.get("target/" + nnFileName);
+        Path nnPath = PathUtils.get("target/" + nnFileName);
         if (Files.exists(nnPath) && isNewer(nnPath, docPath, indexPath, queryPath)) {
             logger.info("read pre-cached exact match vectors from cache file \"" + nnPath + "\"");
             return readExactNN(nnPath);