Skip to content

[Backport] [3.5] Fix flaky org.opensearch.knn.index.codec.jvector.KNNJVectorTests.test.JVectorKnnIndex_mergeEnabled (2nd attempt) (#362) (#363) #365

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
reta merged 1 commit into from
Apr 7, 2026
Merged

[Backport] [3.5] Fix flaky org.opensearch.knn.index.codec.jvector.KNNJVectorTests.test.JVectorKnnIndex_mergeEnabled (2nd attempt) (#362) (#363) #365

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 49 additions & 5 deletions src/main/java/org/opensearch/knn/index/codec/jvector/JVectorFormat.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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
);
}

Expand All @@ -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
);
}

Expand All @@ -98,7 +130,10 @@ public JVectorFormat(
numberOfSubspacesPerVectorSupplier,
minBatchSizeForQuantization,
hierarchyEnabled,
leadingSegmentMergeDisabled
leadingSegmentMergeDisabled,
SIMD_POOL_MERGE,
SIMD_POOL_FLUSH,
PARALLELISM_POOL
);
}

Expand All @@ -111,7 +146,10 @@ public JVectorFormat(
Function<Integer, Integer> numberOfSubspacesPerVectorSupplier,
int minBatchSizeForQuantization,
boolean hierarchyEnabled,
boolean leadingSegmentMergeDisabled
boolean leadingSegmentMergeDisabled,
final ForkJoinPool simdPoolMerge,
final ForkJoinPool simdPoolFlush,
final ForkJoinPool parallelismPool
) {
super(name);
this.maxConn = maxConn;
Expand All @@ -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
Expand All @@ -135,7 +176,10 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
numberOfSubspacesPerVectorSupplier,
minBatchSizeForQuantization,
hierarchyEnabled,
leadingSegmentMergeDisabled
leadingSegmentMergeDisabled,
simdPoolMerge,
simdPoolFlush,
parallelismPool
);
}

Expand Down
44 changes: 25 additions & 19 deletions src/main/java/org/opensearch/knn/index/codec/jvector/JVectorWriter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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(
Expand All @@ -109,7 +110,10 @@ public JVectorWriter(
Function<Integer, Integer> 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;
Expand All @@ -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,
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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()
);

Expand All @@ -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,
Expand Down Expand Up @@ -963,7 +970,7 @@ public void merge() throws IOException {
compactOrdsToRavvOrds.length,
compactOrdsToRavvOrds,
this,
SIMD_POOL_MERGE
simdPoolMerge
);
}

Expand All @@ -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 {
Expand All @@ -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);
}
}
Expand Down Expand Up @@ -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++) {
Expand Down Expand Up @@ -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();

Expand Down
95 changes: 93 additions & 2 deletions src/test/java/org/opensearch/knn/index/codec/jvector/KNNJVectorTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;

Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)) {
Expand Down Expand Up @@ -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.
Expand Down
Loading
Loading