diff --git a/src/main/java/org/opensearch/knn/index/codec/jvector/JVectorFormat.java b/src/main/java/org/opensearch/knn/index/codec/jvector/JVectorFormat.java index 9916392a..afa7deb9 100644 --- a/src/main/java/org/opensearch/knn/index/codec/jvector/JVectorFormat.java +++ b/src/main/java/org/opensearch/knn/index/codec/jvector/JVectorFormat.java @@ -46,6 +46,9 @@ public class JVectorFormat extends KnnVectorsFormat { private final float neighborOverflow; private final boolean hierarchyEnabled; private final boolean leadingSegmentMergeDisabled; + private final ForkJoinPool simdPoolMerge; + private final ForkJoinPool simdPoolFlush; + private final ForkJoinPool parallelismPool; public JVectorFormat() { this( @@ -57,7 +60,10 @@ public JVectorFormat() { JVectorFormat::getDefaultNumberOfSubspacesPerVector, KNNConstants.DEFAULT_MINIMUM_BATCH_SIZE_FOR_QUANTIZATION, KNNConstants.DEFAULT_HIERARCHY_ENABLED, - KNNConstants.DEFAULT_LEADING_SEGMENT_MERGE_DISABLED + KNNConstants.DEFAULT_LEADING_SEGMENT_MERGE_DISABLED, + SIMD_POOL_MERGE, + SIMD_POOL_FLUSH, + PARALLELISM_POOL ); } @@ -75,7 +81,33 @@ public JVectorFormat(int minBatchSizeForQuantization, boolean leadingSegmentMerg JVectorFormat::getDefaultNumberOfSubspacesPerVector, minBatchSizeForQuantization, KNNConstants.DEFAULT_HIERARCHY_ENABLED, - leadingSegmentMergeDisabled + leadingSegmentMergeDisabled, + SIMD_POOL_MERGE, + SIMD_POOL_FLUSH, + PARALLELISM_POOL + ); + } + + public JVectorFormat( + int minBatchSizeForQuantization, + boolean leadingSegmentMergeDisabled, + final ForkJoinPool simdPoolMerge, + final ForkJoinPool simdPoolFlush, + final ForkJoinPool parallelismPool + ) { + this( + NAME, + DEFAULT_MAX_CONN, + DEFAULT_BEAM_WIDTH, + KNNConstants.DEFAULT_NEIGHBOR_OVERFLOW_VALUE.floatValue(), + KNNConstants.DEFAULT_ALPHA_VALUE.floatValue(), + JVectorFormat::getDefaultNumberOfSubspacesPerVector, + minBatchSizeForQuantization, + KNNConstants.DEFAULT_HIERARCHY_ENABLED, + leadingSegmentMergeDisabled, + simdPoolMerge, + simdPoolFlush, + parallelismPool ); } @@ -98,7 +130,10 @@ public JVectorFormat( numberOfSubspacesPerVectorSupplier, minBatchSizeForQuantization, hierarchyEnabled, - leadingSegmentMergeDisabled + leadingSegmentMergeDisabled, + SIMD_POOL_MERGE, + SIMD_POOL_FLUSH, + PARALLELISM_POOL ); } @@ -111,7 +146,10 @@ public JVectorFormat( Function numberOfSubspacesPerVectorSupplier, int minBatchSizeForQuantization, boolean hierarchyEnabled, - boolean leadingSegmentMergeDisabled + boolean leadingSegmentMergeDisabled, + final ForkJoinPool simdPoolMerge, + final ForkJoinPool simdPoolFlush, + final ForkJoinPool parallelismPool ) { super(name); this.maxConn = maxConn; @@ -122,6 +160,9 @@ public JVectorFormat( this.neighborOverflow = neighborOverflow; this.hierarchyEnabled = hierarchyEnabled; this.leadingSegmentMergeDisabled = leadingSegmentMergeDisabled; + this.simdPoolMerge = simdPoolMerge; + this.simdPoolFlush = simdPoolFlush; + this.parallelismPool = parallelismPool; } @Override @@ -135,7 +176,10 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException numberOfSubspacesPerVectorSupplier, minBatchSizeForQuantization, hierarchyEnabled, - leadingSegmentMergeDisabled + leadingSegmentMergeDisabled, + simdPoolMerge, + simdPoolFlush, + parallelismPool ); } diff --git a/src/main/java/org/opensearch/knn/index/codec/jvector/JVectorWriter.java b/src/main/java/org/opensearch/knn/index/codec/jvector/JVectorWriter.java index 8595f99e..328df646 100644 --- a/src/main/java/org/opensearch/knn/index/codec/jvector/JVectorWriter.java +++ b/src/main/java/org/opensearch/knn/index/codec/jvector/JVectorWriter.java @@ -46,9 +46,6 @@ import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED; import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.readVectorEncoding; -import static org.opensearch.knn.index.codec.jvector.JVectorFormat.PARALLELISM_POOL; -import static org.opensearch.knn.index.codec.jvector.JVectorFormat.SIMD_POOL_FLUSH; -import static org.opensearch.knn.index.codec.jvector.JVectorFormat.SIMD_POOL_MERGE; /** * JVectorWriter is responsible for writing vector data into index segments using the JVector library. @@ -98,6 +95,10 @@ public class JVectorWriter extends KnnVectorsWriter { private final boolean hierarchyEnabled; private final boolean leadingSegmentMergeDisabled; + private final ForkJoinPool simdPoolMerge; + private final ForkJoinPool simdPoolFlush; + private final ForkJoinPool parallelismPool; + private boolean finished = false; public JVectorWriter( @@ -109,7 +110,10 @@ public JVectorWriter( Function numberOfSubspacesPerVectorSupplier, int minimumBatchSizeForQuantization, boolean hierarchyEnabled, - boolean leadingSegmentMergeDisabled + boolean leadingSegmentMergeDisabled, + final ForkJoinPool simdPoolMerge, + final ForkJoinPool simdPoolFlush, + final ForkJoinPool parallelismPool ) throws IOException { this.segmentWriteState = segmentWriteState; this.maxConn = maxConn; @@ -120,6 +124,9 @@ public JVectorWriter( this.minimumBatchSizeForQuantization = minimumBatchSizeForQuantization; this.hierarchyEnabled = hierarchyEnabled; this.leadingSegmentMergeDisabled = leadingSegmentMergeDisabled; + this.simdPoolMerge = simdPoolMerge; + this.simdPoolFlush = simdPoolFlush; + this.parallelismPool = parallelismPool; String metaFileName = IndexFileNames.segmentFileName( segmentWriteState.segmentInfo.name, @@ -247,7 +254,7 @@ public void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException { randomAccessVectorValues, fieldInfo, segmentWriteState.segmentInfo.name, - SIMD_POOL_FLUSH + simdPoolFlush ); writeField(field.fieldInfo, randomAccessVectorValues, pqVectors, graphNodeIdToDocMap, graph); @@ -388,7 +395,7 @@ private PQVectors getPQVectors(RandomAccessVectorValues randomAccessVectorValues numberOfClustersPerSubspace, // number of centroids per subspace vectorSimilarityFunction == VectorSimilarityFunction.EUCLIDEAN, // center the dataset UNWEIGHTED, - SIMD_POOL_MERGE, + simdPoolMerge, ForkJoinPool.commonPool() ); @@ -398,7 +405,7 @@ private PQVectors getPQVectors(RandomAccessVectorValues randomAccessVectorValues KNNCounter.KNN_QUANTIZATION_TRAINING_TIME.add(trainingTime); log.info("Encoding and building PQ vectors for field {} for {} vectors", fieldName, randomAccessVectorValues.size()); // PQVectors pqVectors = pq.encodeAll(randomAccessVectorValues, SIMD_POOL); - PQVectors pqVectors = PQVectors.encodeAndBuild(pq, randomAccessVectorValues.size(), randomAccessVectorValues, SIMD_POOL_MERGE); + PQVectors pqVectors = PQVectors.encodeAndBuild(pq, randomAccessVectorValues.size(), randomAccessVectorValues, simdPoolMerge); log.info( "Encoded and built PQ vectors for field {}, original size: {} bytes, compressed size: {} bytes", fieldName, @@ -963,7 +970,7 @@ public void merge() throws IOException { compactOrdsToRavvOrds.length, compactOrdsToRavvOrds, this, - SIMD_POOL_MERGE + simdPoolMerge ); } @@ -980,7 +987,7 @@ public void merge() throws IOException { fieldName ); var bsp = BuildScoreProvider.randomAccessScoreProvider(compactRavv, getVectorSimilarityFunction(fieldInfo)); - var graph = getGraph(bsp, compactRavv, fieldInfo, segmentWriteState.segmentInfo.name, SIMD_POOL_MERGE); + var graph = getGraph(bsp, compactRavv, fieldInfo, segmentWriteState.segmentInfo.name, simdPoolMerge); writeField(fieldInfo, compactRavv, null, compactOrdToDocMap, graph); } } else { @@ -989,7 +996,7 @@ public void merge() throws IOException { var buildScoreProvider = BuildScoreProvider.pqBuildScoreProvider(getVectorSimilarityFunction(fieldInfo), compactPqVectors); // Pre-init the diversity provider here to avoid doing it lazily (as it could block the SIMD threads) buildScoreProvider.diversityProviderFor(0); - var graph = getGraph(buildScoreProvider, compactRavv, fieldInfo, segmentWriteState.segmentInfo.name, SIMD_POOL_MERGE); + var graph = getGraph(buildScoreProvider, compactRavv, fieldInfo, segmentWriteState.segmentInfo.name, simdPoolMerge); writeField(fieldInfo, compactRavv, compactPqVectors, compactOrdToDocMap, graph); } } @@ -1140,20 +1147,18 @@ private boolean tryLeadingSegmentMerge() throws IOException { degreeOverflow, alpha, true, - SIMD_POOL_MERGE, - PARALLELISM_POOL + simdPoolMerge, + parallelismPool ) ) { var vv = heapRavv.threadLocalSupplier(); // parallel graph construction from the merge documents Ids - SIMD_POOL_MERGE.submit( - () -> IntStream.range(leadingGraph.getIdUpperBound(), heapRavv.size()).parallel().forEach(ord -> { - assert heapToGlobalRavvOrds[ord] != GraphNodeIdToDocMap.NO_VECTOR_OR_DELETED_DOC - : "Should be a valid graph node / vector"; - builder.addGraphNode(ord, vv.get().getVector(ord)); - }) - ).join(); + simdPoolMerge.submit(() -> IntStream.range(leadingGraph.getIdUpperBound(), heapRavv.size()).parallel().forEach(ord -> { + assert heapToGlobalRavvOrds[ord] != GraphNodeIdToDocMap.NO_VECTOR_OR_DELETED_DOC + : "Should be a valid graph node / vector"; + builder.addGraphNode(ord, vv.get().getVector(ord)); + })).join(); // mark deleted nodes for (int i = 0; i < numBaseVectors; i++) { @@ -1249,6 +1254,7 @@ public OnHeapGraphIndex getGraph( graphIndexBuilder.addGraphNode(ord, vv.get().getVector(ord)); })).join(); graphIndexBuilder.cleanup(); + graphIndex = (OnHeapGraphIndex) graphIndexBuilder.getGraph(); final long end = Clock.systemDefaultZone().millis(); diff --git a/src/test/java/org/opensearch/knn/index/codec/jvector/KNNJVectorTests.java b/src/test/java/org/opensearch/knn/index/codec/jvector/KNNJVectorTests.java index a8a288ff..977e893c 100644 --- a/src/test/java/org/opensearch/knn/index/codec/jvector/KNNJVectorTests.java +++ b/src/test/java/org/opensearch/knn/index/codec/jvector/KNNJVectorTests.java @@ -14,7 +14,9 @@ import org.apache.lucene.store.FSDirectory; import org.apache.lucene.tests.index.RandomIndexWriter; import org.apache.lucene.tests.util.LuceneTestCase; +import org.junit.After; import org.junit.Assert; +import org.junit.Before; import org.junit.Test; import org.opensearch.knn.TestUtils; import org.opensearch.knn.common.KNNConstants; @@ -28,6 +30,7 @@ import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import static org.opensearch.knn.common.KNNConstants.DEFAULT_LEADING_SEGMENT_MERGE_DISABLED; import static org.opensearch.knn.common.KNNConstants.DEFAULT_MINIMUM_BATCH_SIZE_FOR_QUANTIZATION; import static org.opensearch.knn.index.engine.CommonTestUtils.getCodec; @@ -43,6 +46,22 @@ public class KNNJVectorTests extends LuceneTestCase { private static final String TEST_FIELD = "test_field"; private static final String TEST_ID_FIELD = "id"; + private ForkJoinPool singleThreadGraphMergePool; + + @Before + public void setUp() throws Exception { + super.setUp(); + singleThreadGraphMergePool = new ForkJoinPool(1); /* single threaded */ + } + + @After + public void tearDown() throws Exception { + super.tearDown(); + singleThreadGraphMergePool.shutdown(); + if (singleThreadGraphMergePool.awaitTermination(30, TimeUnit.SECONDS) == false) { + singleThreadGraphMergePool.shutdownNow(); + } + } /** * Test to verify that the JVector codec is able to successfully search for the nearest neighbours @@ -185,7 +204,9 @@ public void test_sorted_index() throws IOException { final String sortFieldName = "sorted_field"; IndexWriterConfig indexWriterConfig = LuceneTestCase.newIndexWriterConfig(); indexWriterConfig.setUseCompoundFile(false); - indexWriterConfig.setCodec(getCodec()); + indexWriterConfig.setCodec( + getCodec(DEFAULT_MINIMUM_BATCH_SIZE_FOR_QUANTIZATION, DEFAULT_LEADING_SEGMENT_MERGE_DISABLED, singleThreadGraphMergePool) + ); indexWriterConfig.setMergePolicy(new ForceMergesOnlyMergePolicy()); // Add index sorting configuration indexWriterConfig.setIndexSort(new Sort(new SortField(sortFieldName, SortField.Type.INT, true))); // true = reverse order @@ -319,9 +340,12 @@ public void testJVectorKnnIndex_mergeEnabled() throws IOException { int totalNumberOfDocs = 10; IndexWriterConfig indexWriterConfig = LuceneTestCase.newIndexWriterConfig(); indexWriterConfig.setUseCompoundFile(false); - indexWriterConfig.setCodec(getCodec()); + indexWriterConfig.setCodec( + getCodec(DEFAULT_MINIMUM_BATCH_SIZE_FOR_QUANTIZATION, DEFAULT_LEADING_SEGMENT_MERGE_DISABLED, singleThreadGraphMergePool) + ); indexWriterConfig.setMergePolicy(new ForceMergesOnlyMergePolicy()); indexWriterConfig.setMergeScheduler(new SerialMergeScheduler()); + indexWriterConfig.setMaxBufferedDocs(totalNumberOfDocs); final Path indexPath = createTempDir(); log.info("Index path: {}", indexPath); try (FSDirectory dir = FSDirectory.open(indexPath); IndexWriter w = new IndexWriter(dir, indexWriterConfig)) { @@ -373,6 +397,73 @@ public void testJVectorKnnIndex_mergeEnabled() throws IOException { } } + @Test + public void testJVectorKnnIndex_mergeDisabled() throws IOException { + int k = 3; // The number of nearest neighbours to gather + int totalNumberOfDocs = 10; + IndexWriterConfig indexWriterConfig = LuceneTestCase.newIndexWriterConfig(); + indexWriterConfig.setUseCompoundFile(false); + indexWriterConfig.setCodec( + getCodec(DEFAULT_MINIMUM_BATCH_SIZE_FOR_QUANTIZATION, DEFAULT_LEADING_SEGMENT_MERGE_DISABLED, singleThreadGraphMergePool) + ); + indexWriterConfig.setMergePolicy(NoMergePolicy.INSTANCE); + indexWriterConfig.setMergeScheduler(new SerialMergeScheduler()); + indexWriterConfig.setMaxBufferedDocs(10); + log.info("Max buffered docs: {}", indexWriterConfig.getMaxBufferedDocs()); + + final Path indexPath = createTempDir(); + final float[] target = new float[] { 0.0f, 0.0f }; + + log.info("Index path: {}", indexPath); + try (FSDirectory dir = FSDirectory.open(indexPath); IndexWriter w = new IndexWriter(dir, indexWriterConfig)) { + for (int i = 1; i < totalNumberOfDocs + 1; i++) { + final float[] source = new float[] { 0.0f, 1.0f * i }; + final Document doc = new Document(); + doc.add(new KnnFloatVectorField("test_field", source, VectorSimilarityFunction.EUCLIDEAN)); + doc.add(new StringField("my_doc_id", Integer.toString(i, 10), Field.Store.YES)); + w.addDocument(doc); + } + log.info("Done writing all files to the file system"); + + w.commit(); + w.flush(); + } + try (FSDirectory dir = FSDirectory.open(indexPath); IndexReader reader = DirectoryReader.open(dir)) { + log.info("We should now have 1 segment with 10 documents"); + Assert.assertEquals(1, reader.getContext().leaves().size()); + Assert.assertEquals(totalNumberOfDocs, reader.numDocs()); + + final IndexSearcher searcher = newSearcher(reader); + KnnFloatVectorQuery knnFloatVectorQuery = getJVectorKnnFloatVectorQuery("test_field", target, k, new MatchAllDocsQuery()); + TopDocs topDocs = searcher.search(knnFloatVectorQuery, k); + assertEquals(k, topDocs.totalHits.value()); + Document doc = reader.storedFields().document(topDocs.scoreDocs[0].doc); + + assertEquals("1", doc.get("my_doc_id")); + Assert.assertEquals( + VectorSimilarityFunction.EUCLIDEAN.compare(target, new float[] { 0.0f, 1.0f }), + topDocs.scoreDocs[0].score, + 0.001f + ); + doc = reader.storedFields().document(topDocs.scoreDocs[1].doc); + assertEquals("2", doc.get("my_doc_id")); + Assert.assertEquals( + VectorSimilarityFunction.EUCLIDEAN.compare(target, new float[] { 0.0f, 2.0f }), + topDocs.scoreDocs[1].score, + 0.001f + ); + doc = reader.storedFields().document(topDocs.scoreDocs[2].doc); + assertEquals("3", doc.get("my_doc_id")); + Assert.assertEquals( + VectorSimilarityFunction.EUCLIDEAN.compare(target, new float[] { 0.0f, 3.0f }), + topDocs.scoreDocs[2].score, + 0.001f + ); + log.info("successfully completed search tests"); + + } + } + /** * Test to verify that the jVector codec is able to successfully search for the nearest neighbors * in the index. diff --git a/src/test/java/org/opensearch/knn/index/engine/CommonTestUtils.java b/src/test/java/org/opensearch/knn/index/engine/CommonTestUtils.java index 1f856ff5..12c91857 100644 --- a/src/test/java/org/opensearch/knn/index/engine/CommonTestUtils.java +++ b/src/test/java/org/opensearch/knn/index/engine/CommonTestUtils.java @@ -146,6 +146,27 @@ public KnnVectorsFormat getKnnVectorsFormatForField(String field) { }; } + public static Codec getCodec(int minBatchSizeForQuantization, boolean leadingSegmentMergeDisabled, ForkJoinPool graphMergePool) { + return new FilterCodec(KNNCodecVersion.V_10_03_0.getCodecName(), new Lucene103Codec()) { + @Override + public KnnVectorsFormat knnVectorsFormat() { + return new PerFieldKnnVectorsFormat() { + + @Override + public KnnVectorsFormat getKnnVectorsFormatForField(String field) { + return new JVectorFormat( + minBatchSizeForQuantization, + leadingSegmentMergeDisabled, + graphMergePool, + graphMergePool, + graphMergePool + ); + } + }; + } + }; + } + /** * Get Stats from KNN Plugin */