Skip to content

[Enhancement] Optimize Merge On KNN Field when soft_delete mode #18050

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

Open
luyuncheng wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[Enhancement] Optimize Merge On KNN Field when soft_delete mode #18050

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
2 changes: 1 addition & 1 deletion server/src/main/java/org/opensearch/index/engine/InternalEngine.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2346,7 +2346,7 @@ private IndexWriterConfig getIndexWriterConfig() {
new SoftDeletesRetentionMergePolicy(
Lucene.SOFT_DELETES_FIELD,
softDeletesPolicy::getRetentionQuery,
new PrunePostingsMergePolicy(mergePolicy, IdFieldMapper.NAME)
new PrunePostingsAndKnnMergePolicy(mergePolicy, IdFieldMapper.NAME)
)
);
boolean shuffleForcedMerge = Booleans.parseBoolean(System.getProperty("opensearch.shuffle_forced_merge", Boolean.TRUE.toString()));
Expand Down
115 changes: 113 additions & 2 deletions ...ndex/engine/PrunePostingsMergePolicy.java → ...ngine/PrunePostingsAndKnnMergePolicy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,15 +32,24 @@

package org.opensearch.index.engine;

import org.apache.lucene.codecs.DocValuesProducer;
import org.apache.lucene.codecs.FieldsProducer;
import org.apache.lucene.codecs.KnnVectorsReader;
import org.apache.lucene.index.BinaryDocValues;
import org.apache.lucene.index.CodecReader;
import org.apache.lucene.index.DocValuesSkipper;
import org.apache.lucene.index.FieldInfo;
import org.apache.lucene.index.FilterCodecReader;
import org.apache.lucene.index.FilterLeafReader;
import org.apache.lucene.index.FilteredTermsEnum;
import org.apache.lucene.index.ImpactsEnum;
import org.apache.lucene.index.MergePolicy;
import org.apache.lucene.index.NumericDocValues;
import org.apache.lucene.index.OneMergeWrappingMergePolicy;
import org.apache.lucene.index.PostingsEnum;
import org.apache.lucene.index.SortedDocValues;
import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.index.SortedSetDocValues;
import org.apache.lucene.index.Terms;
import org.apache.lucene.index.TermsEnum;
import org.apache.lucene.search.DocIdSetIterator;
Expand All @@ -49,6 +58,7 @@

import java.io.IOException;
import java.util.Iterator;
import java.util.Objects;

