diff --git a/CHANGELOG.md b/CHANGELOG.md index 6ea99e42..42fc0865 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), ### Bug Fixes * Revert of support deletes for incremental insertion [240](https://github.com/opensearch-project/opensearch-jvector/pull/240) * Fix visited docs tracking that led to assertions on AbstractKnnVectorQuery side [238] (https://github.com/opensearch-project/opensearch-jvector/pull/238) +* Fix leading segment merge for deletions [242](https://github.com/opensearch-project/opensearch-jvector/pull/242) ### Infrastructure * Update OpenSearch compatibility from version 3.3.0 to 3.3.2 [PR #226](https://github.com/opensearch-project/opensearch-jvector/pull/226) ### Documentation diff --git a/src/main/java/org/opensearch/knn/common/KNNConstants.java b/src/main/java/org/opensearch/knn/common/KNNConstants.java index 6aead6a5..34539980 100644 --- a/src/main/java/org/opensearch/knn/common/KNNConstants.java +++ b/src/main/java/org/opensearch/knn/common/KNNConstants.java @@ -100,6 +100,10 @@ public class KNNConstants { // default public static final Boolean DEFAULT_HIERARCHY_ENABLED = false; + // Parameters that only affect merge + public static final String METHOD_PARAMETER_LEADING_SEGMENT_MERGE_DISABLED = "advanced.leading_segment_merge_disabled"; + public static final boolean DEFAULT_LEADING_SEGMENT_MERGE_DISABLED = false; + // API Constants public static final String CLEAR_CACHE = "clear_cache"; diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN9120Codec/KNN9120PerFieldKnnVectorsFormat.java b/src/main/java/org/opensearch/knn/index/codec/KNN9120Codec/KNN9120PerFieldKnnVectorsFormat.java index 5fc4cf67..8348246f 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN9120Codec/KNN9120PerFieldKnnVectorsFormat.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN9120Codec/KNN9120PerFieldKnnVectorsFormat.java @@ -70,7 +70,8 @@ public KNN9120PerFieldKnnVectorsFormat(final Optional mapperServi knnVectorsFormatParams.getNeighborOverflow(), knnVectorsFormatParams.getNumberOfSubspacesPerVectorSupplier(), knnVectorsFormatParams.getMinBatchSizeForQuantization(), - knnVectorsFormatParams.isHierarchyEnabled() + knnVectorsFormatParams.isHierarchyEnabled(), + knnVectorsFormatParams.isLeadingSegmentMergeDisabled() ); default: throw new IllegalArgumentException("Unsupported java engine: " + knnEngine); diff --git a/src/main/java/org/opensearch/knn/index/codec/jvector/DelayedInitDiversityProvider.java b/src/main/java/org/opensearch/knn/index/codec/jvector/DelayedInitDiversityProvider.java new file mode 100644 index 00000000..1e1a482d --- /dev/null +++ b/src/main/java/org/opensearch/knn/index/codec/jvector/DelayedInitDiversityProvider.java @@ -0,0 +1,54 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.index.codec.jvector; + +import org.opensearch.common.SetOnce; +import org.opensearch.common.SetOnce.AlreadySetException; + +import io.github.jbellis.jvector.graph.NodeArray; +import io.github.jbellis.jvector.graph.diversity.DiversityProvider; +import io.github.jbellis.jvector.util.BitSet; +import lombok.NonNull; + +/** + *

A diversity provider that can be initialized after construction. + * + *

