From d9b585c4e901430ec2efce9bb27a8e4862bf8eff Mon Sep 17 00:00:00 2001 From: luyuncheng Date: Wed, 23 Apr 2025 22:45:03 +0800 Subject: [PATCH] Optimize Merge on knn field Signed-off-by: luyuncheng --- .../index/engine/InternalEngine.java | 2 +- ...va => PrunePostingsAndKnnMergePolicy.java} | 115 +++++++++++++++++- .../index/engine/InternalEngineTests.java | 2 +- ... PrunePostingsAndKnnMergePolicyTests.java} | 109 ++++++++++++++++- 4 files changed, 222 insertions(+), 6 deletions(-) rename server/src/main/java/org/opensearch/index/engine/{PrunePostingsMergePolicy.java => PrunePostingsAndKnnMergePolicy.java} (67%) rename server/src/test/java/org/opensearch/index/engine/{PrunePostingsMergePolicyTests.java => PrunePostingsAndKnnMergePolicyTests.java} (62%) diff --git a/server/src/main/java/org/opensearch/index/engine/InternalEngine.java b/server/src/main/java/org/opensearch/index/engine/InternalEngine.java index 62bfd27516964..10c92eaaefb2a 100644 --- a/server/src/main/java/org/opensearch/index/engine/InternalEngine.java +++ b/server/src/main/java/org/opensearch/index/engine/InternalEngine.java @@ -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())); diff --git a/server/src/main/java/org/opensearch/index/engine/PrunePostingsMergePolicy.java b/server/src/main/java/org/opensearch/index/engine/PrunePostingsAndKnnMergePolicy.java similarity index 67% rename from server/src/main/java/org/opensearch/index/engine/PrunePostingsMergePolicy.java rename to server/src/main/java/org/opensearch/index/engine/PrunePostingsAndKnnMergePolicy.java index 87ac032b7d81b..4779eef914b16 100644 --- a/server/src/main/java/org/opensearch/index/engine/PrunePostingsMergePolicy.java +++ b/server/src/main/java/org/opensearch/index/engine/PrunePostingsAndKnnMergePolicy.java @@ -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; @@ -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 @@ -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 { @@ -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)) { + 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() { + return super.getVectorReader(); + } + @Override public CacheHelper getCoreCacheHelper() { return null; @@ -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(); + } + } } diff --git a/server/src/test/java/org/opensearch/index/engine/InternalEngineTests.java b/server/src/test/java/org/opensearch/index/engine/InternalEngineTests.java index fa501b155e96b..734293a906217 100644 --- a/server/src/test/java/org/opensearch/index/engine/InternalEngineTests.java +++ b/server/src/test/java/org/opensearch/index/engine/InternalEngineTests.java @@ -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") ) ) ) diff --git a/server/src/test/java/org/opensearch/index/engine/PrunePostingsMergePolicyTests.java b/server/src/test/java/org/opensearch/index/engine/PrunePostingsAndKnnMergePolicyTests.java similarity index 62% rename from server/src/test/java/org/opensearch/index/engine/PrunePostingsMergePolicyTests.java rename to server/src/test/java/org/opensearch/index/engine/PrunePostingsAndKnnMergePolicyTests.java index a21094152853d..ec514d61bf441 100644 --- a/server/src/test/java/org/opensearch/index/engine/PrunePostingsMergePolicyTests.java +++ b/server/src/test/java/org/opensearch/index/engine/PrunePostingsAndKnnMergePolicyTests.java @@ -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; @@ -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()) { @@ -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(); @@ -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(); + } + } }