/**
* This merge policy drops id field postings for all delete documents this can be
Expand All @@ -62,9 +72,11 @@
*
* @opensearch.internal
*/
final class PrunePostingsMergePolicy extends OneMergeWrappingMergePolicy {
final class PrunePostingsAndKnnMergePolicy extends OneMergeWrappingMergePolicy {

PrunePostingsMergePolicy(MergePolicy in, String idField) {
public static final String KNN_FIELD = "knn_field";

PrunePostingsAndKnnMergePolicy(MergePolicy in, String idField) {
super(in, toWrap -> new OneMerge(toWrap.segments) {
@Override
public CodecReader wrapForMerge(CodecReader reader) throws IOException {
Expand Down Expand Up @@ -156,6 +168,65 @@ public int size() {
};
}

@Override
public DocValuesProducer getDocValuesReader() {
DocValuesProducer docValuesProducer = super.getDocValuesReader();
if (docValuesProducer == null) {
return null;
}

return new DocValuesProducer() {
@Override
public NumericDocValues getNumeric(FieldInfo fieldInfo) throws IOException {
return docValuesProducer.getNumeric(fieldInfo);
}

@Override
public BinaryDocValues getBinary(FieldInfo fieldInfo) throws IOException {
BinaryDocValues binaryDocValues = docValuesProducer.getBinary(fieldInfo);
if (binaryDocValues != null && fieldInfo.attributes().containsKey(KNN_FIELD)) {

@bugmakerrrrrr bugmakerrrrrr Apr 23, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't know if we should expose something of the plugin within the core. Maybe we can provide an extension point in the engine plugin.

@navneet1v navneet1v Apr 25, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@luyuncheng I think we should just make this generic. Why we want to make this specific to k-NN. Docs containing BDV should just be removed if they are deleted. What is the point of having the deleted docs?

@luyuncheng luyuncheng Apr 25, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should just make this generic. Why we want to make this specific to k-NN. Docs containing BDV should just be removed if they are deleted. What is the point of having the deleted docs?

Good Idea. i think all BinaryDocValues should skipped on merge with soft_deletes.

return new MergeFilterBinaryDocValues(binaryDocValues, liveDocs);
}
return binaryDocValues;
}

@Override
public SortedDocValues getSorted(FieldInfo fieldInfo) throws IOException {
return docValuesProducer.getSorted(fieldInfo);
}

@Override
public SortedNumericDocValues getSortedNumeric(FieldInfo fieldInfo) throws IOException {
return docValuesProducer.getSortedNumeric(fieldInfo);
}

@Override
public SortedSetDocValues getSortedSet(FieldInfo fieldInfo) throws IOException {
return docValuesProducer.getSortedSet(fieldInfo);
}

@Override
public DocValuesSkipper getSkipper(FieldInfo fieldInfo) throws IOException {
return docValuesProducer.getSkipper(fieldInfo);
}

@Override
public void checkIntegrity() throws IOException {
docValuesProducer.checkIntegrity();
}

@Override
public void close() throws IOException {
docValuesProducer.close();
}
};
}

@Override
public KnnVectorsReader getVectorReader() {

@bugmakerrrrrr bugmakerrrrrr Apr 23, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why don't we pruning the vector field?

@navneet1v navneet1v Apr 25, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

+1 on this.

@luyuncheng luyuncheng Apr 25, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i need to override the KnnVectorsReader with different type like 'FloatVectorValues' and 'ByteVectorValues' and handle the iterator with different type and merge logic. so i think it would be better to create new PR and handle KnnVectorsReader carefully

return super.getVectorReader();
}

@Override
public CacheHelper getCoreCacheHelper() {
return null;
Expand Down Expand Up @@ -236,4 +307,44 @@ public BytesRef getPayload() throws IOException {
return in.getPayload();
}
}

private static final class MergeFilterBinaryDocValues extends BinaryDocValues {

protected final BinaryDocValues in;
protected final Bits liveDocs;

protected MergeFilterBinaryDocValues(BinaryDocValues in, Bits liveDocs) {
Objects.requireNonNull(in);
this.in = in;
this.liveDocs = liveDocs;
}

public int docID() {
return this.in.docID();
}

public int nextDoc() throws IOException {
int docId;
do {
docId = in.nextDoc();
} while (docId != DocIdSetIterator.NO_MORE_DOCS && liveDocs.get(docId) == false);
return docId;
}

public int advance(int target) throws IOException {
return this.in.advance(target);
}

public boolean advanceExact(int target) throws IOException {
return this.in.advanceExact(target);
}

public long cost() {
return this.in.cost();
}

public BytesRef binaryValue() throws IOException {
return this.in.binaryValue();
}
}
}
2 changes: 1 addition & 1 deletion server/src/test/java/org/opensearch/index/engine/InternalEngineTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1782,7 +1782,7 @@ public void testLookupVersionWithPrunedAwayIds() throws IOException {
new SoftDeletesRetentionMergePolicy(
Lucene.SOFT_DELETES_FIELD,
MatchAllDocsQuery::new,
new PrunePostingsMergePolicy(indexWriterConfig.getMergePolicy(), "_id")
new PrunePostingsAndKnnMergePolicy(indexWriterConfig.getMergePolicy(), "_id")
)
)
)
Expand Down
109 changes: 107 additions & 2 deletions ...engine/PrunePostingsMergePolicyTests.java → .../PrunePostingsAndKnnMergePolicyTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,9 +34,15 @@

import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.FieldType;
import org.apache.lucene.document.KnnFloatVectorField;
import org.apache.lucene.document.NumericDocValuesField;
import org.apache.lucene.document.StringField;
import org.apache.lucene.index.BinaryDocValues;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.DocValuesType;
import org.apache.lucene.index.FloatVectorValues;
import org.apache.lucene.index.IndexOptions;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.LeafReader;
Expand All @@ -56,7 +62,7 @@

import java.io.IOException;

public class PrunePostingsMergePolicyTests extends OpenSearchTestCase {
public class PrunePostingsAndKnnMergePolicyTests extends OpenSearchTestCase {

public void testPrune() throws IOException {
try (Directory dir = newDirectory()) {
Expand All @@ -65,7 +71,7 @@ public void testPrune() throws IOException {
MergePolicy mp = new SoftDeletesRetentionMergePolicy(
"_soft_deletes",
MatchAllDocsQuery::new,
new PrunePostingsMergePolicy(newLogMergePolicy(), "id")
new PrunePostingsAndKnnMergePolicy(newLogMergePolicy(), "id")
);
iwc.setMergePolicy(new ShuffleForcedMergePolicy(mp));
boolean sorted = randomBoolean();
Expand Down Expand Up @@ -177,4 +183,103 @@ public void testPrune() throws IOException {

}
}

public void testPruneBinaryDocValues() throws IOException {
try (Directory dir = newDirectory()) {
IndexWriterConfig iwc = newIndexWriterConfig();
iwc.setSoftDeletesField("_soft_deletes");
MergePolicy mp = new SoftDeletesRetentionMergePolicy(
"_soft_deletes",
MatchAllDocsQuery::new,
new PrunePostingsAndKnnMergePolicy(newLogMergePolicy(), "id")
);
iwc.setMergePolicy(new ShuffleForcedMergePolicy(mp));
boolean sorted = randomBoolean();
if (sorted) {
iwc.setIndexSort(new Sort(new SortField("sort", SortField.Type.INT)));
}
int numUniqueDocs = randomIntBetween(1, 100);
int numDocs = randomIntBetween(numUniqueDocs, numUniqueDocs * 5);

try (IndexWriter writer = new IndexWriter(dir, iwc)) {
for (int i = 0; i < numDocs; i++) {
if (rarely()) {
writer.flush();
}
if (rarely()) {
writer.forceMerge(1, false);
}
int id = i % numUniqueDocs;
Document doc = new Document();
doc.add(new StringField("id", "" + id, Field.Store.NO));
doc.add(newTextField("text", "the quick brown fox", Field.Store.YES));
doc.add(new NumericDocValuesField("sort", i));
doc.add(new KNNDocValuesFiled("binary", new BytesRef("binary docvalues")));
doc.add(new KnnFloatVectorField("knnVector", new float[] { randomFloat(), randomFloat() }));
writer.softUpdateDocument(new Term("id", "" + id), doc, new NumericDocValuesField("_soft_deletes", 1));
if (i == 0) {
// make sure we have at least 2 segments to ensure we do an actual merge to kick out all postings for
// soft deletes
writer.flush();
}
}

writer.forceMerge(1);
try (DirectoryReader reader = DirectoryReader.open(writer)) {
LeafReader leafReader = reader.leaves().get(0).reader();
assertEquals(numDocs, leafReader.maxDoc());
Terms id = leafReader.terms("id");
TermsEnum iterator = id.iterator();
for (int i = 0; i < numUniqueDocs; i++) {
assertTrue(iterator.seekExact(new BytesRef("" + i)));
assertEquals(1, iterator.docFreq());
}
iterator = leafReader.terms("text").iterator();
assertTrue(iterator.seekExact(new BytesRef("quick")));
assertEquals(leafReader.maxDoc(), iterator.docFreq());
int numValues = 0;
NumericDocValues sort = leafReader.getNumericDocValues("sort");
while (sort.nextDoc() != DocIdSetIterator.NO_MORE_DOCS) {
if (sorted) {
assertEquals(sort.docID(), sort.longValue());
} else {
assertTrue(sort.longValue() >= 0);
assertTrue(sort.longValue() < numDocs);
}
numValues++;
}
assertEquals(numValues, numDocs);

int numBinary = 0;
BinaryDocValues binaryDocValues = leafReader.getBinaryDocValues("binary");
while (binaryDocValues.nextDoc() != DocIdSetIterator.NO_MORE_DOCS) {
numBinary++;
}
assertEquals(numBinary, numUniqueDocs);

int numVectors = 0;
FloatVectorValues floatVectorValues = leafReader.getFloatVectorValues("knnVector");
assertEquals(floatVectorValues.size(), numDocs);
}
}

}
}

public class KNNDocValuesFiled extends Field {
public static final FieldType TYPE = new FieldType();

public KNNDocValuesFiled(String name, BytesRef value) {
super(name, TYPE);
this.fieldsData = value;
}

static {
TYPE.setTokenized(false);
TYPE.setIndexOptions(IndexOptions.NONE);
TYPE.setDocValuesType(DocValuesType.BINARY);
TYPE.putAttribute("knn_field", "true");
TYPE.freeze();
}
}
}