+ * This is useful during leading segment merge when loading the + * neighbors score cache from disk. The score cache is loaded as an + * OnHeapGraphIndex, which requires a diveristy provider during construction. + * However, creating the diversity provider requires knowledge of the ordinals + * contained in the graph which is only available after the graph is loaded. + * Since the diversity provider is not used until the graph is mutated, this + * class can be used to construct the graph without a diversity provider + * available. + */ +class DelayedInitDiversityProvider implements DiversityProvider { + + private SetOnce delegate = new SetOnce<>(); + + /** + * Creates an uninitialized diversity provider. + * Call {@link #initialize} before use. + */ + DelayedInitDiversityProvider() {} + + /** + * Initialize this DiversityProvider with a delegate + * @throws AlreadySetException if already set + */ + void initialize(@NonNull DiversityProvider delegate) throws AlreadySetException { + this.delegate.set(delegate); + } + + @Override + public double retainDiverse(NodeArray neighbors, int maxDegree, int diverseBefore, BitSet selected) { + if (delegate.get() == null) { + throw new IllegalStateException("DelayedInitDiversityProvider was not initialzied (call initialize() before use)"); + } + return delegate.get().retainDiverse(neighbors, maxDegree, diverseBefore, selected); + } +} diff --git a/src/main/java/org/opensearch/knn/index/codec/jvector/GraphNodeIdToDocMap.java b/src/main/java/org/opensearch/knn/index/codec/jvector/GraphNodeIdToDocMap.java index 7fff91e1..6edf88a2 100644 --- a/src/main/java/org/opensearch/knn/index/codec/jvector/GraphNodeIdToDocMap.java +++ b/src/main/java/org/opensearch/knn/index/codec/jvector/GraphNodeIdToDocMap.java @@ -43,10 +43,15 @@ public GraphNodeIdToDocMap(IndexInput in) throws IOException { graphNodeIdsToDocIds = new int[size]; docIdsToGraphNodeIds = new int[maxDocId]; + Arrays.fill(graphNodeIdsToDocIds, -1); + Arrays.fill(docIdsToGraphNodeIds, -1); for (int ord = 0; ord < size; ord++) { final int docId = in.readVInt(); graphNodeIdsToDocIds[ord] = docId; - docIdsToGraphNodeIds[docId] = ord; + if (docId != -1) { + // ignore deleted documents + docIdsToGraphNodeIds[docId] = ord; + } } } @@ -67,19 +72,21 @@ public GraphNodeIdToDocMap(int[] graphNodeIdsToDocIds) { final int maxDocs = maxDocId + 1; // We are going to assume that the number of ordinals is roughly the same as the number of documents in the segment, therefore, // the mapping will not be sparse. - if (maxDocs < graphNodeIdsToDocIds.length) { - throw new IllegalStateException("Max docs " + maxDocs + " is less than the number of ordinals " + graphNodeIdsToDocIds.length); - } - if (maxDocId > graphNodeIdsToDocIds.length) { - log.warn( - "Max doc id {} is greater than the number of ordinals {}, this implies a lot of deleted documents. Or that some documents are missing vectors. Wasting a lot of memory", - maxDocId, + // Note that the merge process may create ephemeral mappings that are sparse. + if (maxDocs < 0.8 * graphNodeIdsToDocIds.length) { + log.info( + "Max docs {} is less than 80% the number of ordinals {}. This is normal if many docs were recently deleted or overwritten. Otherwise if many docs are missing vectors this is a waste of memory.", + maxDocs, graphNodeIdsToDocIds.length ); } this.docIdsToGraphNodeIds = new int[maxDocs]; Arrays.fill(this.docIdsToGraphNodeIds, -1); // -1 means no mapping to ordinal for (int ord = 0; ord < graphNodeIdsToDocIds.length; ord++) { + // -1 means no mapping to docId since document was deleted + if (graphNodeIdsToDocIds[ord] == -1) { + continue; + } this.docIdsToGraphNodeIds[graphNodeIdsToDocIds[ord]] = ord; } } diff --git a/src/main/java/org/opensearch/knn/index/codec/jvector/JVectorDiskANNMethod.java b/src/main/java/org/opensearch/knn/index/codec/jvector/JVectorDiskANNMethod.java index bf4870b7..290f5cb9 100644 --- a/src/main/java/org/opensearch/knn/index/codec/jvector/JVectorDiskANNMethod.java +++ b/src/main/java/org/opensearch/knn/index/codec/jvector/JVectorDiskANNMethod.java @@ -85,6 +85,14 @@ private static MethodComponent initMethodComponent() { (v, context) -> v != null && v > 0 && v <= context.getDimension() ) ) + .addParameter( + METHOD_PARAMETER_LEADING_SEGMENT_MERGE_DISABLED, + new Parameter.BooleanParameter( + METHOD_PARAMETER_LEADING_SEGMENT_MERGE_DISABLED, + DEFAULT_LEADING_SEGMENT_MERGE_DISABLED, + (v, context) -> true + ) + ) .build(); } 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 5d25622d..9916392a 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 @@ -36,6 +36,7 @@ public class JVectorFormat extends KnnVectorsFormat { // Unfortunately, this can't be managed yet by the OpenSearch ThreadPool because it's not supporting {@link ForkJoinPool} types public static final ForkJoinPool SIMD_POOL_MERGE = getPhysicalCoreExecutor(); public static final ForkJoinPool SIMD_POOL_FLUSH = getPhysicalCoreExecutor(); + public static final ForkJoinPool PARALLELISM_POOL = ForkJoinPool.commonPool(); private final int maxConn; private final int beamWidth; @@ -44,6 +45,7 @@ public class JVectorFormat extends KnnVectorsFormat { private final float alpha; private final float neighborOverflow; private final boolean hierarchyEnabled; + private final boolean leadingSegmentMergeDisabled; public JVectorFormat() { this( @@ -54,11 +56,16 @@ public JVectorFormat() { KNNConstants.DEFAULT_ALPHA_VALUE.floatValue(), JVectorFormat::getDefaultNumberOfSubspacesPerVector, KNNConstants.DEFAULT_MINIMUM_BATCH_SIZE_FOR_QUANTIZATION, - KNNConstants.DEFAULT_HIERARCHY_ENABLED + KNNConstants.DEFAULT_HIERARCHY_ENABLED, + KNNConstants.DEFAULT_LEADING_SEGMENT_MERGE_DISABLED ); } public JVectorFormat(int minBatchSizeForQuantization) { + this(minBatchSizeForQuantization, KNNConstants.DEFAULT_LEADING_SEGMENT_MERGE_DISABLED); + } + + public JVectorFormat(int minBatchSizeForQuantization, boolean leadingSegmentMergeDisabled) { this( NAME, DEFAULT_MAX_CONN, @@ -67,7 +74,8 @@ public JVectorFormat(int minBatchSizeForQuantization) { KNNConstants.DEFAULT_ALPHA_VALUE.floatValue(), JVectorFormat::getDefaultNumberOfSubspacesPerVector, minBatchSizeForQuantization, - KNNConstants.DEFAULT_HIERARCHY_ENABLED + KNNConstants.DEFAULT_HIERARCHY_ENABLED, + leadingSegmentMergeDisabled ); } @@ -78,7 +86,8 @@ public JVectorFormat( float alpha, Function numberOfSubspacesPerVectorSupplier, int minBatchSizeForQuantization, - boolean hierarchyEnabled + boolean hierarchyEnabled, + boolean leadingSegmentMergeDisabled ) { this( NAME, @@ -88,7 +97,8 @@ public JVectorFormat( alpha, numberOfSubspacesPerVectorSupplier, minBatchSizeForQuantization, - hierarchyEnabled + hierarchyEnabled, + leadingSegmentMergeDisabled ); } @@ -100,7 +110,8 @@ public JVectorFormat( float alpha, Function numberOfSubspacesPerVectorSupplier, int minBatchSizeForQuantization, - boolean hierarchyEnabled + boolean hierarchyEnabled, + boolean leadingSegmentMergeDisabled ) { super(name); this.maxConn = maxConn; @@ -110,6 +121,7 @@ public JVectorFormat( this.alpha = alpha; this.neighborOverflow = neighborOverflow; this.hierarchyEnabled = hierarchyEnabled; + this.leadingSegmentMergeDisabled = leadingSegmentMergeDisabled; } @Override @@ -122,7 +134,8 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException alpha, numberOfSubspacesPerVectorSupplier, minBatchSizeForQuantization, - hierarchyEnabled + hierarchyEnabled, + leadingSegmentMergeDisabled ); } diff --git a/src/main/java/org/opensearch/knn/index/codec/jvector/JVectorReader.java b/src/main/java/org/opensearch/knn/index/codec/jvector/JVectorReader.java index da7196d9..59c37223 100644 --- a/src/main/java/org/opensearch/knn/index/codec/jvector/JVectorReader.java +++ b/src/main/java/org/opensearch/knn/index/codec/jvector/JVectorReader.java @@ -173,7 +173,8 @@ public void search(String field, float[] target, KnnCollector knnCollector, Acce if (acceptDocs == null) compatibleBits = ord -> true; else { Bits b = acceptDocs.bits(); - compatibleBits = ord -> b == null || b.get(jvectorLuceneDocMap.getLuceneDocId(ord)); + compatibleBits = ord -> b == null + || (jvectorLuceneDocMap.getLuceneDocId(ord) != -1 && b.get(jvectorLuceneDocMap.getLuceneDocId(ord))); } try (var graphSearcher = new GraphSearcher(index)) { 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 492d8a01..bc71b843 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 @@ -5,7 +5,6 @@ package org.opensearch.knn.index.codec.jvector; -import io.github.jbellis.jvector.disk.RandomAccessReader; import io.github.jbellis.jvector.graph.*; import io.github.jbellis.jvector.graph.disk.*; import io.github.jbellis.jvector.graph.disk.feature.Feature; @@ -15,7 +14,6 @@ import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider; import io.github.jbellis.jvector.quantization.PQVectors; import io.github.jbellis.jvector.quantization.ProductQuantization; -import io.github.jbellis.jvector.util.PhysicalCoreExecutor; import io.github.jbellis.jvector.vector.VectorizationProvider; import io.github.jbellis.jvector.vector.types.VectorFloat; import io.github.jbellis.jvector.vector.types.VectorTypeSupport; @@ -33,6 +31,7 @@ import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.store.*; import org.apache.lucene.util.Bits; +import io.github.jbellis.jvector.util.FixedBitSet; import org.apache.lucene.util.IOUtils; import org.apache.lucene.util.RamUsageEstimator; import org.opensearch.knn.plugin.stats.KNNCounter; @@ -47,6 +46,7 @@ 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; @@ -96,6 +96,7 @@ public class JVectorWriter extends KnnVectorsWriter { // as a function of the original dimension private final int minimumBatchSizeForQuantization; // Threshold for the vector count above which we will trigger PQ quantization private final boolean hierarchyEnabled; + private final boolean leadingSegmentMergeDisabled; private boolean finished = false; @@ -107,7 +108,8 @@ public JVectorWriter( float alpha, Function numberOfSubspacesPerVectorSupplier, int minimumBatchSizeForQuantization, - boolean hierarchyEnabled + boolean hierarchyEnabled, + boolean leadingSegmentMergeDisabled ) throws IOException { this.segmentWriteState = segmentWriteState; this.maxConn = maxConn; @@ -117,6 +119,8 @@ public JVectorWriter( this.numberOfSubspacesPerVectorSupplier = numberOfSubspacesPerVectorSupplier; this.minimumBatchSizeForQuantization = minimumBatchSizeForQuantization; this.hierarchyEnabled = hierarchyEnabled; + this.leadingSegmentMergeDisabled = leadingSegmentMergeDisabled; + String metaFileName = IndexFileNames.segmentFileName( segmentWriteState.segmentInfo.name, segmentWriteState.segmentSuffix, @@ -204,16 +208,12 @@ public void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException { log.info("Flushing jVector graph index"); for (FieldWriter field : fields) { final RandomAccessVectorValues randomAccessVectorValues = field.randomAccessVectorValues; - final int[] newToOldOrds = new int[randomAccessVectorValues.size()]; - for (int ord = 0; ord < randomAccessVectorValues.size(); ord++) { - newToOldOrds[ord] = ord; - } final BuildScoreProvider buildScoreProvider; final PQVectors pqVectors; final FieldInfo fieldInfo = field.fieldInfo; if (randomAccessVectorValues.size() >= minimumBatchSizeForQuantization) { log.info("Calculating codebooks and compressed vectors for field {}", fieldInfo.name); - pqVectors = getPQVectors(newToOldOrds, randomAccessVectorValues, fieldInfo); + pqVectors = getPQVectors(randomAccessVectorValues, fieldInfo); buildScoreProvider = BuildScoreProvider.pqBuildScoreProvider(getVectorSimilarityFunction(fieldInfo), pqVectors); } else { log.info( @@ -242,12 +242,11 @@ public void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException { OnHeapGraphIndex graph = getGraph( buildScoreProvider, randomAccessVectorValues, - newToOldOrds, fieldInfo, segmentWriteState.segmentInfo.name, SIMD_POOL_FLUSH ); - writeField(field.fieldInfo, field.randomAccessVectorValues, pqVectors, newToOldOrds, graphNodeIdToDocMap, graph); + writeField(field.fieldInfo, randomAccessVectorValues, pqVectors, graphNodeIdToDocMap, graph); } } @@ -256,7 +255,6 @@ private void writeField( FieldInfo fieldInfo, RandomAccessVectorValues randomAccessVectorValues, PQVectors pqVectors, - int[] newToOldOrds, GraphNodeIdToDocMap graphNodeIdToDocMap, OnHeapGraphIndex graph ) throws IOException { @@ -266,14 +264,7 @@ private void writeField( randomAccessVectorValues.size(), segmentWriteState.segmentInfo.name ); - final var vectorIndexFieldMetadata = writeGraph( - graph, - randomAccessVectorValues, - fieldInfo, - pqVectors, - newToOldOrds, - graphNodeIdToDocMap - ); + final var vectorIndexFieldMetadata = writeGraph(graph, randomAccessVectorValues, fieldInfo, pqVectors, graphNodeIdToDocMap); meta.writeInt(fieldInfo.number); vectorIndexFieldMetadata.toOutput(meta); @@ -306,7 +297,7 @@ private void writeField( /** * Writes the graph and PQ codebooks and compressed vectors to the vector index file * @param graph graph - * @param randomAccessVectorValues random access vector values + * @param randomAccessVectorValues random access vector values with remapped ordinals * @param fieldInfo field info * @return Tuple of start offset and length of the graph * @throws IOException IOException @@ -316,7 +307,6 @@ private VectorIndexFieldMetadata writeGraph( RandomAccessVectorValues randomAccessVectorValues, FieldInfo fieldInfo, PQVectors pqVectors, - int[] newToOldOrds, GraphNodeIdToDocMap graphNodeIdToDocMap ) throws IOException { // field data file, which contains the graph @@ -351,7 +341,7 @@ private VectorIndexFieldMetadata writeGraph( ) { var suppliers = Feature.singleStateFactory( FeatureId.INLINE_VECTORS, - nodeId -> new InlineVectors.State(randomAccessVectorValues.getVector(newToOldOrds[nodeId])) + nodeId -> new InlineVectors.State(randomAccessVectorValues.getVector(nodeId)) ); writer.write(suppliers); long endGraphOffset = jVectorIndexWriter.position(); @@ -381,8 +371,7 @@ private VectorIndexFieldMetadata writeGraph( } } - private PQVectors getPQVectors(int[] newToOldOrds, RandomAccessVectorValues randomAccessVectorValues, FieldInfo fieldInfo) - throws IOException { + private PQVectors getPQVectors(RandomAccessVectorValues randomAccessVectorValues, FieldInfo fieldInfo) throws IOException { final String fieldName = fieldInfo.name; final VectorSimilarityFunction vectorSimilarityFunction = fieldInfo.getVectorSimilarityFunction(); log.info("Computing PQ codebooks for field {} for {} vectors", fieldName, randomAccessVectorValues.size()); @@ -406,7 +395,7 @@ private PQVectors getPQVectors(int[] newToOldOrds, RandomAccessVectorValues rand 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, newToOldOrds.length, newToOldOrds, randomAccessVectorValues, SIMD_POOL_MERGE); + PQVectors pqVectors = PQVectors.encodeAndBuild(pq, randomAccessVectorValues.size(), randomAccessVectorValues, SIMD_POOL_MERGE); log.info( "Encoded and built PQ vectors for field {}, original size: {} bytes, compressed size: {} bytes", fieldName, @@ -589,6 +578,11 @@ class RandomAccessMergedFloatVectorValues implements RandomAccessVectorValues { private static final int READER_ID = 0; private static final int READER_ORD = 1; private static final int LEADING_READER_IDX = 0; + // a number in [0.0, 1.0] that indicates how sparse heap graph ordinals are allowed + // to become before we break out of leading segment merge to make the ordinals compact again. + // Ordinal sparsity has a memory cost (in terms map memory usage) + // during leading segment merge. + private static final double MIN_HEAP_GRAPH_ORDINAL_DENSITY = 0.4; private final VectorTypeSupport VECTOR_TYPE_SUPPORT = VectorizationProvider.getInstance().getVectorTypeSupport(); @@ -605,13 +599,25 @@ class RandomAccessMergedFloatVectorValues implements RandomAccessVectorValues { // Total number of documents including those without values private final int totalDocsCount; + private final int totalLiveVectorsInLeadingReader; + private final int totalLiveVectorsInOtherReaders; + private final int totalLiveVectorsCount; + // Vector dimension private final int dimension; private final FieldInfo fieldInfo; private final MergeState mergeState; + + // The "graphNode" ordinal space includes all vectors in the leading segment and live vectors in the other segments. + // The "compact" ordinal space includes only live vectors from all segments. + // "graphNode" and "compact" ordinal spaces may not have the same relative order in certain cases. + // The RavvOrd space is the keyspace of ravvOrdToReaderMapping[][] declared earlier. private final GraphNodeIdToDocMap graphNodeIdToDocMap; + private final GraphNodeIdToDocMap compactOrdToDocMap; private final int[] graphNodeIdsToRavvOrds; - private boolean deletesFound = false; + private final int[] compactOrdsToRavvOrds; + + private final FixedBitSet[] liveGraphNodesPerReader; /** * Creates a random access view over merged float vector values. @@ -636,9 +642,19 @@ public RandomAccessMergedFloatVectorValues(FieldInfo fieldInfo, MergeState merge List allReaders = new ArrayList<>(); final MergeState.DocMap[] docMaps = mergeState.docMaps.clone(); final Bits[] liveDocs = mergeState.liveDocs.clone(); + this.liveGraphNodesPerReader = new FixedBitSet[liveDocs.length]; + // initialize liveGraphNodesPerReader + for (int i = 0; i < liveGraphNodesPerReader.length; i++) { + final int length; + if (liveDocs[i] == null) { + length = mergeState.maxDocs[i]; + } else { + length = liveDocs[i].length(); + } + liveGraphNodesPerReader[i] = new FixedBitSet(length); + liveGraphNodesPerReader[i].set(0, length); + } final int[] baseOrds = new int[mergeState.knnVectorsReaders.length]; - final int[] deletedOrds = new int[mergeState.knnVectorsReaders.length]; // counts the number of deleted documents in each reader - // that previously had a vector // Find the leading reader, count the total number of live vectors, and the base ordinals for each reader for (int i = 0; i < mergeState.knnVectorsReaders.length; i++) { @@ -657,8 +673,8 @@ public RandomAccessMergedFloatVectorValues(FieldInfo fieldInfo, MergeState merge if (liveDocs[i] == null || liveDocs[i].get(it.docID())) { liveVectorCountInReader++; } else { - deletedOrds[i]++; - deletesFound = true; + // This vector is deleted so we need to mark it as deleted in the liveGraphNodesPerReader + liveGraphNodesPerReader[i].clear(it.index()); } } if (liveVectorCountInReader >= vectorsCountInLeadingReader) { @@ -698,6 +714,10 @@ public RandomAccessMergedFloatVectorValues(FieldInfo fieldInfo, MergeState merge final int tempBaseOrd = baseOrds[LEADING_READER_IDX]; baseOrds[LEADING_READER_IDX] = baseOrds[tempLeadingReaderIdx]; baseOrds[tempLeadingReaderIdx] = tempBaseOrd; + // swap liveGraphNodesPerReader + final FixedBitSet tempLiveGraphNodes = liveGraphNodesPerReader[LEADING_READER_IDX]; + liveGraphNodesPerReader[LEADING_READER_IDX] = liveGraphNodesPerReader[tempLeadingReaderIdx]; + liveGraphNodesPerReader[tempLeadingReaderIdx] = tempLiveGraphNodes; } this.perReaderFloatVectorValues = new JVectorFloatVectorValues[readers.length]; @@ -706,111 +726,104 @@ public RandomAccessMergedFloatVectorValues(FieldInfo fieldInfo, MergeState merge // Build mapping from global ordinal to [readerIndex, readerOrd] this.ravvOrdToReaderMapping = new int[totalDocsCount][2]; - int documentsIterated = 0; - // Will be used to build the new graphNodeIdToDocMap with the new graph node id to docId mapping. // This mapping should not be used to access the vectors at any time during construction, but only after the merge is complete // and the new segment is created and used by searchers. - final int[] graphNodeIdToDocIds = new int[totalLiveVectorsCount]; - this.graphNodeIdsToRavvOrds = new int[totalLiveVectorsCount]; + // note: totalVectorsCount is an unnecessarily loose bound for the "graphNodeId" space, see init of graphNodeIdsToRavvOrds below + final int[] graphNodeIdToDocIds = new int[totalVectorsCount]; + final int[] compactOrdToDocIds = new int[totalLiveVectorsCount]; + Arrays.fill(graphNodeIdToDocIds, -1); + Arrays.fill(compactOrdToDocIds, -1); + // The graph node id to ravv ordinal mapping, for the leading reader this mapping is always identity, for the other readers we + // only need to map the live vectors + // Therefore the size of this array is the total number of vectors in the leading reader + the number of live vectors in all + // other readers + totalLiveVectorsInLeadingReader = liveGraphNodesPerReader[LEADING_READER_IDX].cardinality(); + totalLiveVectorsInOtherReaders = totalLiveVectorsCount - totalLiveVectorsInLeadingReader; + this.totalLiveVectorsCount = totalLiveVectorsCount; + this.graphNodeIdsToRavvOrds = new int[totalLiveVectorsInOtherReaders + readers[LEADING_READER_IDX].getFloatVectorValues( + fieldName + ).size()]; + this.compactOrdsToRavvOrds = new int[totalLiveVectorsCount]; + // initialize the graphNodeIdsToRavvOrds with -1 to indicate that there is no mapping + Arrays.fill(graphNodeIdsToRavvOrds, -1); + Arrays.fill(compactOrdsToRavvOrds, -1); + int documentsIterated = 0; int graphNodeId = 0; - if (deletesFound) { - // If there are deletes, we need to build a new graph from scratch and compact the graph node ids - // TODO: remove this logic once we support incremental graph building with deletes see - // https://github.com/opensearch-project/opensearch-jvector/issues/171 - for (int readerIdx = 0; readerIdx < readers.length; readerIdx++) { - final JVectorFloatVectorValues values = (JVectorFloatVectorValues) readers[readerIdx].getFloatVectorValues(fieldName); - perReaderFloatVectorValues[readerIdx] = values; - // For each vector in this reader - KnnVectorValues.DocIndexIterator it = values.iterator(); - - for (int docId = it.nextDoc(); docId != DocIdSetIterator.NO_MORE_DOCS; docId = it.nextDoc()) { - if (docMaps[readerIdx].get(docId) == -1) { - log.warn( - "Document {} in reader {} is not mapped to a global ordinal from the merge docMaps. Will skip this document for now", - docId, - readerIdx - ); - } else { - // Mapping from ravv ordinals to [readerIndex, readerOrd] - // Map graph node id to ravv ordinal - // Map graph node id to doc id - final int newGlobalDocId = docMaps[readerIdx].get(docId); - final int ravvLocalOrd = it.index(); - final int ravvGlobalOrd = ravvLocalOrd + baseOrds[readerIdx]; - graphNodeIdToDocIds[graphNodeId] = newGlobalDocId; - graphNodeIdsToRavvOrds[graphNodeId] = ravvGlobalOrd; - graphNodeId++; - ravvOrdToReaderMapping[ravvGlobalOrd][READER_ID] = readerIdx; // Reader index - ravvOrdToReaderMapping[ravvGlobalOrd][READER_ORD] = ravvLocalOrd; // Ordinal in reader - } + int compactNodeId = 0; + // We can reuse the existing graph and simply remap the ravv ordinals to the new global doc ids + // for the leading reader we must preserve the original node Ids and map them to the corresponding ravv vectors originally + // used to build the graph + // This is necessary because we are later going to expand that graph with new vectors from the other readers. + // The leading reader is ALWAYS the first one in the readers array + // Deletes will be applied after the graph is built on heap using the cleanup method + final JVectorFloatVectorValues leadingReaderValues = (JVectorFloatVectorValues) readers[LEADING_READER_IDX] + .getFloatVectorValues(fieldName); + perReaderFloatVectorValues[LEADING_READER_IDX] = leadingReaderValues; + var leadingReaderIt = leadingReaderValues.iterator(); + for (int docId = leadingReaderIt.nextDoc(); docId != DocIdSetIterator.NO_MORE_DOCS; docId = leadingReaderIt.nextDoc()) { + final int newGlobalDocId = docMaps[LEADING_READER_IDX].get(docId); + // we increment anyway even if deleted because we are going to apply cleanup later to remove deleted nodes from the leading + // graph that will be used incremented + graphNodeId++; + if (newGlobalDocId == -1) { + log.debug( + "Document {} in reader {} is not mapped to a global ordinal from the merge docMaps. This means it's deleted, Will skip this document for now", + docId, + LEADING_READER_IDX + ); + } - documentsIterated++; - } + final int ravvLocalOrd = leadingReaderIt.index(); + final int ravvGlobalOrd = ravvLocalOrd + baseOrds[LEADING_READER_IDX]; + graphNodeIdToDocIds[ravvLocalOrd] = newGlobalDocId; + graphNodeIdsToRavvOrds[ravvLocalOrd] = ravvGlobalOrd; + if (newGlobalDocId != -1) { + // the "compact" ordinal space doesn't include any deletes + compactOrdsToRavvOrds[compactNodeId] = ravvGlobalOrd; + compactOrdToDocIds[compactNodeId] = newGlobalDocId; + compactNodeId += 1; } - } else { - // If there are no deletes, we can reuse the existing graph and simply remap the ravv ordinals to the new global doc ids - // for the leading reader we must preserve the original node Ids and map them to the corresponding ravv vectors originally - // used to build the graph - // This is necessary because we are later going to expand that graph with new vectors from the other readers. - // The leading reader is ALWAYS the first one in the readers array - final JVectorFloatVectorValues leadingReaderValues = (JVectorFloatVectorValues) readers[LEADING_READER_IDX] - .getFloatVectorValues(fieldName); - perReaderFloatVectorValues[LEADING_READER_IDX] = leadingReaderValues; - var leadingReaderIt = leadingReaderValues.iterator(); - for (int docId = leadingReaderIt.nextDoc(); docId != DocIdSetIterator.NO_MORE_DOCS; docId = leadingReaderIt.nextDoc()) { - final int newGlobalDocId = docMaps[LEADING_READER_IDX].get(docId); - if (newGlobalDocId == -1) { - log.warn( + ravvOrdToReaderMapping[ravvGlobalOrd][READER_ID] = LEADING_READER_IDX; // Reader index + ravvOrdToReaderMapping[ravvGlobalOrd][READER_ORD] = ravvLocalOrd; // Ordinal in reader + + documentsIterated++; + } + + // For the remaining readers we map the graph node id to the ravv ordinal in the order they appear + for (int readerIdx = 1; readerIdx < readers.length; readerIdx++) { + final JVectorFloatVectorValues values = (JVectorFloatVectorValues) readers[readerIdx].getFloatVectorValues(fieldName); + perReaderFloatVectorValues[readerIdx] = values; + // For each vector in this reader + KnnVectorValues.DocIndexIterator it = values.iterator(); + + for (int docId = it.nextDoc(); docId != DocIdSetIterator.NO_MORE_DOCS; docId = it.nextDoc()) { + if (docMaps[readerIdx].get(docId) == -1) { + log.debug( "Document {} in reader {} is not mapped to a global ordinal from the merge docMaps. Will skip this document for now", docId, - LEADING_READER_IDX + readerIdx ); } else { - final int ravvLocalOrd = leadingReaderIt.index(); - final int ravvGlobalOrd = ravvLocalOrd + baseOrds[LEADING_READER_IDX]; - graphNodeIdToDocIds[ravvLocalOrd] = newGlobalDocId; - graphNodeIdsToRavvOrds[ravvLocalOrd] = ravvGlobalOrd; + // Mapping from ravv ordinals to [readerIndex, readerOrd] + // Map graph node id to ravv ordinal + // Map graph node id to doc id + final int newGlobalDocId = docMaps[readerIdx].get(docId); + final int ravvLocalOrd = it.index(); + final int ravvGlobalOrd = ravvLocalOrd + baseOrds[readerIdx]; + graphNodeIdToDocIds[graphNodeId] = newGlobalDocId; + graphNodeIdsToRavvOrds[graphNodeId] = ravvGlobalOrd; + compactOrdsToRavvOrds[compactNodeId] = ravvGlobalOrd; + compactOrdToDocIds[compactNodeId] = newGlobalDocId; + compactNodeId++; graphNodeId++; - ravvOrdToReaderMapping[ravvGlobalOrd][READER_ID] = LEADING_READER_IDX; // Reader index + ravvOrdToReaderMapping[ravvGlobalOrd][READER_ID] = readerIdx; // Reader index ravvOrdToReaderMapping[ravvGlobalOrd][READER_ORD] = ravvLocalOrd; // Ordinal in reader } documentsIterated++; } - - // For the remaining readers we map the graph node id to the ravv ordinal in the order they appear - for (int readerIdx = 1; readerIdx < readers.length; readerIdx++) { - final JVectorFloatVectorValues values = (JVectorFloatVectorValues) readers[readerIdx].getFloatVectorValues(fieldName); - perReaderFloatVectorValues[readerIdx] = values; - // For each vector in this reader - KnnVectorValues.DocIndexIterator it = values.iterator(); - - for (int docId = it.nextDoc(); docId != DocIdSetIterator.NO_MORE_DOCS; docId = it.nextDoc()) { - if (docMaps[readerIdx].get(docId) == -1) { - log.warn( - "Document {} in reader {} is not mapped to a global ordinal from the merge docMaps. Will skip this document for now", - docId, - readerIdx - ); - } else { - // Mapping from ravv ordinals to [readerIndex, readerOrd] - // Map graph node id to ravv ordinal - // Map graph node id to doc id - final int newGlobalDocId = docMaps[readerIdx].get(docId); - final int ravvLocalOrd = it.index(); - final int ravvGlobalOrd = ravvLocalOrd + baseOrds[readerIdx]; - graphNodeIdToDocIds[graphNodeId] = newGlobalDocId; - graphNodeIdsToRavvOrds[graphNodeId] = ravvGlobalOrd; - graphNodeId++; - ravvOrdToReaderMapping[ravvGlobalOrd][READER_ID] = readerIdx; // Reader index - ravvOrdToReaderMapping[ravvGlobalOrd][READER_ORD] = ravvLocalOrd; // Ordinal in reader - } - - documentsIterated++; - } - } } if (documentsIterated < totalVectorsCount) { @@ -825,6 +838,7 @@ public RandomAccessMergedFloatVectorValues(FieldInfo fieldInfo, MergeState merge } this.graphNodeIdToDocMap = new GraphNodeIdToDocMap(graphNodeIdToDocIds); + this.compactOrdToDocMap = new GraphNodeIdToDocMap(compactOrdToDocIds); log.debug("Created RandomAccessMergedFloatVectorValues with {} total vectors from {} readers", size, readers.length); } @@ -855,12 +869,11 @@ public void merge() throws IOException { // Get PQ compressor for leading reader final int totalVectorsCount = size; final String fieldName = fieldInfo.name; - final PQVectors pqVectors; - final OnHeapGraphIndex graph; + final PQVectors compactPqVectors; // Get the leading reader PerFieldKnnVectorsFormat.FieldsReader fieldsReader = (PerFieldKnnVectorsFormat.FieldsReader) readers[LEADING_READER_IDX]; JVectorReader leadingReader = (JVectorReader) fieldsReader.getFieldReader(fieldName); - final BuildScoreProvider buildScoreProvider; + final RemappedRandomAccessVectorValues compactRavv = new RemappedRandomAccessVectorValues(this, compactOrdsToRavvOrds); // Check if the leading reader has pre-existing PQ codebooks and if so, refine them with the remaining vectors if (leadingReader.getProductQuantizationForField(fieldInfo.name).isEmpty()) { // No pre-existing codebooks, check if we have enough vectors to trigger quantization @@ -869,14 +882,14 @@ public void merge() throws IOException { fieldName, mergeState.segmentInfo.name ); - if (this.size() >= minimumBatchSizeForQuantization) { + if (totalLiveVectorsCount >= minimumBatchSizeForQuantization) { log.info( "Calculating new codebooks and compressed vectors for field: {}, with totalVectorCount: {}, above minimumBatchSizeForQuantization: {}", fieldName, totalVectorsCount, minimumBatchSizeForQuantization ); - pqVectors = getPQVectors(graphNodeIdsToRavvOrds, this, fieldInfo); + compactPqVectors = getPQVectors(compactRavv, fieldInfo); } else { log.info( "Not enough vectors found for field: {}, totalVectorCount: {}, is below minimumBatchSizeForQuantization: {}", @@ -884,7 +897,7 @@ public void merge() throws IOException { totalVectorsCount, minimumBatchSizeForQuantization ); - pqVectors = null; + compactPqVectors = null; } } else { log.info( @@ -906,69 +919,222 @@ public void merge() throws IOException { final long trainingTime = end - start; log.info("Refined PQ codebooks for field {}, in {} millis", fieldName, trainingTime); KNNCounter.KNN_QUANTIZATION_TRAINING_TIME.add(trainingTime); - pqVectors = PQVectors.encodeAndBuild( + compactPqVectors = PQVectors.encodeAndBuild( leadingCompressor, - graphNodeIdsToRavvOrds.length, - graphNodeIdsToRavvOrds, + // graphNodeIdsToRavvOrds.length, + // graphNodeIdsToRavvOrds, + compactOrdsToRavvOrds.length, + compactOrdsToRavvOrds, this, SIMD_POOL_MERGE ); } - if (pqVectors == null) { - buildScoreProvider = BuildScoreProvider.randomAccessScoreProvider( - this, - graphNodeIdsToRavvOrds, - getVectorSimilarityFunction(fieldInfo) - ); - // graph = getGraph(buildScoreProvider, this, newToOldOrds, fieldInfo, segmentWriteState.segmentInfo.name); - if (!deletesFound) { - final String segmentName = segmentWriteState.segmentInfo.name; + if (compactPqVectors == null) { + final String segmentName = segmentWriteState.segmentInfo.name; + log.info("No PQ codebooks found, will merge with full-precision vectors: field {} in segment {}", fieldName, segmentName); + + boolean ok = tryLeadingSegmentMerge(); + if (!ok) { + // leading segment merge was skipped log.info( - "No deletes found, and no PQ codebooks found, expanding previous graph with additional vectors for field {} in segment {}", - fieldName, - segmentName - ); - final RandomAccessReader leadingOnHeapGraphReader = leadingReader.getNeighborsScoreCacheForField(fieldName); - final int numBaseVectors = leadingReader.getFloatVectorValues(fieldName).size(); - graph = (OnHeapGraphIndex) buildAndMergeNewNodes( - leadingOnHeapGraphReader, - this, - buildScoreProvider, - numBaseVectors, - graphNodeIdsToRavvOrds, - beamWidth, - degreeOverflow, - alpha, - hierarchyEnabled - ); - } else { - log.info("Deletes found, and no PQ codebooks found, building new graph from scratch"); - graph = getGraph( - buildScoreProvider, - this, - graphNodeIdsToRavvOrds, - fieldInfo, - segmentWriteState.segmentInfo.name, - SIMD_POOL_MERGE + "Merging segments by building graph from scratch (skipping leading segment merge) for segment {}, on field {}", + segmentName, + fieldName ); + var bsp = BuildScoreProvider.randomAccessScoreProvider(compactRavv, getVectorSimilarityFunction(fieldInfo)); + var graph = getGraph(bsp, compactRavv, fieldInfo, segmentWriteState.segmentInfo.name, SIMD_POOL_MERGE); + writeField(fieldInfo, compactRavv, null, compactOrdToDocMap, graph); } } else { log.info("PQ codebooks found, building graph from scratch with PQ vectors"); - buildScoreProvider = BuildScoreProvider.pqBuildScoreProvider(getVectorSimilarityFunction(fieldInfo), pqVectors); + // We're building from scratch, so we can use the "compact" ordinal space directly + 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); - graph = getGraph( - buildScoreProvider, - this, - graphNodeIdsToRavvOrds, - fieldInfo, - segmentWriteState.segmentInfo.name, - SIMD_POOL_MERGE - ); + var graph = getGraph(buildScoreProvider, compactRavv, fieldInfo, segmentWriteState.segmentInfo.name, SIMD_POOL_MERGE); + writeField(fieldInfo, compactRavv, compactPqVectors, compactOrdToDocMap, graph); + } + } + + /** + *

Perform leading segment merge. + * + *

+ * In some cases leading segment merge should be skipped, and this method will return false. + * This is the case when: + * - Leading segment merge is disabled through configuration + * - There is a risk of integer overflow due to sparsity of the OnHeapGraph + * - The OnHeapGraph is too sparse relative to it's size + * + * @return a boolean value indicating if leading segment merge was performed + */ + private boolean tryLeadingSegmentMerge() throws IOException { + if (leadingSegmentMergeDisabled) { + log.info("Leading segment merge is disabled, skipping"); + return false; } - writeField(fieldInfo, this, pqVectors, graphNodeIdsToRavvOrds, graphNodeIdToDocMap, graph); + var leadingFieldsReader = (PerFieldKnnVectorsFormat.FieldsReader) readers[LEADING_READER_IDX]; + var leadingReader = (JVectorReader) leadingFieldsReader.getFieldReader(fieldInfo.name); + var graphReader = leadingReader.getNeighborsScoreCacheForField(fieldInfo.name); + + // A diversity provider is a required argument to OnHeapGraphIndex::load, + // however creating a VamanaDiversityProvider requires a buildScore provider + // which in turn requires knowledge of the ordinal -> vector mapping. + // Since the heap graph can contain ordinal "holes", this information + // is not available until AFTER the graph is loaded. + // Using a DelayedInitDiversityProvider avoids this chicken-and-egg problem. + // This is okay since the diversity provider is only used during mutations. + var diversityProvider = new DelayedInitDiversityProvider(); + + var numBaseVectors = leadingReader.getFloatVectorValues(fieldInfo.name).size(); + + try (OnHeapGraphIndex leadingGraph = OnHeapGraphIndex.load(graphReader, this.dimension(), degreeOverflow, diversityProvider)) { + // Since we're performing leading segment merge, we need to load the OnHeapGraphIndex + // corresponding to the leading segment, then mutate it. + // Since ordinals of OnHeapGraphIndex cannot be compacted, holes in the ordinals + // will build up over time. + + int leadingGraphIdUpperBound = leadingGraph.getIdUpperBound(); + long heapOrdUpperBoundLong = leadingGraphIdUpperBound + (long) totalLiveVectorsInOtherReaders; + + // even though Lucene should ensure that the sum of maxDoc across all segments + // is under IndexWriter.MAX_DOCS, our leading segment heap graph has holes + // that Lucene won't know about, which may push heapOrdUpperBound out of range. + // IndexWriter.MAX_DOCS = Integer.MAX_VALUE - 128 + if (heapOrdUpperBoundLong > IndexWriter.MAX_DOCS) { + log.warn( + "New ordinal upper bound for neighbor score cache is too large. " + + "This may indicate many deletes, or that you're approaching the maximum segment size. " + + "Will skip leading segment merge." + ); + return false; // indicate skipped + } + + var heapGraphOrdinalDensity = totalLiveVectorsCount / (double) heapOrdUpperBoundLong; + if (heapGraphOrdinalDensity < MIN_HEAP_GRAPH_ORDINAL_DENSITY) { + log.warn( + "Heap ordinals will be insufficiently dense ({} < {}). " + + "Will skip leading segment merge. (totalLiveVectors={}, heapOrdUpperBound={})", + heapGraphOrdinalDensity, + MIN_HEAP_GRAPH_ORDINAL_DENSITY, + totalLiveVectorsCount, + heapOrdUpperBoundLong + ); + return false; + } + + log.info("Starting leading segment merge for segment {} on field {}", segmentWriteState.segmentInfo.name, fieldInfo.name); + + // we already checked that heapOrdUpperBoundLong <= IndexWriter.MAX_DOCS < Integer.MAX_VALUE + int heapOrdUpperBound = Math.toIntExact(heapOrdUpperBoundLong); + + // While creating score providers and updating the graph, we need to supply a RAVV that + // takes account the accumulated holes in the OnHeapGraph, so we need some mappings + // to and from the ordinal space of the OnHeapGraphIndex (the heap ordinals) + + // The "mid" ordinal space is the same as the "graphNodeId" ordinal space + // which includes all vectors from the leading segment and live vectors from other segments. + // We need this mapping to delete vectors later. + var midToHeapOrds = new int[graphNodeIdsToRavvOrds.length]; + var heapToGlobalRavvOrds = new int[heapOrdUpperBound]; + Arrays.fill(midToHeapOrds, -1); + Arrays.fill(heapToGlobalRavvOrds, -1); + + // The "final" ord space is what happens to the ordinals on the heap graph + // on being written as an OnDiskGraphIndex. + // JVector automatically compacts the ordinals while preserving the order. + // Note that this may NOT be the same as the "compact" ordinal space calculated earler, + // (although it is also compact) + var finalOrdToDocId = new int[totalLiveVectorsCount]; + + int midOrd = 0; + int finalOrd = 0; + for (int heapOrd = 0; heapOrd < leadingGraphIdUpperBound; heapOrd++) { + if (!leadingGraph.containsNode(heapOrd)) { + continue; + } + midToHeapOrds[midOrd] = heapOrd; + heapToGlobalRavvOrds[heapOrd] = graphNodeIdsToRavvOrds[midOrd]; + // the old ordinal space of the leading reader is not exactly the "mid" ordinal space + // but by definition they match for the leading reader, so `.get(midOrd)` is valid + if (liveGraphNodesPerReader[LEADING_READER_IDX].get(midOrd)) { + finalOrdToDocId[finalOrd] = graphNodeIdToDocMap.getLuceneDocId(midOrd); + finalOrd++; + } + midOrd++; + } + + for (int heapOrd = leadingGraphIdUpperBound; heapOrd < heapOrdUpperBound; heapOrd++) { + midToHeapOrds[midOrd] = heapOrd; + heapToGlobalRavvOrds[heapOrd] = graphNodeIdsToRavvOrds[midOrd]; + finalOrdToDocId[finalOrd] = graphNodeIdToDocMap.getLuceneDocId(midOrd); + finalOrd++; + midOrd++; + } + + if (midOrd != midToHeapOrds.length || finalOrd != finalOrdToDocId.length) { + log.error( + "Got midOrd_limit={} (wanted {}), finalOrd_limit={} (wanted {})", + midOrd, + midToHeapOrds.length, + finalOrd, + finalOrdToDocId.length + ); + throw new IllegalStateException("failed to fill one of the maps, this is a bug"); + } + + var heapRavv = new RemappedRandomAccessVectorValues(this, heapToGlobalRavvOrds); + + var leadingBsp = BuildScoreProvider.randomAccessScoreProvider(heapRavv, getVectorSimilarityFunction(fieldInfo)); + + // we left this uninitialized earlier, but we're ready to set it up now + // just in time to mutate the graph + diversityProvider.initialize(new VamanaDiversityProvider(leadingBsp, alpha)); + + OnHeapGraphIndex graph; + try ( + GraphIndexBuilder builder = new GraphIndexBuilder( + leadingBsp, + heapRavv.dimension(), + leadingGraph, + beamWidth, + degreeOverflow, + alpha, + true, + SIMD_POOL_MERGE, + PARALLELISM_POOL + ) + ) { + 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 -> { + builder.addGraphNode(ord, vv.get().getVector(ord)); + }) + ).join(); + + // mark deleted nodes + for (int i = 0; i < numBaseVectors; i++) { + if (!liveGraphNodesPerReader[LEADING_READER_IDX].get(i)) { + // we need to convert from the "mid" to the "heap" ordinal space to avoid errors + builder.markNodeDeleted(midToHeapOrds[i]); + } + } + + builder.cleanup(); + + graph = (OnHeapGraphIndex) builder.getGraph(); + } + + // Note that the ordinals for the OnDiskGraphIndex will automatically be compacted + // But the OnHeapGraphIndex will not + var finalOrdToDocMap = new GraphNodeIdToDocMap(finalOrdToDocId); + writeField(fieldInfo, heapRavv, null, finalOrdToDocMap, graph); + return true; + } } @Override @@ -983,7 +1149,7 @@ public int dimension() { @Override public VectorFloat getVector(int ord) { - if (ord < 0 || ord >= totalDocsCount) { + if (ord < 0 || ord >= ravvOrdToReaderMapping.length) { throw new IllegalArgumentException("Ordinal out of bounds: " + ord); } @@ -1014,7 +1180,6 @@ public RandomAccessVectorValues copy() { public OnHeapGraphIndex getGraph( BuildScoreProvider buildScoreProvider, RandomAccessVectorValues randomAccessVectorValues, - int[] newToOldOrds, FieldInfo fieldInfo, String segmentName, ForkJoinPool SIMD_POOL @@ -1041,8 +1206,8 @@ public OnHeapGraphIndex getGraph( log.info("Building graph from merged float vector"); // parallel graph construction from the merge documents Ids - SIMD_POOL.submit(() -> IntStream.range(0, newToOldOrds.length).parallel().forEach(ord -> { - graphIndexBuilder.addGraphNode(ord, vv.get().getVector(newToOldOrds[ord])); + SIMD_POOL.submit(() -> IntStream.range(0, randomAccessVectorValues.size()).parallel().forEach(ord -> { + graphIndexBuilder.addGraphNode(ord, vv.get().getVector(ord)); })).join(); graphIndexBuilder.cleanup(); graphIndex = (OnHeapGraphIndex) graphIndexBuilder.getGraph(); @@ -1097,43 +1262,4 @@ public RandomAccessVectorValues copy() { } } - static ImmutableGraphIndex buildAndMergeNewNodes( - RandomAccessReader in, - RandomAccessVectorValues newVectors, - BuildScoreProvider buildScoreProvider, - int startingNodeOffset, - int[] graphToRavvOrdMap, - int beamWidth, - float overflowRatio, - float alpha, - boolean addHierarchy - ) throws IOException { - - var diversityProvider = new VamanaDiversityProvider(buildScoreProvider, alpha); - - try (OnHeapGraphIndex graph = OnHeapGraphIndex.load(in, newVectors.dimension(), overflowRatio, diversityProvider);) { - - GraphIndexBuilder builder = new GraphIndexBuilder( - buildScoreProvider, - newVectors.dimension(), - graph, - beamWidth, - overflowRatio, - alpha, - true, - PhysicalCoreExecutor.pool(), - ForkJoinPool.commonPool() - ); - - var vv = newVectors.threadLocalSupplier(); - - // parallel graph construction from the merge documents Ids - PhysicalCoreExecutor.pool().submit(() -> IntStream.range(startingNodeOffset, newVectors.size()).parallel().forEach(ord -> { - builder.addGraphNode(ord, vv.get().getVector(graphToRavvOrdMap[ord])); - })).join(); - - builder.cleanup(); - return builder.getGraph(); - } - } } diff --git a/src/main/java/org/opensearch/knn/index/codec/params/KNNVectorsFormatParams.java b/src/main/java/org/opensearch/knn/index/codec/params/KNNVectorsFormatParams.java index c5e8d21a..4076cde1 100644 --- a/src/main/java/org/opensearch/knn/index/codec/params/KNNVectorsFormatParams.java +++ b/src/main/java/org/opensearch/knn/index/codec/params/KNNVectorsFormatParams.java @@ -26,6 +26,7 @@ public class KNNVectorsFormatParams { private boolean hierarchyEnabled; private Function numberOfSubspacesPerVectorSupplier; private final SpaceType spaceType; + private boolean leadingSegmentMergeDisabled; public KNNVectorsFormatParams(final Map params, int defaultMaxConnections, int defaultBeamWidth) { this( @@ -58,6 +59,7 @@ public KNNVectorsFormatParams( initHierarchyEnabled(params, defaultHierarchyEnabled); initNumberOfSubspacesPerVectorSupplier(params); this.spaceType = spaceType; + initLeadingSegmentMergeDisabled(params, KNNConstants.DEFAULT_LEADING_SEGMENT_MERGE_DISABLED); } public boolean validate(final Map params) { @@ -120,4 +122,12 @@ private void initNumberOfSubspacesPerVectorSupplier(final Map pa } this.numberOfSubspacesPerVectorSupplier = JVectorFormat::getDefaultNumberOfSubspacesPerVector; } + + private void initLeadingSegmentMergeDisabled(final Map params, boolean defaultLsmDisabled) { + if (params != null && params.containsKey(KNNConstants.METHOD_PARAMETER_LEADING_SEGMENT_MERGE_DISABLED)) { + this.hierarchyEnabled = (boolean) params.get(KNNConstants.METHOD_PARAMETER_LEADING_SEGMENT_MERGE_DISABLED); + return; + } + this.leadingSegmentMergeDisabled = defaultLsmDisabled; + } } diff --git a/src/test/java/org/opensearch/knn/index/ThreadLeakFiltersForTests.java b/src/test/java/org/opensearch/knn/index/ThreadLeakFiltersForTests.java index 6598a00f..f7df4bbd 100644 --- a/src/test/java/org/opensearch/knn/index/ThreadLeakFiltersForTests.java +++ b/src/test/java/org/opensearch/knn/index/ThreadLeakFiltersForTests.java @@ -21,6 +21,7 @@ public boolean reject(Thread thread) { || threadGroup.getName().startsWith("TGRP-KNNJVectorTests") || threadGroup.getName().startsWith("TGRP-JVectorConcurrentQueryTests") || threadGroup.getName().startsWith("TGRP-JVectorMergeWithDeletedDocsTests") + || threadGroup.getName().startsWith("TGRP-JVectorWriterMergeTests") || threadGroup.getName().startsWith("TGRP-MemoryUsageAnalysisTests")); } } diff --git a/src/test/java/org/opensearch/knn/index/codec/jvector/GraphNodeIdToDocMapTests.java b/src/test/java/org/opensearch/knn/index/codec/jvector/GraphNodeIdToDocMapTests.java index ac9947c3..8c3c5959 100644 --- a/src/test/java/org/opensearch/knn/index/codec/jvector/GraphNodeIdToDocMapTests.java +++ b/src/test/java/org/opensearch/knn/index/codec/jvector/GraphNodeIdToDocMapTests.java @@ -48,10 +48,25 @@ public void testConstructorWithSequentialDocIds() { } } - @Test(expected = IllegalStateException.class) - public void testConstructorThrowsWhenMaxDocsLessThanOrdinals() { - int[] invalidMapping = { -1 }; // This would cause issues in Arrays.stream().max() - new GraphNodeIdToDocMap(invalidMapping); + @Test + public void testConstructorWithDeletedDocuments() { + int[] ordinalsToDocIds = { -1, 2, -1, 6 }; // Mix of deleted (-1) and valid docs + GraphNodeIdToDocMap docMap = new GraphNodeIdToDocMap(ordinalsToDocIds); + + // Test ordinal to doc ID mapping (including deleted docs) + assertEquals(-1, docMap.getLuceneDocId(0)); // deleted + assertEquals(2, docMap.getLuceneDocId(1)); + assertEquals(-1, docMap.getLuceneDocId(2)); // deleted + assertEquals(6, docMap.getLuceneDocId(3)); + + // Test doc ID to ordinal mapping (only valid docs should map back) + assertEquals(1, docMap.getJVectorNodeId(2)); + assertEquals(3, docMap.getJVectorNodeId(6)); + + // Unmapped doc IDs return -1 + assertEquals(-1, docMap.getJVectorNodeId(0)); + assertEquals(-1, docMap.getJVectorNodeId(1)); + assertEquals(-1, docMap.getJVectorNodeId(3)); } @Test diff --git a/src/test/java/org/opensearch/knn/index/codec/jvector/JVectorWriterMergeTests.java b/src/test/java/org/opensearch/knn/index/codec/jvector/JVectorWriterMergeTests.java new file mode 100644 index 00000000..7c10e25a --- /dev/null +++ b/src/test/java/org/opensearch/knn/index/codec/jvector/JVectorWriterMergeTests.java @@ -0,0 +1,442 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.index.codec.jvector; + +import static org.opensearch.knn.index.engine.CommonTestUtils.getCodec; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.Arrays; +import java.util.List; +import java.util.PriorityQueue; +import java.util.Set; +import java.util.stream.Collectors; + +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field.Store; +import org.apache.lucene.document.IntField; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.store.FSDirectory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.util.BitSet; +import org.apache.lucene.util.FixedBitSet; +import org.junit.Assert; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.opensearch.knn.TestUtils; +import org.opensearch.knn.common.KNNConstants; +import org.opensearch.knn.index.ThreadLeakFiltersForTests; + +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakFilters; + +import io.github.jbellis.jvector.util.DocIdSetIterator; +import lombok.AllArgsConstructor; +import lombok.Builder; +import lombok.Builder.Default; +import lombok.Singular; + +@ThreadLeakFilters(defaultFilters = true, filters = { ThreadLeakFiltersForTests.class }) +@LuceneTestCase.SuppressSysoutChecks(bugUrl = "") +public class JVectorWriterMergeTests extends LuceneTestCase { + + static final int RANDOM_SEED = 42; + static final String ID_FIELD = "doc_id"; + static final String VECTOR_FIELD = "vectors"; + static final String SEG_IN_ROUND = "segment_in_round"; + static final String ROUND_ID = "round_id"; + + /** start inclusive, end exclusive */ + @AllArgsConstructor + static class DeletionRange { + int start; + int end; + } + + @Builder + static class MergeTestRound { + @Default + List segmentSizes = List.of(); + @Default + List deletionRanges = List.of(); + } + + @Builder + static class MergeTestScenario { + @Singular + List rounds; + @Default + int minPqThreshold = KNNConstants.DEFAULT_MINIMUM_BATCH_SIZE_FOR_QUANTIZATION; + @Default + int dimension = 128; + @Default + int nQueries = 10; + @Default + int topK = 10; + @Default + int overqueryFactor = KNNConstants.DEFAULT_OVER_QUERY_FACTOR; + @Default + double minimumRecall = 0.99; + @Default + boolean leadingSegmentMergeDisabled = KNNConstants.DEFAULT_LEADING_SEGMENT_MERGE_DISABLED; + } + + @Rule + public TemporaryFolder tempDir = new TemporaryFolder(); + + void runScenario(MergeTestScenario scenario) throws IOException { + // don't worry about deletions when generating base vectors, they will just become holes in the ids later + int nBase = scenario.rounds.stream().mapToInt(r -> r.segmentSizes.stream().mapToInt(x -> x).sum()).sum(); + var baseVecs = TestUtils.randomlyGenerateStandardVectors(nBase, scenario.dimension, RANDOM_SEED); + var queryVecs = TestUtils.randomlyGenerateStandardVectors(scenario.nQueries, scenario.dimension, RANDOM_SEED + 1); + + var liveVecs = new FixedBitSet(nBase); + liveVecs.clear(); + + // Path indexPath = createTempDir(); + IndexWriterConfig iwc = LuceneTestCase.newIndexWriterConfig(); + iwc.setUseCompoundFile(false); + iwc.setCodec(getCodec(scenario.minPqThreshold, scenario.leadingSegmentMergeDisabled)); + iwc.setMergePolicy(new ForceMergesOnlyMergePolicy(false)); + + try (var fsd = FSDirectory.open(tempDir.getRoot().toPath()); var writer = new IndexWriter(fsd, iwc);) { + int vectorOffset = 0; + for (int roundId = 0; roundId < scenario.rounds.size(); roundId++) { + var round = scenario.rounds.get(roundId); + + for (int segInRound = 0; segInRound < round.segmentSizes.size(); segInRound++) { + var segmentSize = round.segmentSizes.get(segInRound); + + for (int i = 0; i < segmentSize; i++) { + int id = vectorOffset + i; + Document doc = new Document(); + doc.add(new KnnFloatVectorField(VECTOR_FIELD, baseVecs[id], VectorSimilarityFunction.EUCLIDEAN)); + doc.add(new IntField(ID_FIELD, id, Store.YES)); + doc.add(new IntField(SEG_IN_ROUND, segInRound, Store.YES)); + doc.add(new IntField(ROUND_ID, roundId, Store.YES)); + writer.addDocument(doc); + + liveVecs.set(id); + } + writer.flush(); + writer.commit(); + vectorOffset += segmentSize; + } + + for (var dl : round.deletionRanges) { + var query = IntField.newRangeQuery(ID_FIELD, dl.start, dl.end - 1); // end inclusive + writer.deleteDocuments(query); + liveVecs.clear(dl.start, dl.end); + writer.commit(); + } + + writer.forceMerge(1); + writer.commit(); + + try (var reader = DirectoryReader.open(writer)) { + // assertEquals("Should have one segment after merge", 1, reader.getContext().leaves().size()); + Assert.assertTrue("Should have one segment after merge", 1 >= reader.getContext().leaves().size()); + assertEquals(liveVecs.cardinality(), reader.numDocs()); + + var searcher = LuceneTestCase.newSearcher(reader); + + var totalCorrect = 0; + for (var queryVec : queryVecs) { + var gt = computeL2GroundTruth(baseVecs, liveVecs, queryVec, scenario.topK); + var query = new JVectorKnnFloatVectorQuery( + VECTOR_FIELD, + queryVec, + scenario.topK, + scenario.overqueryFactor, + 0.0f, + 0.0f, + false + ); + var results = searcher.search(query, scenario.topK).scoreDocs; + + var pred = Arrays.stream(results).map(r -> { + try { + return reader.storedFields().document(r.doc).getField(ID_FIELD).numericValue().intValue(); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + }).collect(Collectors.toSet()); + + assertEquals(scenario.topK, pred.size()); + + var correct = 0; + for (var doc : pred) { + if (gt.contains(doc)) { + correct += 1; + } + } + totalCorrect += correct; + } + + double recall = totalCorrect / (double) (scenario.topK * queryVecs.length); + Assert.assertTrue( + String.format("Recall too low [RoundId %d], got %f < %f", roundId, recall, scenario.minimumRecall), + recall >= scenario.minimumRecall + ); + } + } + } + } + + Set computeL2GroundTruth(float[][] baseVectors, BitSet liveVectors, float[] query, int k) { + if (liveVectors.length() == 0) { + return Set.of(); + } + var pq = new PriorityQueue(k, (a, b) -> Float.compare(a.score, b.score)); + for (int i = liveVectors.nextSetBit(0); i != DocIdSetIterator.NO_MORE_DOCS; i = liveVectors.nextSetBit(i + 1)) { + var score = VectorSimilarityFunction.EUCLIDEAN.compare(baseVectors[i], query); + var sd = new ScoreDoc(i, score); + + if (pq.size() < k) { + pq.add(sd); + } else { + if (pq.peek().score < sd.score) { + pq.poll(); + pq.add(sd); + } + } + // otherwise nextBitSet(length()) throws assertion for a FixedBitSet + if (i + 1 >= liveVectors.length()) { + break; + } + } + return pq.stream().map(sd -> sd.doc).collect(Collectors.toSet()); + } + + @Test + public void simpleTest() throws IOException { + var scenario = MergeTestScenario.builder() + .minPqThreshold(Integer.MAX_VALUE) + .round(MergeTestRound.builder().segmentSizes(List.of(100)).build()) + .minimumRecall(1.0) + .build(); + runScenario(scenario); + } + + @Test + public void simpleMergeNoPQ() throws IOException { + var scenario = MergeTestScenario.builder() + .minPqThreshold(Integer.MAX_VALUE) + .round(MergeTestRound.builder().segmentSizes(List.of(100, 200)).build()) + .minimumRecall(1.0) + .build(); + runScenario(scenario); + } + + @Test + public void manySegmentMergeNoPq() throws IOException { + var scenario = MergeTestScenario.builder() + .minPqThreshold(Integer.MAX_VALUE) + .round(MergeTestRound.builder().segmentSizes(List.of(100, 200, 50, 250, 1)).build()) + .minimumRecall(1.0) + .build(); + runScenario(scenario); + } + + @Test + public void multiPhaseSegmentMergeNoPq() throws IOException { + var scenario = MergeTestScenario.builder() + .minPqThreshold(Integer.MAX_VALUE) + .round(MergeTestRound.builder().segmentSizes(List.of(10, 200, 50, 25, 1)).build()) + .round(MergeTestRound.builder().segmentSizes(List.of(10, 20, 150)).build()) + .round( + MergeTestRound.builder() + // note: this segment is large, should become new leading segment + .segmentSizes(List.of(20, 1000, 1)) + .build() + ) + .overqueryFactor(20) // can't achieve 1.0 recall otherwise + .minimumRecall(1.0) + .build(); + runScenario(scenario); + } + + @Test + public void simpleDeletionTest() throws IOException { + var scenario = MergeTestScenario.builder() + .minPqThreshold(Integer.MAX_VALUE) + .round(MergeTestRound.builder().segmentSizes(List.of(100)).deletionRanges(List.of(new DeletionRange(20, 70))).build()) + .minimumRecall(1.0) + .build(); + runScenario(scenario); + } + + @Test + public void multiPhaseSegmentMergeWithDeletionsNoPq() throws IOException { + var scenario = MergeTestScenario.builder() + .minPqThreshold(Integer.MAX_VALUE) + .round( + MergeTestRound.builder() + .segmentSizes(List.of(10, 200, 50, 100, 100)) + .deletionRanges( + List.of( + new DeletionRange(0, 5), + new DeletionRange(210, 260), // wipe out an enitre segment + new DeletionRange(330, 400), // deletion range spanning segments + new DeletionRange(410, 411) + ) + ) + .build() + ) + .round( + MergeTestRound.builder() + .segmentSizes(List.of(10, 20, 300)) + .deletionRanges( + List.of( + new DeletionRange(405, 600), // go wild + new DeletionRange(10, 11) // I just don't like index 10 anymore + // btw remember that deletions don't renumber ordinals as far as this test class is concerned + // max ordinal (and holes) just keep growing as you add more rounds + ) + ) + .build() + ) + .round( + MergeTestRound.builder() + .segmentSizes(List.of(20, 1000, 1)) + .deletionRanges(List.of(new DeletionRange(998, 1023), new DeletionRange(1023, 1100), new DeletionRange(1502, 1504))) + .build() + ) + .overqueryFactor(20) + .minimumRecall(0.99) + .build(); + runScenario(scenario); + } + + @Test + public void multiPhaseMergeWithDeletesProgressivePQ() throws IOException { + var scenario = MergeTestScenario.builder() + .minPqThreshold(1000) // start using PQ once the vector count crosses this + .round( + MergeTestRound.builder() + .segmentSizes(List.of(10, 200, 50, 100, 100)) // count 460 + .deletionRanges( + List.of( + // total deletions 126 + new DeletionRange(0, 5), // count 5 + new DeletionRange(210, 260), // count 50 + new DeletionRange(330, 400), // count 70 + new DeletionRange(410, 411) // count 1 + ) + ) + .build() // count 334 + ) + .round( + MergeTestRound.builder() + .segmentSizes(List.of(10, 20, 300)) // count 330 + .deletionRanges( + List.of( + // total deletions 195 + new DeletionRange(405, 600), // count 195, BUT this range has some overlap so actually 194 + new DeletionRange(10, 11) // count 1 + ) + ) + .build() // +135, total count 469 + ) + .round( + MergeTestRound.builder() + .segmentSizes(List.of(20, 2000, 1)) // I can't do more math, so we'll just add 2000 vectors to ensure PQ + .deletionRanges(List.of(new DeletionRange(998, 1023), new DeletionRange(1023, 1100), new DeletionRange(2502, 2504))) + .build() + ) + .round(MergeTestRound.builder().segmentSizes(List.of(400)).build()) // an extra round just because + .overqueryFactor(20) + .minimumRecall(0.98) + .build(); + runScenario(scenario); + } + + @Test + public void multiPhaseMergeWithDeletesAlwaysPQ() throws IOException { + var scenario = MergeTestScenario.builder() + .minPqThreshold(1) // always use PQ + .round( + MergeTestRound.builder() + .segmentSizes(List.of(10, 200, 50, 100, 100)) + .deletionRanges( + List.of( + new DeletionRange(0, 5), + new DeletionRange(210, 260), + new DeletionRange(330, 400), + new DeletionRange(410, 411) + ) + ) + .build() + ) + .round( + MergeTestRound.builder() + .segmentSizes(List.of(10, 20, 300)) + .deletionRanges(List.of(new DeletionRange(405, 600), new DeletionRange(10, 11))) + .build() + ) + .round( + MergeTestRound.builder() + .segmentSizes(List.of(20, 2000, 1)) + .deletionRanges(List.of(new DeletionRange(998, 1023), new DeletionRange(1023, 1100), new DeletionRange(2502, 2504))) + .build() + ) + .overqueryFactor(20) + .minimumRecall(0.99) + .build(); + runScenario(scenario); + } + + @Test + public void testLeadingSegmentMergeDisabled() throws IOException { + // we'll progressively increase to beyond the PQ threshold + var scenario = MergeTestScenario.builder() + .minPqThreshold(1000) // start using PQ once the vector count crosses this + .leadingSegmentMergeDisabled(true) // and disable leading segment merge + .round( + MergeTestRound.builder() + .segmentSizes(List.of(10, 200, 50, 100, 100)) // count 460 + .deletionRanges( + List.of( + // total deletions 126 + new DeletionRange(0, 5), // count 5 + new DeletionRange(210, 260), // count 50 + new DeletionRange(330, 400), // count 70 + new DeletionRange(410, 411) // count 1 + ) + ) + .build() // count 334 + ) + .round( + MergeTestRound.builder() + .segmentSizes(List.of(10, 20, 300)) // count 330 + .deletionRanges( + List.of( + // total deletions 195 + new DeletionRange(405, 600), // count 195, BUT this range has some overlap so actually 194 + new DeletionRange(10, 11) // count 1 + ) + ) + .build() // +135, total count 469 + ) + .round(MergeTestRound.builder().segmentSizes(List.of(50, 200)).build()) // bonus round + .round( + MergeTestRound.builder() + .segmentSizes(List.of(20, 2000, 1)) // Add 2000 vectors to ensure PQ + .deletionRanges(List.of(new DeletionRange(998, 1023), new DeletionRange(1023, 1100), new DeletionRange(2502, 2504))) + .build() + ) + .round(MergeTestRound.builder().segmentSizes(List.of(400)).build()) + .overqueryFactor(20) + .minimumRecall(0.99) + .build(); + runScenario(scenario); + } +} 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 d5b9a6e2..7f24a06c 100644 --- a/src/test/java/org/opensearch/knn/index/engine/CommonTestUtils.java +++ b/src/test/java/org/opensearch/knn/index/engine/CommonTestUtils.java @@ -49,6 +49,7 @@ import static org.opensearch.knn.common.KNNConstants.DISK_ANN; import static org.opensearch.knn.common.KNNConstants.VECTOR_DATA_TYPE_FIELD; import static org.opensearch.knn.index.KNNSettings.KNN_INDEX; +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; public class CommonTestUtils { @@ -123,10 +124,14 @@ public static String createIndexMapping(int dimension, SpaceType spaceType, Vect } public static Codec getCodec() { - return getCodec(DEFAULT_MINIMUM_BATCH_SIZE_FOR_QUANTIZATION); + return getCodec(DEFAULT_MINIMUM_BATCH_SIZE_FOR_QUANTIZATION, DEFAULT_LEADING_SEGMENT_MERGE_DISABLED); } public static Codec getCodec(int minBatchSizeForQuantization) { + return getCodec(minBatchSizeForQuantization, DEFAULT_LEADING_SEGMENT_MERGE_DISABLED); + } + + public static Codec getCodec(int minBatchSizeForQuantization, boolean leadingSegmentMergeDisabled) { return new FilterCodec(KNNCodecVersion.V_10_03_0.getCodecName(), new Lucene103Codec()) { @Override public KnnVectorsFormat knnVectorsFormat() { @@ -134,7 +139,7 @@ public KnnVectorsFormat knnVectorsFormat() { @Override public KnnVectorsFormat getKnnVectorsFormatForField(String field) { - return new JVectorFormat(minBatchSizeForQuantization); + return new JVectorFormat(minBatchSizeForQuantization, leadingSegmentMergeDisabled); } }; } diff --git a/src/test/java/org/opensearch/knn/index/engine/JVectorEngineIT.java b/src/test/java/org/opensearch/knn/index/engine/JVectorEngineIT.java index fe371e78..323be924 100644 --- a/src/test/java/org/opensearch/knn/index/engine/JVectorEngineIT.java +++ b/src/test/java/org/opensearch/knn/index/engine/JVectorEngineIT.java @@ -104,6 +104,7 @@ public void testAddDoc() throws Exception { .field(KNNConstants.METHOD_PARAMETER_NEIGHBOR_OVERFLOW, 2.0) .field(KNNConstants.METHOD_PARAMETER_HIERARCHY_ENABLED, false) .field(KNNConstants.METHOD_PARAMETER_MIN_BATCH_SIZE_FOR_QUANTIZATION, 1000) + .field(KNNConstants.METHOD_PARAMETER_LEADING_SEGMENT_MERGE_DISABLED, false) .endObject() .endObject() .endObject()