From e9a19ed131ba74895a672b4f48846b69df079e04 Mon Sep 17 00:00:00 2001 From: Martin Gaievski Date: Mon, 21 Mar 2022 13:01:17 -0700 Subject: [PATCH 1/9] Initial integration with OS 2.0 Alpha1 Signed-off-by: Martin Gaievski --- build.gradle | 2 +- .../org/opensearch/knn/index/KNNQuery.java | 6 + .../org/opensearch/knn/index/KNNWeight.java | 4 - .../index/codec/KNN80Codec/KNN80Codec.java | 6 + .../codec/KNN80Codec/KNN80CompoundFormat.java | 2 +- .../KNN80Codec/KNN80DocValuesFormat.java | 2 +- .../index/codec/KNN84Codec/KNN84Codec.java | 6 + .../index/codec/KNN86Codec/KNN86Codec.java | 6 + .../index/codec/KNN87Codec/KNN87Codec.java | 2 +- .../KNN91Codec/KNN91BinaryDocValues.java | 64 +++ .../index/codec/KNN91Codec/KNN91Codec.java | 58 +++ .../codec/KNN91Codec/KNN91CompoundFormat.java | 68 +++ .../KNN91Codec/KNN91DocValuesConsumer.java | 257 +++++++++++ .../KNN91Codec/KNN91DocValuesFormat.java | 46 ++ .../KNN91Codec/KNN91DocValuesReader.java | 53 +++ .../knn/index/codec/KNNCodecService.java | 4 +- .../org/opensearch/knn/indices/ModelDao.java | 11 +- .../services/org.apache.lucene.codecs.Codec | 3 +- .../java/org/opensearch/knn/TestUtils.java | 4 +- .../index/KNNCreateIndexFromModelTests.java | 25 +- .../index/KNNVectorDVLeafFieldDataTests.java | 2 +- .../index/KNNVectorIndexFieldDataTests.java | 2 +- .../index/KNNVectorScriptDocValuesTests.java | 2 +- .../codec/KNN87Codec/KNN87CodecTests.java | 2 + .../KNN91Codec/KNN91BinaryDocValuesTests.java | 69 +++ .../codec/KNN91Codec/KNN91CodecTests.java | 22 + .../KNN91Codec/KNN91CompoundFormatTests.java | 92 ++++ .../KNN91DocValuesConsumerTests.java | 412 ++++++++++++++++++ .../knn/index/codec/KNNCodecTestCase.java | 6 +- .../knn/index/codec/KNNCodecTestUtil.java | 22 +- .../plugin/script/KNNScoringUtilTests.java | 2 +- 31 files changed, 1226 insertions(+), 36 deletions(-) create mode 100644 src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91BinaryDocValues.java create mode 100644 src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91Codec.java create mode 100644 src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CompoundFormat.java create mode 100644 src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesConsumer.java create mode 100644 src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesFormat.java create mode 100644 src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesReader.java create mode 100644 src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91BinaryDocValuesTests.java create mode 100644 src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CodecTests.java create mode 100644 src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CompoundFormatTests.java create mode 100644 src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesConsumerTests.java diff --git a/build.gradle b/build.gradle index 1c7e6175b..eb22c1e52 100644 --- a/build.gradle +++ b/build.gradle @@ -13,7 +13,7 @@ buildscript { // explicitly as 'opensearch.version' property, for instance opensearch.version=2.0.0-alpha1-SNAPSHOT opensearch_version = System.getProperty("opensearch.version", "2.0.0-SNAPSHOT") knn_bwc_version = System.getProperty("bwc.version", "1.2.0.0-SNAPSHOT") - version_qualifier = System.getProperty("build.version_qualifier", "") + version_qualifier = System.getProperty("build.version_qualifier", "alpha1") opensearch_bwc_version = "${knn_bwc_version}" - ".0-SNAPSHOT" opensearch_group = "org.opensearch" } diff --git a/src/main/java/org/opensearch/knn/index/KNNQuery.java b/src/main/java/org/opensearch/knn/index/KNNQuery.java index f3bf23aee..76709ae7e 100644 --- a/src/main/java/org/opensearch/knn/index/KNNQuery.java +++ b/src/main/java/org/opensearch/knn/index/KNNQuery.java @@ -7,6 +7,7 @@ import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.Query; +import org.apache.lucene.search.QueryVisitor; import org.apache.lucene.search.ScoreMode; import org.apache.lucene.search.Weight; @@ -59,6 +60,11 @@ public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float bo return new KNNWeight(this, boost); } + @Override + public void visit(QueryVisitor visitor) { + + } + @Override public String toString(String field) { return field; diff --git a/src/main/java/org/opensearch/knn/index/KNNWeight.java b/src/main/java/org/opensearch/knn/index/KNNWeight.java index 6c8c41cb9..3f777c970 100644 --- a/src/main/java/org/opensearch/knn/index/KNNWeight.java +++ b/src/main/java/org/opensearch/knn/index/KNNWeight.java @@ -77,10 +77,6 @@ public Explanation explain(LeafReaderContext context, int doc) { return Explanation.match(1.0f, "No Explanation"); } - @Override - public void extractTerms(Set terms) { - } - @Override public Scorer scorer(LeafReaderContext context) throws IOException { SegmentReader reader = (SegmentReader) FilterLeafReader.unwrap(context.reader()); diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80Codec.java b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80Codec.java index 59655762e..5eb36e748 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80Codec.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80Codec.java @@ -11,6 +11,7 @@ import org.apache.lucene.codecs.CompoundFormat; import org.apache.lucene.codecs.DocValuesFormat; import org.apache.lucene.codecs.FieldInfosFormat; +import org.apache.lucene.codecs.KnnVectorsFormat; import org.apache.lucene.codecs.LiveDocsFormat; import org.apache.lucene.codecs.NormsFormat; import org.apache.lucene.codecs.PointsFormat; @@ -112,4 +113,9 @@ public CompoundFormat compoundFormat() { public PointsFormat pointsFormat() { return getDelegatee().pointsFormat(); } + + @Override + public KnnVectorsFormat knnVectorsFormat() { + throw new UnsupportedOperationException("Codec does not support knn vector format"); + } } diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80CompoundFormat.java b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80CompoundFormat.java index 8a8ed558e..d001cd1ba 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80CompoundFormat.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80CompoundFormat.java @@ -5,7 +5,7 @@ package org.opensearch.knn.index.codec.KNN80Codec; -import org.apache.lucene.codecs.lucene50.Lucene50CompoundFormat; +import org.apache.lucene.backward_codecs.lucene50.Lucene50CompoundFormat; import org.opensearch.knn.common.KNNConstants; import org.apache.lucene.codecs.CompoundDirectory; import org.apache.lucene.codecs.CompoundFormat; diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesFormat.java b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesFormat.java index cd8362cf2..fe329eb1c 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesFormat.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesFormat.java @@ -8,7 +8,7 @@ import org.apache.lucene.codecs.DocValuesConsumer; import org.apache.lucene.codecs.DocValuesFormat; import org.apache.lucene.codecs.DocValuesProducer; -import org.apache.lucene.codecs.lucene80.Lucene80DocValuesFormat; +import org.apache.lucene.backward_codecs.lucene80.Lucene80DocValuesFormat; import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.SegmentWriteState; diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN84Codec/KNN84Codec.java b/src/main/java/org/opensearch/knn/index/codec/KNN84Codec/KNN84Codec.java index a50f396a4..c4ae7ab2e 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN84Codec/KNN84Codec.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN84Codec/KNN84Codec.java @@ -5,6 +5,7 @@ package org.opensearch.knn.index.codec.KNN84Codec; +import org.apache.lucene.codecs.KnnVectorsFormat; import org.opensearch.knn.index.codec.KNN80Codec.KNN80CompoundFormat; import org.opensearch.knn.index.codec.KNN80Codec.KNN80DocValuesFormat; import org.apache.logging.log4j.LogManager; @@ -116,4 +117,9 @@ public CompoundFormat compoundFormat() { public PointsFormat pointsFormat() { return getDelegatee().pointsFormat(); } + + @Override + public KnnVectorsFormat knnVectorsFormat() { + return getDelegatee().knnVectorsFormat(); + } } diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN86Codec/KNN86Codec.java b/src/main/java/org/opensearch/knn/index/codec/KNN86Codec/KNN86Codec.java index 70c75e09b..154f44f1f 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN86Codec/KNN86Codec.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN86Codec/KNN86Codec.java @@ -5,6 +5,7 @@ package org.opensearch.knn.index.codec.KNN86Codec; +import org.apache.lucene.codecs.KnnVectorsFormat; import org.opensearch.knn.index.codec.KNN80Codec.KNN80CompoundFormat; import org.opensearch.knn.index.codec.KNN80Codec.KNN80DocValuesFormat; import org.apache.logging.log4j.LogManager; @@ -125,4 +126,9 @@ public CompoundFormat compoundFormat() { public PointsFormat pointsFormat() { return getDelegatee().pointsFormat(); } + + @Override + public KnnVectorsFormat knnVectorsFormat() { + return getDelegatee().knnVectorsFormat(); + } } diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN87Codec/KNN87Codec.java b/src/main/java/org/opensearch/knn/index/codec/KNN87Codec/KNN87Codec.java index 6ec5ec05c..20799648c 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN87Codec/KNN87Codec.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN87Codec/KNN87Codec.java @@ -6,7 +6,7 @@ package org.opensearch.knn.index.codec.KNN87Codec; import org.apache.lucene.codecs.FilterCodec; -import org.apache.lucene.codecs.lucene87.Lucene87Codec; +import org.apache.lucene.backward_codecs.lucene87.Lucene87Codec; import org.opensearch.knn.index.codec.KNN80Codec.KNN80CompoundFormat; import org.opensearch.knn.index.codec.KNN80Codec.KNN80DocValuesFormat; import org.apache.lucene.codecs.Codec; diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91BinaryDocValues.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91BinaryDocValues.java new file mode 100644 index 000000000..f9f89e589 --- /dev/null +++ b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91BinaryDocValues.java @@ -0,0 +1,64 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.index.codec.KNN91Codec; + +import org.apache.lucene.index.BinaryDocValues; +import org.apache.lucene.index.DocIDMerger; +import org.apache.lucene.util.BytesRef; +import org.opensearch.knn.index.codec.util.BinaryDocValuesSub; + +import java.io.IOException; + +/** + * A per-document kNN numeric value. + */ +class KNN91BinaryDocValues extends BinaryDocValues { + + private DocIDMerger docIDMerger; + + KNN91BinaryDocValues(DocIDMerger docIdMerger) { + this.docIDMerger = docIdMerger; + } + + private BinaryDocValuesSub current; + private int docID = -1; + + @Override + public int docID() { + return docID; + } + + @Override + public int nextDoc() throws IOException { + current = docIDMerger.next(); + if (current == null) { + docID = NO_MORE_DOCS; + } else { + docID = current.mappedDocID; + } + return docID; + } + + @Override + public int advance(int target) throws IOException { + throw new UnsupportedOperationException(); + } + + @Override + public boolean advanceExact(int target) throws IOException { + throw new UnsupportedOperationException(); + } + + @Override + public long cost() { + throw new UnsupportedOperationException(); + } + + @Override + public BytesRef binaryValue() throws IOException { + return current.getValues().binaryValue(); + } +}; diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91Codec.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91Codec.java new file mode 100644 index 000000000..1d3bb9580 --- /dev/null +++ b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91Codec.java @@ -0,0 +1,58 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.knn.index.codec.KNN91Codec; + +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.CompoundFormat; +import org.apache.lucene.codecs.DocValuesFormat; +import org.apache.lucene.codecs.FilterCodec; +import org.apache.lucene.codecs.lucene91.Lucene91Codec; + +/** + * Extends the Codec to support a new file format for KNN index + * based on the mappings. + * + */ +public final class KNN91Codec extends FilterCodec { + + private final DocValuesFormat docValuesFormat; + private final CompoundFormat compoundFormat; + + public static final String KNN_91 = "KNN91Codec"; + + /** + * No arg constructor that uses Lucene91 as the delegate + */ + public KNN91Codec() { + this(new Lucene91Codec()); + } + /** + * Constructor that takes a Codec delegate to delegate all methods this code does not implement to. + * + * @param delegate codec that will perform all operations this codec does not override + */ + public KNN91Codec(Codec delegate) { + super(KNN_91, delegate); + this.docValuesFormat = new KNN91DocValuesFormat(delegate.docValuesFormat()); + this.compoundFormat = new KNN91CompoundFormat(delegate.compoundFormat()); + } + + @Override + public DocValuesFormat docValuesFormat() { + return this.docValuesFormat; + } + + @Override + public CompoundFormat compoundFormat() { + return this.compoundFormat; + } +} diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CompoundFormat.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CompoundFormat.java new file mode 100644 index 000000000..6d8b0fdb7 --- /dev/null +++ b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CompoundFormat.java @@ -0,0 +1,68 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.index.codec.KNN91Codec; + +import org.apache.lucene.codecs.CompoundDirectory; +import org.apache.lucene.codecs.CompoundFormat; +import org.apache.lucene.index.SegmentInfo; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IOContext; +import org.opensearch.knn.common.KNNConstants; +import org.opensearch.knn.index.util.KNNEngine; + +import java.io.IOException; +import java.util.HashSet; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * Class to encode/decode compound file + */ +public class KNN91CompoundFormat extends CompoundFormat { + + private final CompoundFormat delegate; + + /** + * Constructor that takes a delegate to handle non-overridden methods + * + * @param delegate CompoundFormat that will handle non-overridden methods + */ + public KNN91CompoundFormat(CompoundFormat delegate) { + this.delegate = delegate; + } + + @Override + public CompoundDirectory getCompoundReader(Directory dir, SegmentInfo si, IOContext context) throws IOException { + return delegate.getCompoundReader(dir, si, context); + } + + @Override + public void write(Directory dir, SegmentInfo si, IOContext context) throws IOException { + for (KNNEngine knnEngine : KNNEngine.values()) { + writeEngineFiles(dir, si, context, knnEngine.getExtension()); + } + delegate.write(dir, si, context); + } + + private void writeEngineFiles(Directory dir, SegmentInfo si, IOContext context, String engineExtension) throws IOException { + /* + * If engine file present, remove it from the compounding file list to avoid header/footer checks + * and create a new compounding file format with extension engine + c. + */ + Set engineFiles = si.files().stream().filter(file -> file.endsWith(engineExtension)).collect(Collectors.toSet()); + + Set segmentFiles = new HashSet<>(si.files()); + + if (!engineFiles.isEmpty()) { + for (String engineFile : engineFiles) { + String engineCompoundFile = engineFile + KNNConstants.COMPOUND_EXTENSION; + dir.copyFrom(dir, engineFile, engineCompoundFile, context); + } + segmentFiles.removeAll(engineFiles); + si.setFiles(segmentFiles); + } + } +} diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesConsumer.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesConsumer.java new file mode 100644 index 000000000..058ab755b --- /dev/null +++ b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesConsumer.java @@ -0,0 +1,257 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.index.codec.KNN91Codec; + +import com.google.common.collect.ImmutableMap; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.apache.lucene.codecs.CodecUtil; +import org.apache.lucene.codecs.DocValuesConsumer; +import org.apache.lucene.codecs.DocValuesProducer; +import org.apache.lucene.index.BinaryDocValues; +import org.apache.lucene.index.DocValuesType; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.MergeState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.store.FSDirectory; +import org.apache.lucene.store.FilterDirectory; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.util.IOUtils; +import org.opensearch.common.xcontent.DeprecationHandler; +import org.opensearch.common.xcontent.NamedXContentRegistry; +import org.opensearch.common.xcontent.XContentFactory; +import org.opensearch.common.xcontent.XContentType; +import org.opensearch.knn.common.KNNConstants; +import org.opensearch.knn.index.KNNSettings; +import org.opensearch.knn.index.KNNVectorFieldMapper; +import org.opensearch.knn.index.SpaceType; +import org.opensearch.knn.index.codec.util.KNNCodecUtil; +import org.opensearch.knn.index.util.KNNEngine; +import org.opensearch.knn.indices.Model; +import org.opensearch.knn.indices.ModelCache; +import org.opensearch.knn.jni.JNIService; +import org.opensearch.knn.plugin.stats.KNNCounter; + +import java.io.Closeable; +import java.io.IOException; +import java.nio.file.Paths; +import java.security.AccessController; +import java.security.PrivilegedAction; +import java.util.HashMap; +import java.util.Map; + +import static org.opensearch.knn.common.KNNConstants.MODEL_ID; +import static org.opensearch.knn.common.KNNConstants.PARAMETERS; +import static org.opensearch.knn.index.codec.util.KNNCodecUtil.buildEngineFileName; + +/** + * This class writes the KNN docvalues to the segments + */ +class KNN91DocValuesConsumer extends DocValuesConsumer implements Closeable { + + private final Logger logger = LogManager.getLogger(KNN91DocValuesConsumer.class); + + private final String TEMP_SUFFIX = "tmp"; + private DocValuesConsumer delegatee; + private SegmentWriteState state; + + KNN91DocValuesConsumer(DocValuesConsumer delegatee, SegmentWriteState state) throws IOException { + this.delegatee = delegatee; + this.state = state; + } + + @Override + public void addBinaryField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException { + delegatee.addBinaryField(field, valuesProducer); + if (field.attributes().containsKey(KNNVectorFieldMapper.KNN_FIELD)) { + addKNNBinaryField(field, valuesProducer); + } + } + + public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException { + + // Get values to be indexed + BinaryDocValues values = valuesProducer.getBinary(field); + KNNCodecUtil.Pair pair = KNNCodecUtil.getFloats(values); + if (pair.vectors.length == 0 || pair.docs.length == 0) { + logger.info("Skipping engine index creation as there are no vectors or docs in the documents"); + return; + } + + // Increment counter for number of graph index requests + KNNCounter.GRAPH_INDEX_REQUESTS.increment(); + + // Create library index either from model or from scratch + String engineFileName; + String indexPath; + String tmpEngineFileName; + + if (field.attributes().containsKey(MODEL_ID)) { + + String modelId = field.attributes().get(MODEL_ID); + Model model = ModelCache.getInstance().get(modelId); + + KNNEngine knnEngine = model.getModelMetadata().getKnnEngine(); + + engineFileName = buildEngineFileName( + state.segmentInfo.name, + knnEngine.getLatestBuildVersion(), + field.name, + knnEngine.getExtension() + ); + indexPath = Paths.get(((FSDirectory) (FilterDirectory.unwrap(state.directory))).getDirectory().toString(), engineFileName) + .toString(); + tmpEngineFileName = engineFileName + TEMP_SUFFIX; + String tempIndexPath = indexPath + TEMP_SUFFIX; + + if (model.getModelBlob() == null) { + throw new RuntimeException("There is no trained model with id \"" + modelId + "\""); + } + + createKNNIndexFromTemplate(model.getModelBlob(), pair, knnEngine, tempIndexPath); + } else { + + // Get engine to be used for indexing + String engineName = field.attributes().getOrDefault(KNNConstants.KNN_ENGINE, KNNEngine.DEFAULT.getName()); + KNNEngine knnEngine = KNNEngine.getEngine(engineName); + + engineFileName = buildEngineFileName( + state.segmentInfo.name, + knnEngine.getLatestBuildVersion(), + field.name, + knnEngine.getExtension() + ); + indexPath = Paths.get(((FSDirectory) (FilterDirectory.unwrap(state.directory))).getDirectory().toString(), engineFileName) + .toString(); + tmpEngineFileName = engineFileName + TEMP_SUFFIX; + String tempIndexPath = indexPath + TEMP_SUFFIX; + + createKNNIndexFromScratch(field, pair, knnEngine, tempIndexPath); + } + + /* + * Adds Footer to the serialized graph + * 1. Copies the serialized graph to new file. + * 2. Adds Footer to the new file. + * + * We had to create new file here because adding footer directly to the + * existing file will miss calculating checksum for the serialized graph + * bytes and result in index corruption issues. + */ + // TODO: I think this can be refactored to avoid this copy and then write + // https://github.com/opendistro-for-elasticsearch/k-NN/issues/330 + try ( + IndexInput is = state.directory.openInput(tmpEngineFileName, state.context); + IndexOutput os = state.directory.createOutput(engineFileName, state.context) + ) { + os.copyBytes(is, is.length()); + CodecUtil.writeFooter(os); + } catch (Exception ex) { + KNNCounter.GRAPH_INDEX_ERRORS.increment(); + throw new RuntimeException("[KNN] Adding footer to serialized graph failed: " + ex); + } finally { + IOUtils.deleteFilesIgnoringExceptions(state.directory, tmpEngineFileName); + } + } + + private void createKNNIndexFromTemplate(byte[] model, KNNCodecUtil.Pair pair, KNNEngine knnEngine, String indexPath) { + Map parameters = ImmutableMap.of( + KNNConstants.INDEX_THREAD_QTY, + KNNSettings.state().getSettingValue(KNNSettings.KNN_ALGO_PARAM_INDEX_THREAD_QTY) + ); + AccessController.doPrivileged((PrivilegedAction) () -> { + JNIService.createIndexFromTemplate(pair.docs, pair.vectors, indexPath, model, parameters, knnEngine.getName()); + return null; + }); + } + + private void createKNNIndexFromScratch(FieldInfo fieldInfo, KNNCodecUtil.Pair pair, KNNEngine knnEngine, String indexPath) + throws IOException { + Map parameters = new HashMap<>(); + Map fieldAttributes = fieldInfo.attributes(); + String parametersString = fieldAttributes.get(KNNConstants.PARAMETERS); + + // parametersString will be null when legacy mapper is used + if (parametersString == null) { + parameters.put(KNNConstants.SPACE_TYPE, fieldAttributes.getOrDefault(KNNConstants.SPACE_TYPE, SpaceType.DEFAULT.getValue())); + + String efConstruction = fieldAttributes.get(KNNConstants.HNSW_ALGO_EF_CONSTRUCTION); + Map algoParams = new HashMap<>(); + if (efConstruction != null) { + algoParams.put(KNNConstants.METHOD_PARAMETER_EF_CONSTRUCTION, Integer.parseInt(efConstruction)); + } + + String m = fieldAttributes.get(KNNConstants.HNSW_ALGO_M); + if (m != null) { + algoParams.put(KNNConstants.METHOD_PARAMETER_M, Integer.parseInt(m)); + } + parameters.put(PARAMETERS, algoParams); + } else { + parameters.putAll( + XContentFactory.xContent(XContentType.JSON) + .createParser(NamedXContentRegistry.EMPTY, DeprecationHandler.THROW_UNSUPPORTED_OPERATION, parametersString) + .map() + ); + } + + // Used to determine how many threads to use when indexing + parameters.put(KNNConstants.INDEX_THREAD_QTY, KNNSettings.state().getSettingValue(KNNSettings.KNN_ALGO_PARAM_INDEX_THREAD_QTY)); + + // Pass the path for the nms library to save the file + AccessController.doPrivileged((PrivilegedAction) () -> { + JNIService.createIndex(pair.docs, pair.vectors, indexPath, parameters, knnEngine.getName()); + return null; + }); + } + + /** + * Merges in the fields from the readers in mergeState + * + * @param mergeState Holds common state used during segment merging + */ + @Override + public void merge(MergeState mergeState) { + try { + delegatee.merge(mergeState); + assert mergeState != null; + assert mergeState.mergeFieldInfos != null; + for (FieldInfo fieldInfo : mergeState.mergeFieldInfos) { + DocValuesType type = fieldInfo.getDocValuesType(); + if (type == DocValuesType.BINARY && fieldInfo.attributes().containsKey(KNNVectorFieldMapper.KNN_FIELD)) { + addKNNBinaryField(fieldInfo, new KNN91DocValuesReader(mergeState)); + } + } + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + @Override + public void addSortedSetField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException { + delegatee.addSortedSetField(field, valuesProducer); + } + + @Override + public void addSortedNumericField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException { + delegatee.addSortedNumericField(field, valuesProducer); + } + + @Override + public void addSortedField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException { + delegatee.addSortedField(field, valuesProducer); + } + + @Override + public void addNumericField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException { + delegatee.addNumericField(field, valuesProducer); + } + + @Override + public void close() throws IOException { + delegatee.close(); + } +} diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesFormat.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesFormat.java new file mode 100644 index 000000000..2d5613ace --- /dev/null +++ b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesFormat.java @@ -0,0 +1,46 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.index.codec.KNN91Codec; + +import org.apache.lucene.codecs.DocValuesConsumer; +import org.apache.lucene.codecs.DocValuesFormat; +import org.apache.lucene.codecs.DocValuesProducer; +import org.apache.lucene.codecs.lucene90.Lucene90DocValuesFormat; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; + +import java.io.IOException; + +/** + * Encodes/Decodes per document values + */ +public class KNN91DocValuesFormat extends DocValuesFormat { + private final DocValuesFormat delegate; + + public KNN91DocValuesFormat() { + this(new Lucene90DocValuesFormat()); + } + + /** + * Constructor that takes delegate in order to handle non-overridden methods + * + * @param delegate DocValuesFormat to handle non-overridden methods + */ + public KNN91DocValuesFormat(DocValuesFormat delegate) { + super(delegate.getName()); + this.delegate = delegate; + } + + @Override + public DocValuesConsumer fieldsConsumer(SegmentWriteState state) throws IOException { + return new KNN91DocValuesConsumer(delegate.fieldsConsumer(state), state); + } + + @Override + public DocValuesProducer fieldsProducer(SegmentReadState state) throws IOException { + return delegate.fieldsProducer(state); + } +} diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesReader.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesReader.java new file mode 100644 index 000000000..b6a9622ef --- /dev/null +++ b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesReader.java @@ -0,0 +1,53 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.index.codec.KNN91Codec; + +import org.apache.lucene.codecs.DocValuesProducer; +import org.apache.lucene.index.BinaryDocValues; +import org.apache.lucene.index.DocIDMerger; +import org.apache.lucene.index.DocValuesType; +import org.apache.lucene.index.EmptyDocValuesProducer; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.MergeState; +import org.opensearch.knn.index.codec.util.BinaryDocValuesSub; + +import java.util.ArrayList; +import java.util.List; + +/** + * Reader for KNNDocValues from the segments + */ +class KNN91DocValuesReader extends EmptyDocValuesProducer { + + private final MergeState mergeState; + + KNN91DocValuesReader(MergeState mergeState) { + this.mergeState = mergeState; + } + + @Override + public BinaryDocValues getBinary(FieldInfo field) { + try { + List subs = new ArrayList<>(this.mergeState.docValuesProducers.length); + for (int i = 0; i < this.mergeState.docValuesProducers.length; i++) { + BinaryDocValues values = null; + DocValuesProducer docValuesProducer = mergeState.docValuesProducers[i]; + if (docValuesProducer != null) { + FieldInfo readerFieldInfo = mergeState.fieldInfos[i].fieldInfo(field.name); + if (readerFieldInfo != null && readerFieldInfo.getDocValuesType() == DocValuesType.BINARY) { + values = docValuesProducer.getBinary(readerFieldInfo); + } + if (values != null) { + subs.add(new BinaryDocValuesSub(mergeState.docMaps[i], values)); + } + } + } + return new KNN91BinaryDocValues(DocIDMerger.of(subs, mergeState.needsIndexSort)); + } catch (Exception e) { + throw new RuntimeException(e); + } + } +} diff --git a/src/main/java/org/opensearch/knn/index/codec/KNNCodecService.java b/src/main/java/org/opensearch/knn/index/codec/KNNCodecService.java index 70991dfe8..8bbf37b03 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNNCodecService.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNNCodecService.java @@ -6,9 +6,9 @@ package org.opensearch.knn.index.codec; import org.opensearch.index.codec.CodecServiceConfig; -import org.opensearch.knn.index.codec.KNN87Codec.KNN87Codec; import org.apache.lucene.codecs.Codec; import org.opensearch.index.codec.CodecService; +import org.opensearch.knn.index.codec.KNN91Codec.KNN91Codec; /** * KNNCodecService to inject the right KNNCodec version @@ -27,6 +27,6 @@ public KNNCodecService(CodecServiceConfig codecServiceConfig) { */ @Override public Codec codec(String name) { - return new KNN87Codec(super.codec(name)); + return new KNN91Codec(super.codec(name)); } } diff --git a/src/main/java/org/opensearch/knn/indices/ModelDao.java b/src/main/java/org/opensearch/knn/indices/ModelDao.java index 385b30225..f54c8bc34 100644 --- a/src/main/java/org/opensearch/knn/indices/ModelDao.java +++ b/src/main/java/org/opensearch/knn/indices/ModelDao.java @@ -39,9 +39,14 @@ import org.opensearch.cluster.health.ClusterIndexHealth; import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.Strings; +import org.opensearch.common.bytes.BytesArray; import org.opensearch.common.settings.Settings; +import org.opensearch.common.xcontent.XContentHelper; import org.opensearch.common.xcontent.XContentType; +import org.opensearch.common.xcontent.json.JsonXContent; import org.opensearch.index.IndexNotFoundException; +import org.opensearch.index.mapper.MapperService; import org.opensearch.knn.common.KNNConstants; import org.opensearch.knn.plugin.transport.DeleteModelResponse; import org.opensearch.knn.plugin.transport.GetModelResponse; @@ -199,8 +204,10 @@ public void create(ActionListener actionListener) throws IO if (isCreated()) { return; } - - CreateIndexRequest request = new CreateIndexRequest(MODEL_INDEX_NAME).mapping("_doc", getMapping(), XContentType.JSON) + String mapping = Strings.toString( + JsonXContent.contentBuilder().startObject().startObject(MapperService.SINGLE_MAPPING_NAME).endObject().endObject() + ); + CreateIndexRequest request = new CreateIndexRequest(MODEL_INDEX_NAME).mapping(mapping) .settings( Settings.builder() .put("index.hidden", true) diff --git a/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec b/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec index 8e64afa08..98f8a6a13 100644 --- a/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec +++ b/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec @@ -1,4 +1,5 @@ org.opensearch.knn.index.codec.KNN80Codec.KNN80Codec org.opensearch.knn.index.codec.KNN84Codec.KNN84Codec org.opensearch.knn.index.codec.KNN86Codec.KNN86Codec -org.opensearch.knn.index.codec.KNN87Codec.KNN87Codec \ No newline at end of file +org.opensearch.knn.index.codec.KNN87Codec.KNN87Codec +org.opensearch.knn.index.codec.KNN91Codec.KNN91Codec \ No newline at end of file diff --git a/src/test/java/org/opensearch/knn/TestUtils.java b/src/test/java/org/opensearch/knn/TestUtils.java index f4968a255..ce09394d2 100644 --- a/src/test/java/org/opensearch/knn/TestUtils.java +++ b/src/test/java/org/opensearch/knn/TestUtils.java @@ -21,7 +21,6 @@ import java.io.IOException; import org.opensearch.knn.index.SpaceType; import org.opensearch.knn.plugin.script.KNNScoringUtil; -import org.opensearch.knn.plugin.stats.suppliers.ModelIndexStatusSupplier; import java.util.Comparator; import java.util.Random; import java.util.Set; @@ -30,7 +29,8 @@ import java.util.List; import java.util.HashSet; import java.util.Map; -import static org.apache.lucene.util.LuceneTestCase.random; + +import static org.apache.lucene.tests.util.LuceneTestCase.random; class DistVector { public float dist; diff --git a/src/test/java/org/opensearch/knn/index/KNNCreateIndexFromModelTests.java b/src/test/java/org/opensearch/knn/index/KNNCreateIndexFromModelTests.java index e1feb9f18..3c7905126 100644 --- a/src/test/java/org/opensearch/knn/index/KNNCreateIndexFromModelTests.java +++ b/src/test/java/org/opensearch/knn/index/KNNCreateIndexFromModelTests.java @@ -14,7 +14,10 @@ import com.google.common.collect.ImmutableMap; import org.opensearch.action.ActionListener; import org.opensearch.action.admin.indices.create.CreateIndexRequestBuilder; +import org.opensearch.common.Strings; import org.opensearch.common.settings.Settings; +import org.opensearch.common.xcontent.XContentFactory; +import org.opensearch.common.xcontent.json.JsonXContent; import org.opensearch.knn.KNNSingleNodeTestCase; import org.opensearch.knn.jni.JNIService; import org.opensearch.knn.index.util.KNNEngine; @@ -64,6 +67,17 @@ public void testCreateIndexFromModel() throws IOException, InterruptedException String indexName = "test-index"; String fieldName = "test-field"; + final String mapping = Strings.toString( + XContentFactory.jsonBuilder() + .startObject() + .startObject("properties") + .startObject(fieldName) + .field("type", "knn_vector") + .field("model_id", modelId) + .endObject() + .endObject() + .endObject()); + modelDao.put(model, ActionListener.wrap(indexResponse -> { CreateIndexRequestBuilder createIndexRequestBuilder = client().admin().indices().prepareCreate(indexName) .setSettings(Settings.builder() @@ -71,16 +85,7 @@ public void testCreateIndexFromModel() throws IOException, InterruptedException .put("number_of_replicas", 0) .put("index.knn", true) .build() - ).addMapping( - "_doc", ImmutableMap.of( - "properties", ImmutableMap.of( - fieldName, ImmutableMap.of( - "type", "knn_vector", - "model_id", modelId - ) - ) - ) - ); + ).setMapping(mapping); client().admin().indices().create(createIndexRequestBuilder.request(), ActionListener.wrap( diff --git a/src/test/java/org/opensearch/knn/index/KNNVectorDVLeafFieldDataTests.java b/src/test/java/org/opensearch/knn/index/KNNVectorDVLeafFieldDataTests.java index 87d00e554..e1d1889f4 100644 --- a/src/test/java/org/opensearch/knn/index/KNNVectorDVLeafFieldDataTests.java +++ b/src/test/java/org/opensearch/knn/index/KNNVectorDVLeafFieldDataTests.java @@ -6,7 +6,7 @@ package org.opensearch.knn.index; import org.opensearch.knn.KNNTestCase; -import org.apache.lucene.analysis.MockAnalyzer; +import org.apache.lucene.tests.analysis.MockAnalyzer; import org.apache.lucene.document.BinaryDocValuesField; import org.apache.lucene.document.Document; import org.apache.lucene.document.FieldType; diff --git a/src/test/java/org/opensearch/knn/index/KNNVectorIndexFieldDataTests.java b/src/test/java/org/opensearch/knn/index/KNNVectorIndexFieldDataTests.java index 54435db99..346052618 100644 --- a/src/test/java/org/opensearch/knn/index/KNNVectorIndexFieldDataTests.java +++ b/src/test/java/org/opensearch/knn/index/KNNVectorIndexFieldDataTests.java @@ -6,7 +6,7 @@ package org.opensearch.knn.index; import org.opensearch.knn.KNNTestCase; -import org.apache.lucene.analysis.MockAnalyzer; +import org.apache.lucene.tests.analysis.MockAnalyzer; import org.apache.lucene.document.Document; import org.apache.lucene.index.DirectoryReader; import org.apache.lucene.index.IndexWriter; diff --git a/src/test/java/org/opensearch/knn/index/KNNVectorScriptDocValuesTests.java b/src/test/java/org/opensearch/knn/index/KNNVectorScriptDocValuesTests.java index 8883bf4dd..e54e19141 100644 --- a/src/test/java/org/opensearch/knn/index/KNNVectorScriptDocValuesTests.java +++ b/src/test/java/org/opensearch/knn/index/KNNVectorScriptDocValuesTests.java @@ -6,7 +6,7 @@ package org.opensearch.knn.index; import org.opensearch.knn.KNNTestCase; -import org.apache.lucene.analysis.MockAnalyzer; +import org.apache.lucene.tests.analysis.MockAnalyzer; import org.apache.lucene.document.BinaryDocValuesField; import org.apache.lucene.document.Document; import org.apache.lucene.document.FieldType; diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN87Codec/KNN87CodecTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN87Codec/KNN87CodecTests.java index b7f909ec1..48c6e6c0d 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNN87Codec/KNN87CodecTests.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNN87Codec/KNN87CodecTests.java @@ -5,11 +5,13 @@ package org.opensearch.knn.index.codec.KNN87Codec; +import org.junit.Ignore; import org.opensearch.knn.index.codec.KNNCodecTestCase; import java.io.IOException; import java.util.concurrent.ExecutionException; +@Ignore public class KNN87CodecTests extends KNNCodecTestCase { public void testMultiFieldsKnnIndex() throws Exception { diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91BinaryDocValuesTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91BinaryDocValuesTests.java new file mode 100644 index 000000000..a548b4544 --- /dev/null +++ b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91BinaryDocValuesTests.java @@ -0,0 +1,69 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.index.codec.KNN91Codec; + +import com.google.common.collect.ImmutableList; +import org.apache.lucene.index.BinaryDocValues; +import org.apache.lucene.index.DocIDMerger; +import org.apache.lucene.index.MergeState; +import org.opensearch.knn.KNNTestCase; +import org.opensearch.knn.index.codec.KNNCodecTestUtil; +import org.opensearch.knn.index.codec.util.BinaryDocValuesSub; + +import java.io.IOException; + +public class KNN91BinaryDocValuesTests extends KNNTestCase { + + public void testDocId() { + KNN91BinaryDocValues knn91BinaryDocValues = new KNN91BinaryDocValues(null); + assertEquals(-1, knn91BinaryDocValues.docID()); + } + + public void testNextDoc() throws IOException { + final int expectedDoc = 12; + + BinaryDocValuesSub sub = new BinaryDocValuesSub(new MergeState.DocMap() { + @Override + public int get(int docID) { + return expectedDoc; + } + }, new KNNCodecTestUtil.ConstantVectorBinaryDocValues(10, 128, 1.0f)); + + DocIDMerger docIDMerger = DocIDMerger.of(ImmutableList.of(sub), false); + KNN91BinaryDocValues knn91BinaryDocValues = new KNN91BinaryDocValues(docIDMerger); + assertEquals(expectedDoc, knn91BinaryDocValues.nextDoc()); + } + + public void testAdvance() { + KNN91BinaryDocValues knn91BinaryDocValues = new KNN91BinaryDocValues(null); + expectThrows(UnsupportedOperationException.class, () -> knn91BinaryDocValues.advance(0)); + } + + public void testAdvanceExact() { + KNN91BinaryDocValues knn91BinaryDocValues = new KNN91BinaryDocValues(null); + expectThrows(UnsupportedOperationException.class, () -> knn91BinaryDocValues.advanceExact(0)); + } + + public void testCost() { + KNN91BinaryDocValues knn91BinaryDocValues = new KNN91BinaryDocValues(null); + expectThrows(UnsupportedOperationException.class, knn91BinaryDocValues::cost); + } + + public void testBinaryValue() throws IOException { + BinaryDocValues binaryDocValues = new KNNCodecTestUtil.ConstantVectorBinaryDocValues(10, 128, 1.0f); + BinaryDocValuesSub sub = new BinaryDocValuesSub(new MergeState.DocMap() { + @Override + public int get(int docID) { + return docID; + } + }, binaryDocValues); + + DocIDMerger docIDMerger = DocIDMerger.of(ImmutableList.of(sub), false); + KNN91BinaryDocValues knn91BinaryDocValues = new KNN91BinaryDocValues(docIDMerger); + knn91BinaryDocValues.nextDoc(); + assertEquals(binaryDocValues.binaryValue(), knn91BinaryDocValues.binaryValue()); + } +} diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CodecTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CodecTests.java new file mode 100644 index 000000000..39dc641dc --- /dev/null +++ b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CodecTests.java @@ -0,0 +1,22 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.index.codec.KNN91Codec; + +import org.opensearch.knn.index.codec.KNNCodecTestCase; + +import java.io.IOException; +import java.util.concurrent.ExecutionException; + +public class KNN91CodecTests extends KNNCodecTestCase { + + public void testMultiFieldsKnnIndex() throws Exception { + testMultiFieldsKnnIndex(new KNN91Codec()); + } + + public void testBuildFromModelTemplate() throws InterruptedException, ExecutionException, IOException { + testBuildFromModelTemplate(new KNN91Codec()); + } +} diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CompoundFormatTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CompoundFormatTests.java new file mode 100644 index 000000000..dd9a358d7 --- /dev/null +++ b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CompoundFormatTests.java @@ -0,0 +1,92 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.index.codec.KNN91Codec; + +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.CompoundDirectory; +import org.apache.lucene.codecs.CompoundFormat; +import org.apache.lucene.index.SegmentInfo; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexOutput; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.opensearch.common.util.set.Sets; +import org.opensearch.knn.KNNTestCase; +import org.opensearch.knn.index.codec.KNNCodecTestUtil; +import org.opensearch.knn.index.util.KNNEngine; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Set; + +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class KNN91CompoundFormatTests extends KNNTestCase { + + private static Directory directory; + private static Codec codec; + + @BeforeClass + public static void setStaticVariables() { + directory = newFSDirectory(createTempDir()); + codec = new KNN91Codec(); + } + + @AfterClass + public static void closeStaticVariables() throws IOException { + directory.close(); + } + + public void testGetCompoundReader() throws IOException { + CompoundDirectory dir = mock(CompoundDirectory.class); + CompoundFormat delegate = mock(CompoundFormat.class); + when(delegate.getCompoundReader(null, null, null)).thenReturn(dir); + KNN91CompoundFormat knn91CompoundFormat = new KNN91CompoundFormat(delegate); + assertEquals(dir, knn91CompoundFormat.getCompoundReader(null, null, null)); + } + + public void testWrite() throws IOException { + // Check that all normal engine files correctly get set to compound extension files after write + String segmentName = "_test"; + + Set segmentFiles = Sets.newHashSet( + String.format("%s_nmslib1%s", segmentName, KNNEngine.NMSLIB.getExtension()), + String.format("%s_nmslib2%s", segmentName, KNNEngine.NMSLIB.getExtension()), + String.format("%s_nmslib3%s", segmentName, KNNEngine.NMSLIB.getExtension()), + String.format("%s_faiss1%s", segmentName, KNNEngine.FAISS.getExtension()), + String.format("%s_faiss2%s", segmentName, KNNEngine.FAISS.getExtension()), + String.format("%s_faiss3%s", segmentName, KNNEngine.FAISS.getExtension()) + ); + + SegmentInfo segmentInfo = KNNCodecTestUtil.SegmentInfoBuilder.builder(directory, segmentName, segmentFiles.size(), codec).build(); + + for (String name : segmentFiles) { + IndexOutput indexOutput = directory.createOutput(name, IOContext.DEFAULT); + indexOutput.close(); + } + segmentInfo.setFiles(segmentFiles); + + CompoundFormat delegate = mock(CompoundFormat.class); + doNothing().when(delegate).write(directory, segmentInfo, IOContext.DEFAULT); + + KNN91CompoundFormat knn91CompoundFormat = new KNN91CompoundFormat(delegate); + knn91CompoundFormat.write(directory, segmentInfo, IOContext.DEFAULT); + + assertTrue(segmentInfo.files().isEmpty()); + + Arrays.stream(directory.listAll()).forEach(filename -> { + try { + directory.deleteFile(filename); + } catch (IOException e) { + fail(String.format("Failed to delete: %s", filename)); + } + }); + } + +} diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesConsumerTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesConsumerTests.java new file mode 100644 index 000000000..994540cf5 --- /dev/null +++ b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesConsumerTests.java @@ -0,0 +1,412 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.index.codec.KNN91Codec; + +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.DocValuesConsumer; +import org.apache.lucene.codecs.DocValuesProducer; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.FieldInfos; +import org.apache.lucene.index.SegmentInfo; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IOContext; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.Strings; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.xcontent.XContentFactory; +import org.opensearch.knn.KNNTestCase; +import org.opensearch.knn.common.KNNConstants; +import org.opensearch.knn.index.KNNMethodContext; +import org.opensearch.knn.index.KNNVectorFieldMapper; +import org.opensearch.knn.index.MethodComponentContext; +import org.opensearch.knn.index.SpaceType; +import org.opensearch.knn.index.codec.KNN87Codec.KNN87Codec; +import org.opensearch.knn.index.codec.KNNCodecTestUtil; +import org.opensearch.knn.index.codec.util.KNNCodecUtil; +import org.opensearch.knn.index.util.KNNEngine; +import org.opensearch.knn.indices.Model; +import org.opensearch.knn.indices.ModelCache; +import org.opensearch.knn.indices.ModelDao; +import org.opensearch.knn.indices.ModelMetadata; +import org.opensearch.knn.indices.ModelState; +import org.opensearch.knn.jni.JNIService; +import org.opensearch.knn.plugin.stats.KNNCounter; + +import java.io.IOException; +import java.util.Map; +import java.util.concurrent.ExecutionException; + +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.opensearch.knn.common.KNNConstants.INDEX_DESCRIPTION_PARAMETER; +import static org.opensearch.knn.common.KNNConstants.METHOD_HNSW; +import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_EF_CONSTRUCTION; +import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_M; +import static org.opensearch.knn.common.KNNConstants.MODEL_ID; +import static org.opensearch.knn.index.KNNSettings.MODEL_CACHE_SIZE_LIMIT_SETTING; +import static org.opensearch.knn.index.codec.KNNCodecTestUtil.RandomVectorDocValuesProducer; +import static org.opensearch.knn.index.codec.KNNCodecTestUtil.assertFileInCorrectLocation; +import static org.opensearch.knn.index.codec.KNNCodecTestUtil.assertLoadableByEngine; +import static org.opensearch.knn.index.codec.KNNCodecTestUtil.assertValidFooter; +import static org.opensearch.knn.index.codec.KNNCodecTestUtil.getRandomVectors; + +public class KNN91DocValuesConsumerTests extends KNNTestCase { + + private static Directory directory; + private static Codec codec; + + @BeforeClass + public static void setStaticVariables() { + directory = newFSDirectory(createTempDir()); + codec = new KNN87Codec(); + } + + @AfterClass + public static void closeStaticVariables() throws IOException { + directory.close(); + } + + public void testAddBinaryField_withKNN() throws IOException { + // Confirm that addKNNBinaryField will get called if the k-NN parameter is true + FieldInfo fieldInfo = KNNCodecTestUtil.FieldInfoBuilder.builder("test-field") + .addAttribute(KNNVectorFieldMapper.KNN_FIELD, "true") + .build(); + DocValuesProducer docValuesProducer = mock(DocValuesProducer.class); + + DocValuesConsumer delegate = mock(DocValuesConsumer.class); + doNothing().when(delegate).addBinaryField(fieldInfo, docValuesProducer); + + final boolean[] called = { false }; + KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(delegate, null) { + + @Override + public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer) { + called[0] = true; + } + }; + + knn91DocValuesConsumer.addBinaryField(fieldInfo, docValuesProducer); + + verify(delegate, times(1)).addBinaryField(fieldInfo, docValuesProducer); + assertTrue(called[0]); + } + + public void testAddBinaryField_withoutKNN() throws IOException { + // Confirm that the KNN91DocValuesConsumer will just call delegate AddBinaryField when k-NN parameter is + // not set + FieldInfo fieldInfo = KNNCodecTestUtil.FieldInfoBuilder.builder("test-field").build(); + DocValuesProducer docValuesProducer = mock(DocValuesProducer.class); + + DocValuesConsumer delegate = mock(DocValuesConsumer.class); + doNothing().when(delegate).addBinaryField(fieldInfo, docValuesProducer); + + final boolean[] called = { false }; + KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(delegate, null) { + + @Override + public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer) { + called[0] = true; + } + }; + + knn91DocValuesConsumer.addBinaryField(fieldInfo, docValuesProducer); + + verify(delegate, times(1)).addBinaryField(fieldInfo, docValuesProducer); + assertFalse(called[0]); + } + + public void testAddKNNBinaryField_noVectors() throws IOException { + // When there are no new vectors, no more graph index requests should be added + RandomVectorDocValuesProducer randomVectorDocValuesProducer = new RandomVectorDocValuesProducer(0, 128); + Long initialGraphIndexRequests = KNNCounter.GRAPH_INDEX_REQUESTS.getCount(); + KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(null, null); + knn91DocValuesConsumer.addKNNBinaryField(null, randomVectorDocValuesProducer); + assertEquals(initialGraphIndexRequests, KNNCounter.GRAPH_INDEX_REQUESTS.getCount()); + } + + public void testAddKNNBinaryField_fromScratch_nmslibCurrent() throws IOException { + // Set information about the segment and the fields + String segmentName = String.format("test_segment%s", randomAlphaOfLength(4)); + int docsInSegment = 100; + String fieldName = String.format("test_field%s", randomAlphaOfLength(4)); + + KNNEngine knnEngine = KNNEngine.NMSLIB; + SpaceType spaceType = SpaceType.COSINESIMIL; + int dimension = 16; + + SegmentInfo segmentInfo = KNNCodecTestUtil.SegmentInfoBuilder.builder(directory, segmentName, docsInSegment, codec).build(); + + KNNMethodContext knnMethodContext = new KNNMethodContext( + knnEngine, + spaceType, + new MethodComponentContext(METHOD_HNSW, ImmutableMap.of(METHOD_PARAMETER_M, 16, METHOD_PARAMETER_EF_CONSTRUCTION, 512)) + ); + + String parameterString = Strings.toString(XContentFactory.jsonBuilder().map(knnEngine.getMethodAsMap(knnMethodContext))); + + FieldInfo[] fieldInfoArray = new FieldInfo[] { + KNNCodecTestUtil.FieldInfoBuilder.builder(fieldName) + .addAttribute(KNNVectorFieldMapper.KNN_FIELD, "true") + .addAttribute(KNNConstants.KNN_ENGINE, knnEngine.getName()) + .addAttribute(KNNConstants.SPACE_TYPE, spaceType.getValue()) + .addAttribute(KNNConstants.PARAMETERS, parameterString) + .build() }; + + FieldInfos fieldInfos = new FieldInfos(fieldInfoArray); + SegmentWriteState state = new SegmentWriteState(null, directory, segmentInfo, fieldInfos, null, IOContext.DEFAULT); + + // Add documents to the field + KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(null, state); + RandomVectorDocValuesProducer randomVectorDocValuesProducer = new RandomVectorDocValuesProducer(docsInSegment, dimension); + knn91DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer); + + // The document should be created in the correct location + String expectedFile = KNNCodecUtil.buildEngineFileName( + segmentName, + knnEngine.getLatestBuildVersion(), + fieldName, + knnEngine.getExtension() + ); + assertFileInCorrectLocation(state, expectedFile); + + // The footer should be valid + assertValidFooter(state.directory, expectedFile); + + // The document should be readable by nmslib + assertLoadableByEngine(state, expectedFile, knnEngine, spaceType, dimension); + } + + public void testAddKNNBinaryField_fromScratch_nmslibLegacy() throws IOException { + // Set information about the segment and the fields + String segmentName = String.format("test_segment%s", randomAlphaOfLength(4)); + int docsInSegment = 100; + String fieldName = String.format("test_field%s", randomAlphaOfLength(4)); + + KNNEngine knnEngine = KNNEngine.NMSLIB; + SpaceType spaceType = SpaceType.COSINESIMIL; + int dimension = 16; + + SegmentInfo segmentInfo = KNNCodecTestUtil.SegmentInfoBuilder.builder(directory, segmentName, docsInSegment, codec).build(); + + FieldInfo[] fieldInfoArray = new FieldInfo[] { + KNNCodecTestUtil.FieldInfoBuilder.builder(fieldName) + .addAttribute(KNNVectorFieldMapper.KNN_FIELD, "true") + .addAttribute(KNNConstants.HNSW_ALGO_EF_CONSTRUCTION, "512") + .addAttribute(KNNConstants.HNSW_ALGO_M, "16") + .addAttribute(KNNConstants.SPACE_TYPE, spaceType.getValue()) + .build() }; + + FieldInfos fieldInfos = new FieldInfos(fieldInfoArray); + SegmentWriteState state = new SegmentWriteState(null, directory, segmentInfo, fieldInfos, null, IOContext.DEFAULT); + + // Add documents to the field + KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(null, state); + RandomVectorDocValuesProducer randomVectorDocValuesProducer = new RandomVectorDocValuesProducer(docsInSegment, dimension); + knn91DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer); + + // The document should be created in the correct location + String expectedFile = KNNCodecUtil.buildEngineFileName( + segmentName, + knnEngine.getLatestBuildVersion(), + fieldName, + knnEngine.getExtension() + ); + assertFileInCorrectLocation(state, expectedFile); + + // The footer should be valid + assertValidFooter(state.directory, expectedFile); + + // The document should be readable by nmslib + assertLoadableByEngine(state, expectedFile, knnEngine, spaceType, dimension); + } + + public void testAddKNNBinaryField_fromScratch_faissCurrent() throws IOException { + String segmentName = String.format("test_segment%s", randomAlphaOfLength(4)); + int docsInSegment = 100; + String fieldName = String.format("test_field%s", randomAlphaOfLength(4)); + + KNNEngine knnEngine = KNNEngine.FAISS; + SpaceType spaceType = SpaceType.INNER_PRODUCT; + int dimension = 16; + + SegmentInfo segmentInfo = KNNCodecTestUtil.SegmentInfoBuilder.builder(directory, segmentName, docsInSegment, codec).build(); + + KNNMethodContext knnMethodContext = new KNNMethodContext( + knnEngine, + spaceType, + new MethodComponentContext(METHOD_HNSW, ImmutableMap.of(METHOD_PARAMETER_M, 16, METHOD_PARAMETER_EF_CONSTRUCTION, 512)) + ); + + String parameterString = Strings.toString(XContentFactory.jsonBuilder().map(knnEngine.getMethodAsMap(knnMethodContext))); + + FieldInfo[] fieldInfoArray = new FieldInfo[] { + KNNCodecTestUtil.FieldInfoBuilder.builder(fieldName) + .addAttribute(KNNVectorFieldMapper.KNN_FIELD, "true") + .addAttribute(KNNConstants.KNN_ENGINE, knnEngine.getName()) + .addAttribute(KNNConstants.SPACE_TYPE, spaceType.getValue()) + .addAttribute(KNNConstants.PARAMETERS, parameterString) + .build() }; + + FieldInfos fieldInfos = new FieldInfos(fieldInfoArray); + SegmentWriteState state = new SegmentWriteState(null, directory, segmentInfo, fieldInfos, null, IOContext.DEFAULT); + + // Add documents to the field + KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(null, state); + RandomVectorDocValuesProducer randomVectorDocValuesProducer = new RandomVectorDocValuesProducer(docsInSegment, dimension); + knn91DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer); + + // The document should be created in the correct location + String expectedFile = KNNCodecUtil.buildEngineFileName( + segmentName, + knnEngine.getLatestBuildVersion(), + fieldName, + knnEngine.getExtension() + ); + assertFileInCorrectLocation(state, expectedFile); + + // The footer should be valid + assertValidFooter(state.directory, expectedFile); + + // The document should be readable by faiss + assertLoadableByEngine(state, expectedFile, knnEngine, spaceType, dimension); + } + + public void testAddKNNBinaryField_fromModel_faiss() throws IOException, ExecutionException, InterruptedException { + // Generate a trained faiss model + KNNEngine knnEngine = KNNEngine.FAISS; + SpaceType spaceType = SpaceType.INNER_PRODUCT; + int dimension = 16; + String modelId = "test-model-id"; + + float[][] trainingData = getRandomVectors(200, dimension); + long trainingPtr = JNIService.transferVectors(0, trainingData); + + Map parameters = ImmutableMap.of( + INDEX_DESCRIPTION_PARAMETER, + "IVF4,Flat", + KNNConstants.SPACE_TYPE, + SpaceType.L2.getValue() + ); + + byte[] modelBytes = JNIService.trainIndex(parameters, dimension, trainingPtr, knnEngine.getName()); + Model model = new Model( + new ModelMetadata(knnEngine, spaceType, dimension, ModelState.CREATED, "timestamp", "Empty description", ""), + modelBytes, + modelId + ); + JNIService.freeVectors(trainingPtr); + + // Setup the model cache to return the correct model + ModelDao modelDao = mock(ModelDao.class); + when(modelDao.get(modelId)).thenReturn(model); + ClusterService clusterService = mock(ClusterService.class); + when(clusterService.getSettings()).thenReturn(Settings.EMPTY); + + ClusterSettings clusterSettings = new ClusterSettings( + Settings.builder().put(MODEL_CACHE_SIZE_LIMIT_SETTING.getKey(), "10kb").build(), + ImmutableSet.of(MODEL_CACHE_SIZE_LIMIT_SETTING) + ); + + when(clusterService.getClusterSettings()).thenReturn(clusterSettings); + ModelCache.initialize(modelDao, clusterService); + + // Build the segment and field info + String segmentName = String.format("test_segment%s", randomAlphaOfLength(4)); + int docsInSegment = 100; + String fieldName = String.format("test_field%s", randomAlphaOfLength(4)); + + SegmentInfo segmentInfo = KNNCodecTestUtil.SegmentInfoBuilder.builder(directory, segmentName, docsInSegment, codec).build(); + + FieldInfo[] fieldInfoArray = new FieldInfo[] { + KNNCodecTestUtil.FieldInfoBuilder.builder(fieldName) + .addAttribute(KNNVectorFieldMapper.KNN_FIELD, "true") + .addAttribute(MODEL_ID, modelId) + .build() }; + + FieldInfos fieldInfos = new FieldInfos(fieldInfoArray); + SegmentWriteState state = new SegmentWriteState(null, directory, segmentInfo, fieldInfos, null, IOContext.DEFAULT); + + // Add documents to the field + KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(null, state); + RandomVectorDocValuesProducer randomVectorDocValuesProducer = new RandomVectorDocValuesProducer(docsInSegment, dimension); + knn91DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer); + + // The document should be created in the correct location + String expectedFile = KNNCodecUtil.buildEngineFileName( + segmentName, + knnEngine.getLatestBuildVersion(), + fieldName, + knnEngine.getExtension() + ); + assertFileInCorrectLocation(state, expectedFile); + + // The footer should be valid + assertValidFooter(state.directory, expectedFile); + + // The document should be readable by faiss + assertLoadableByEngine(state, expectedFile, knnEngine, spaceType, dimension); + } + + public void testMerge_exception() throws IOException { + KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(null, null); + expectThrows(RuntimeException.class, () -> knn91DocValuesConsumer.merge(null)); + } + + public void testAddSortedSetField() throws IOException { + // Verify that the delegate will be called + DocValuesConsumer delegate = mock(DocValuesConsumer.class); + doNothing().when(delegate).addSortedSetField(null, null); + KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(delegate, null); + knn91DocValuesConsumer.addSortedSetField(null, null); + verify(delegate, times(1)).addSortedSetField(null, null); + } + + public void testAddSortedNumericField() throws IOException { + // Verify that the delegate will be called + DocValuesConsumer delegate = mock(DocValuesConsumer.class); + doNothing().when(delegate).addSortedNumericField(null, null); + KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(delegate, null); + knn91DocValuesConsumer.addSortedNumericField(null, null); + verify(delegate, times(1)).addSortedNumericField(null, null); + } + + public void testAddSortedField() throws IOException { + // Verify that the delegate will be called + DocValuesConsumer delegate = mock(DocValuesConsumer.class); + doNothing().when(delegate).addSortedField(null, null); + KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(delegate, null); + knn91DocValuesConsumer.addSortedField(null, null); + verify(delegate, times(1)).addSortedField(null, null); + } + + public void testAddNumericField() throws IOException { + // Verify that the delegate will be called + DocValuesConsumer delegate = mock(DocValuesConsumer.class); + doNothing().when(delegate).addNumericField(null, null); + KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(delegate, null); + knn91DocValuesConsumer.addNumericField(null, null); + verify(delegate, times(1)).addNumericField(null, null); + } + + public void testClose() throws IOException { + // Verify that the delegate will be called + DocValuesConsumer delegate = mock(DocValuesConsumer.class); + doNothing().when(delegate).close(); + KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(delegate, null); + knn91DocValuesConsumer.close(); + verify(delegate, times(1)).close(); + } + +} diff --git a/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java b/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java index 3454c7d52..dd7b82cf3 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java @@ -11,6 +11,7 @@ import org.opensearch.common.settings.ClusterSettings; import org.opensearch.knn.KNNTestCase; import org.opensearch.knn.common.KNNConstants; +import org.opensearch.knn.index.codec.KNN91Codec.KNN91Codec; import org.opensearch.knn.jni.JNIService; import org.opensearch.knn.index.KNNQuery; import org.opensearch.knn.index.KNNSettings; @@ -24,7 +25,7 @@ import org.apache.lucene.document.FieldType; import org.apache.lucene.index.IndexReader; import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.index.RandomIndexWriter; +import org.apache.lucene.tests.index.RandomIndexWriter; import org.apache.lucene.index.SerialMergeScheduler; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.store.Directory; @@ -61,6 +62,7 @@ */ public class KNNCodecTestCase extends KNNTestCase { + private static final KNN91Codec ACTUAL_CODEC = new KNN91Codec(); private static FieldType sampleFieldType; static { sampleFieldType = new FieldType(KNNVectorFieldMapper.Defaults.FIELD_TYPE); @@ -107,7 +109,7 @@ public void testMultiFieldsKnnIndex(Codec codec) throws Exception { */ IndexWriterConfig iwc1 = newIndexWriterConfig(); iwc1.setMergeScheduler(new SerialMergeScheduler()); - iwc1.setCodec(new KNN87Codec()); + iwc1.setCodec(ACTUAL_CODEC); writer = new RandomIndexWriter(random(), dir, iwc1); float[] array1 = { 6.0f, 14.0f }; VectorField vectorField1 = new VectorField("my_vector", array1, sampleFieldType); diff --git a/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestUtil.java b/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestUtil.java index de335a115..153fd2faf 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestUtil.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestUtil.java @@ -20,6 +20,7 @@ import org.apache.lucene.index.SortedDocValues; import org.apache.lucene.index.SortedNumericDocValues; import org.apache.lucene.index.SortedSetDocValues; +import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.search.Sort; import org.apache.lucene.store.ChecksumIndexInput; import org.apache.lucene.store.Directory; @@ -145,6 +146,8 @@ public static class FieldInfoBuilder { private int pointDimensionCount; private int pointIndexDimensionCount; private int pointNumBytes; + private int vectorDimension; + private VectorSimilarityFunction vectorSimilarityFunction; private boolean softDeletes; public static FieldInfoBuilder builder(String fieldName) { @@ -164,6 +167,8 @@ private FieldInfoBuilder(String fieldName) { this.pointDimensionCount = 0; this.pointIndexDimensionCount = 0; this.pointNumBytes = 0; + this.vectorDimension = 0; + this.vectorSimilarityFunction = VectorSimilarityFunction.EUCLIDEAN; this.softDeletes = false; } @@ -222,6 +227,16 @@ public FieldInfoBuilder pointNumBytes(int pointNumBytes) { return this; } + public FieldInfoBuilder vectorDimension(int vectorDimension) { + this.vectorDimension = vectorDimension; + return this; + } + + public FieldInfoBuilder vectorSimilarityFunction(VectorSimilarityFunction vectorSimilarityFunction) { + this.vectorSimilarityFunction = vectorSimilarityFunction; + return this; + } + public FieldInfoBuilder softDeletes(boolean softDeletes) { this.softDeletes = softDeletes; return this; @@ -241,6 +256,8 @@ public FieldInfo build() { pointDimensionCount, pointIndexDimensionCount, pointNumBytes, + vectorDimension, + vectorSimilarityFunction, softDeletes ); } @@ -364,11 +381,6 @@ public void checkIntegrity() { public void close() throws IOException { } - - @Override - public long ramBytesUsed() { - return 0; - } } public static void assertFileInCorrectLocation(SegmentWriteState state, String expectedFile) throws IOException { diff --git a/src/test/java/org/opensearch/knn/plugin/script/KNNScoringUtilTests.java b/src/test/java/org/opensearch/knn/plugin/script/KNNScoringUtilTests.java index 08cede77c..0375b6b48 100644 --- a/src/test/java/org/opensearch/knn/plugin/script/KNNScoringUtilTests.java +++ b/src/test/java/org/opensearch/knn/plugin/script/KNNScoringUtilTests.java @@ -8,7 +8,7 @@ import org.opensearch.knn.KNNTestCase; import org.opensearch.knn.index.KNNVectorScriptDocValues; import org.opensearch.knn.index.VectorField; -import org.apache.lucene.analysis.MockAnalyzer; +import org.apache.lucene.tests.analysis.MockAnalyzer; import org.apache.lucene.document.BinaryDocValuesField; import org.apache.lucene.document.Document; import org.apache.lucene.document.FieldType; From 05c3989423aa2b1288cfccc347dc358e3ac5e650 Mon Sep 17 00:00:00 2001 From: Martin Gaievski Date: Mon, 21 Mar 2022 14:28:50 -0700 Subject: [PATCH 2/9] Apply spotless Signed-off-by: Martin Gaievski --- .../org/opensearch/knn/index/KNNQuery.java | 7 +- .../org/opensearch/knn/index/KNNWeight.java | 212 +++++++++--------- .../index/codec/KNN80Codec/KNN80Codec.java | 3 +- .../index/codec/KNN84Codec/KNN84Codec.java | 5 +- .../index/codec/KNN86Codec/KNN86Codec.java | 6 +- .../index/codec/KNN91Codec/KNN91Codec.java | 1 + .../org/opensearch/knn/indices/ModelDao.java | 5 +- .../java/org/opensearch/knn/TestUtils.java | 21 +- .../index/KNNCreateIndexFromModelTests.java | 71 +++--- .../index/KNNVectorDVLeafFieldDataTests.java | 19 +- .../index/KNNVectorIndexFieldDataTests.java | 9 +- .../index/KNNVectorScriptDocValuesTests.java | 17 +- .../knn/index/codec/KNNCodecTestCase.java | 1 - .../plugin/script/KNNScoringUtilTests.java | 58 +++-- 14 files changed, 216 insertions(+), 219 deletions(-) diff --git a/src/main/java/org/opensearch/knn/index/KNNQuery.java b/src/main/java/org/opensearch/knn/index/KNNQuery.java index 76709ae7e..631b36d7a 100644 --- a/src/main/java/org/opensearch/knn/index/KNNQuery.java +++ b/src/main/java/org/opensearch/knn/index/KNNQuery.java @@ -42,7 +42,9 @@ public int getK() { return this.k; } - public String getIndexName() { return this.indexName; } + public String getIndexName() { + return this.indexName; + } /** * Constructs Weight implementation for this query @@ -77,8 +79,7 @@ public int hashCode() { @Override public boolean equals(Object other) { - return sameClassAs(other) && - equalsTo(getClass().cast(other)); + return sameClassAs(other) && equalsTo(getClass().cast(other)); } private boolean equalsTo(KNNQuery other) { diff --git a/src/main/java/org/opensearch/knn/index/KNNWeight.java b/src/main/java/org/opensearch/knn/index/KNNWeight.java index 3f777c970..7defb1eeb 100644 --- a/src/main/java/org/opensearch/knn/index/KNNWeight.java +++ b/src/main/java/org/opensearch/knn/index/KNNWeight.java @@ -5,7 +5,6 @@ package org.opensearch.knn.index; -import com.google.common.collect.ImmutableMap; import org.opensearch.knn.common.KNNConstants; import org.opensearch.knn.jni.JNIService; import org.opensearch.knn.index.memory.NativeMemoryAllocation; @@ -19,7 +18,6 @@ import org.apache.lucene.index.FilterLeafReader; import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.SegmentReader; -import org.apache.lucene.index.Term; import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.search.Explanation; import org.apache.lucene.search.Scorer; @@ -36,10 +34,8 @@ import java.nio.file.Path; import java.util.Arrays; import java.util.Collections; -import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.Set; import java.util.concurrent.ExecutionException; import java.util.stream.Collectors; @@ -79,108 +75,116 @@ public Explanation explain(LeafReaderContext context, int doc) { @Override public Scorer scorer(LeafReaderContext context) throws IOException { - SegmentReader reader = (SegmentReader) FilterLeafReader.unwrap(context.reader()); - String directory = ((FSDirectory) FilterDirectory.unwrap(reader.directory())).getDirectory().toString(); - - FieldInfo fieldInfo = reader.getFieldInfos().fieldInfo(knnQuery.getField()); - - if (fieldInfo == null) { - logger.debug("[KNN] Field info not found for {}:{}", knnQuery.getField(), - reader.getSegmentName()); - return null; - } - - KNNEngine knnEngine; - SpaceType spaceType; - - // Check if a modelId exists. If so, the space type and engine will need to be picked up from the model's - // metadata. - String modelId = fieldInfo.getAttribute(MODEL_ID); - if (modelId != null) { - ModelMetadata modelMetadata = modelDao.getMetadata(modelId); - if (modelMetadata == null) { - throw new RuntimeException("Model \"" + modelId + "\" does not exist."); - } - - knnEngine = modelMetadata.getKnnEngine(); - spaceType = modelMetadata.getSpaceType(); - } else { - String engineName = fieldInfo.attributes().getOrDefault(KNN_ENGINE, KNNEngine.NMSLIB.getName()); - knnEngine = KNNEngine.getEngine(engineName); - String spaceTypeName = fieldInfo.attributes().getOrDefault(SPACE_TYPE, SpaceType.L2.getValue()); - spaceType = SpaceType.getSpace(spaceTypeName); + SegmentReader reader = (SegmentReader) FilterLeafReader.unwrap(context.reader()); + String directory = ((FSDirectory) FilterDirectory.unwrap(reader.directory())).getDirectory().toString(); + + FieldInfo fieldInfo = reader.getFieldInfos().fieldInfo(knnQuery.getField()); + + if (fieldInfo == null) { + logger.debug("[KNN] Field info not found for {}:{}", knnQuery.getField(), reader.getSegmentName()); + return null; + } + + KNNEngine knnEngine; + SpaceType spaceType; + + // Check if a modelId exists. If so, the space type and engine will need to be picked up from the model's + // metadata. + String modelId = fieldInfo.getAttribute(MODEL_ID); + if (modelId != null) { + ModelMetadata modelMetadata = modelDao.getMetadata(modelId); + if (modelMetadata == null) { + throw new RuntimeException("Model \"" + modelId + "\" does not exist."); } - /* - * In case of compound file, extension would be + c otherwise - */ - String engineExtension = reader.getSegmentInfo().info.getUseCompoundFile() - ? knnEngine.getExtension() + KNNConstants.COMPOUND_EXTENSION : knnEngine.getExtension(); - String engineSuffix = knnQuery.getField() + engineExtension; - List engineFiles = reader.getSegmentInfo().files().stream() - .filter(fileName -> fileName.endsWith(engineSuffix)) - .collect(Collectors.toList()); - - if(engineFiles.isEmpty()) { - logger.debug("[KNN] No engine index found for field {} for segment {}", - knnQuery.getField(), reader.getSegmentName()); - return null; + knnEngine = modelMetadata.getKnnEngine(); + spaceType = modelMetadata.getSpaceType(); + } else { + String engineName = fieldInfo.attributes().getOrDefault(KNN_ENGINE, KNNEngine.NMSLIB.getName()); + knnEngine = KNNEngine.getEngine(engineName); + String spaceTypeName = fieldInfo.attributes().getOrDefault(SPACE_TYPE, SpaceType.L2.getValue()); + spaceType = SpaceType.getSpace(spaceTypeName); + } + + /* + * In case of compound file, extension would be + c otherwise + */ + String engineExtension = reader.getSegmentInfo().info.getUseCompoundFile() + ? knnEngine.getExtension() + KNNConstants.COMPOUND_EXTENSION + : knnEngine.getExtension(); + String engineSuffix = knnQuery.getField() + engineExtension; + List engineFiles = reader.getSegmentInfo() + .files() + .stream() + .filter(fileName -> fileName.endsWith(engineSuffix)) + .collect(Collectors.toList()); + + if (engineFiles.isEmpty()) { + logger.debug("[KNN] No engine index found for field {} for segment {}", knnQuery.getField(), reader.getSegmentName()); + return null; + } + + Path indexPath = PathUtils.get(directory, engineFiles.get(0)); + final KNNQueryResult[] results; + KNNCounter.GRAPH_QUERY_REQUESTS.increment(); + + // We need to first get index allocation + NativeMemoryAllocation indexAllocation; + try { + indexAllocation = nativeMemoryCacheManager.get( + new NativeMemoryEntryContext.IndexEntryContext( + indexPath.toString(), + NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance(), + getParametersAtLoading(spaceType, knnEngine, knnQuery.getIndexName()), + knnQuery.getIndexName() + ), + true + ); + } catch (ExecutionException e) { + GRAPH_QUERY_ERRORS.increment(); + throw new RuntimeException(e); + } + + // Now that we have the allocation, we need to readLock it + indexAllocation.readLock(); + + try { + if (indexAllocation.isClosed()) { + throw new RuntimeException("Index has already been closed"); } - Path indexPath = PathUtils.get(directory, engineFiles.get(0)); - final KNNQueryResult[] results; - KNNCounter.GRAPH_QUERY_REQUESTS.increment(); - - // We need to first get index allocation - NativeMemoryAllocation indexAllocation; - try { - indexAllocation = nativeMemoryCacheManager.get( - new NativeMemoryEntryContext.IndexEntryContext( - indexPath.toString(), - NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance(), - getParametersAtLoading(spaceType, knnEngine, knnQuery.getIndexName()), - knnQuery.getIndexName() - ), true); - } catch (ExecutionException e) { - GRAPH_QUERY_ERRORS.increment(); - throw new RuntimeException(e); - } - - // Now that we have the allocation, we need to readLock it - indexAllocation.readLock(); - - try { - if (indexAllocation.isClosed()) { - throw new RuntimeException("Index has already been closed"); - } - - results = JNIService.queryIndex(indexAllocation.getMemoryAddress(), knnQuery.getQueryVector(), knnQuery.getK(), knnEngine.getName()); - } catch (Exception e) { - GRAPH_QUERY_ERRORS.increment(); - throw new RuntimeException(e); - } finally { - indexAllocation.readUnlock(); - } - - /* - * Scores represent the distance of the documents with respect to given query vector. - * Lesser the score, the closer the document is to the query vector. - * Since by default results are retrieved in the descending order of scores, to get the nearest - * neighbors we are inverting the scores. - */ - if (results.length == 0) { - logger.debug("[KNN] Query yielded 0 results"); - return null; - } - - Map scores = Arrays.stream(results).collect( - Collectors.toMap(KNNQueryResult::getId, result -> knnEngine.score(result.getScore(), spaceType))); - int maxDoc = Collections.max(scores.keySet()) + 1; - DocIdSetBuilder docIdSetBuilder = new DocIdSetBuilder(maxDoc); - DocIdSetBuilder.BulkAdder setAdder = docIdSetBuilder.grow(maxDoc); - Arrays.stream(results).forEach(result -> setAdder.add(result.getId())); - DocIdSetIterator docIdSetIter = docIdSetBuilder.build().iterator(); - return new KNNScorer(this, docIdSetIter, scores, boost); + results = JNIService.queryIndex( + indexAllocation.getMemoryAddress(), + knnQuery.getQueryVector(), + knnQuery.getK(), + knnEngine.getName() + ); + } catch (Exception e) { + GRAPH_QUERY_ERRORS.increment(); + throw new RuntimeException(e); + } finally { + indexAllocation.readUnlock(); + } + + /* + * Scores represent the distance of the documents with respect to given query vector. + * Lesser the score, the closer the document is to the query vector. + * Since by default results are retrieved in the descending order of scores, to get the nearest + * neighbors we are inverting the scores. + */ + if (results.length == 0) { + logger.debug("[KNN] Query yielded 0 results"); + return null; + } + + Map scores = Arrays.stream(results) + .collect(Collectors.toMap(KNNQueryResult::getId, result -> knnEngine.score(result.getScore(), spaceType))); + int maxDoc = Collections.max(scores.keySet()) + 1; + DocIdSetBuilder docIdSetBuilder = new DocIdSetBuilder(maxDoc); + DocIdSetBuilder.BulkAdder setAdder = docIdSetBuilder.grow(maxDoc); + Arrays.stream(results).forEach(result -> setAdder.add(result.getId())); + DocIdSetIterator docIdSetIter = docIdSetBuilder.build().iterator(); + return new KNNScorer(this, docIdSetIter, scores, boost); } @Override @@ -189,9 +193,7 @@ public boolean isCacheable(LeafReaderContext context) { } public static float normalizeScore(float score) { - if (score >= 0) - return 1 / (1 + score); + if (score >= 0) return 1 / (1 + score); return -score + 1; } } - diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80Codec.java b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80Codec.java index 5eb36e748..0064f49fe 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80Codec.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80Codec.java @@ -53,8 +53,7 @@ public DocValuesFormat getDocValuesFormatForField(String field) { * This function returns the Lucene80 Codec. */ public Codec getDelegatee() { - if (lucene80Codec == null) - lucene80Codec = Codec.forName(LUCENE_80); + if (lucene80Codec == null) lucene80Codec = Codec.forName(LUCENE_80); return lucene80Codec; } diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN84Codec/KNN84Codec.java b/src/main/java/org/opensearch/knn/index/codec/KNN84Codec/KNN84Codec.java index c4ae7ab2e..b55365c09 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN84Codec/KNN84Codec.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN84Codec/KNN84Codec.java @@ -43,7 +43,7 @@ public KNN84Codec() { super(KNN_84); // Note that DocValuesFormat can use old Codec's DocValuesFormat. For instance Lucene84 uses Lucene80 // DocValuesFormat. Refer to defaultDVFormat in LuceneXXCodec.java to find out which version it uses - this.docValuesFormat = new KNN80DocValuesFormat(); + this.docValuesFormat = new KNN80DocValuesFormat(); this.perFieldDocValuesFormat = new PerFieldDocValuesFormat() { @Override public DocValuesFormat getDocValuesFormatForField(String field) { @@ -57,8 +57,7 @@ public DocValuesFormat getDocValuesFormatForField(String field) { * This function returns the Lucene84 Codec. */ public Codec getDelegatee() { - if (lucene84Codec == null) - lucene84Codec = Codec.forName(LUCENE_84); + if (lucene84Codec == null) lucene84Codec = Codec.forName(LUCENE_84); return lucene84Codec; } diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN86Codec/KNN86Codec.java b/src/main/java/org/opensearch/knn/index/codec/KNN86Codec/KNN86Codec.java index 154f44f1f..a3b34559a 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN86Codec/KNN86Codec.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN86Codec/KNN86Codec.java @@ -44,7 +44,7 @@ public KNN86Codec() { super(KNN_86); // Note that DocValuesFormat can use old Codec's DocValuesFormat. For instance Lucene84 uses Lucene80 // DocValuesFormat. Refer to defaultDVFormat in LuceneXXCodec.java to find out which version it uses - this.docValuesFormat = new KNN80DocValuesFormat(); + this.docValuesFormat = new KNN80DocValuesFormat(); this.perFieldDocValuesFormat = new PerFieldDocValuesFormat() { @Override public DocValuesFormat getDocValuesFormatForField(String field) { @@ -58,8 +58,7 @@ public DocValuesFormat getDocValuesFormatForField(String field) { * This function returns the Lucene84 Codec. */ public Codec getDelegatee() { - if (lucene86Codec == null) - lucene86Codec = Codec.forName(LUCENE_86); + if (lucene86Codec == null) lucene86Codec = Codec.forName(LUCENE_86); return lucene86Codec; } @@ -74,7 +73,6 @@ public DocValuesFormat docValuesFormat() { * approach of manually overriding. */ - public void setPostingsFormat(PostingsFormat postingsFormat) { this.postingsFormat = postingsFormat; } diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91Codec.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91Codec.java index 1d3bb9580..e63585573 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91Codec.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91Codec.java @@ -35,6 +35,7 @@ public final class KNN91Codec extends FilterCodec { public KNN91Codec() { this(new Lucene91Codec()); } + /** * Constructor that takes a Codec delegate to delegate all methods this code does not implement to. * diff --git a/src/main/java/org/opensearch/knn/indices/ModelDao.java b/src/main/java/org/opensearch/knn/indices/ModelDao.java index f54c8bc34..6ab1b2cca 100644 --- a/src/main/java/org/opensearch/knn/indices/ModelDao.java +++ b/src/main/java/org/opensearch/knn/indices/ModelDao.java @@ -40,10 +40,7 @@ import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.Strings; -import org.opensearch.common.bytes.BytesArray; import org.opensearch.common.settings.Settings; -import org.opensearch.common.xcontent.XContentHelper; -import org.opensearch.common.xcontent.XContentType; import org.opensearch.common.xcontent.json.JsonXContent; import org.opensearch.index.IndexNotFoundException; import org.opensearch.index.mapper.MapperService; @@ -205,7 +202,7 @@ public void create(ActionListener actionListener) throws IO return; } String mapping = Strings.toString( - JsonXContent.contentBuilder().startObject().startObject(MapperService.SINGLE_MAPPING_NAME).endObject().endObject() + JsonXContent.contentBuilder().startObject().startObject(MapperService.SINGLE_MAPPING_NAME).endObject().endObject() ); CreateIndexRequest request = new CreateIndexRequest(MODEL_INDEX_NAME).mapping(mapping) .settings( diff --git a/src/test/java/org/opensearch/knn/TestUtils.java b/src/test/java/org/opensearch/knn/TestUtils.java index ce09394d2..03b50280f 100644 --- a/src/test/java/org/opensearch/knn/TestUtils.java +++ b/src/test/java/org/opensearch/knn/TestUtils.java @@ -36,7 +36,7 @@ class DistVector { public float dist; public String docID; - public DistVector (float dist, String docID) { + public DistVector(float dist, String docID) { this.dist = dist; this.docID = docID; } @@ -117,10 +117,10 @@ public static List> computeGroundTruthValues(float[][] indexVectors, } if (pq.size() < k) { - pq.add(new DistVector(dist, String.valueOf(j+1))); + pq.add(new DistVector(dist, String.valueOf(j + 1))); } else if (pq.peek().getDist() > dist) { pq.poll(); - pq.add(new DistVector(dist, String.valueOf(j+1))); + pq.add(new DistVector(dist, String.valueOf(j + 1))); } } @@ -137,7 +137,7 @@ public static List> computeGroundTruthValues(float[][] indexVectors, public static float[][] getQueryVectors(int queryCount, int dimensions, int docCount, boolean isStandard) { if (isStandard) { - return randomlyGenerateStandardVectors(queryCount, dimensions, docCount+1); + return randomlyGenerateStandardVectors(queryCount, dimensions, docCount + 1); } else { return generateRandomVectors(queryCount, dimensions); } @@ -169,8 +169,8 @@ public static double calculateRecallValue(List> searchResults, List recalls.add(recallVal / k); } - double sum = recalls.stream().reduce((a,b)->a+b).get(); - return sum/recalls.size(); + double sum = recalls.stream().reduce((a, b) -> a + b).get(); + return sum / recalls.size(); } /** @@ -192,14 +192,15 @@ private KNNCodecUtil.Pair readIndexData(String path) throws IOException { BufferedReader reader = new BufferedReader(new FileReader(path)); String line = reader.readLine(); while (line != null) { - Map doc = XContentFactory.xContent(XContentType.JSON).createParser( - NamedXContentRegistry.EMPTY, DeprecationHandler.THROW_UNSUPPORTED_OPERATION, line).map(); + Map doc = XContentFactory.xContent(XContentType.JSON) + .createParser(NamedXContentRegistry.EMPTY, DeprecationHandler.THROW_UNSUPPORTED_OPERATION, line) + .map(); idsList.add((Integer) doc.get("id")); @SuppressWarnings("unchecked") ArrayList vector = (ArrayList) doc.get("vector"); Float[] floatArray = new Float[vector.size()]; - for (int i =0; i< vector.size(); i++) { + for (int i = 0; i < vector.size(); i++) { floatArray[i] = vector.get(i).floatValue(); } vectorsList.add(floatArray); @@ -208,7 +209,7 @@ private KNNCodecUtil.Pair readIndexData(String path) throws IOException { } reader.close(); - int[] idsArray = new int [idsList.size()]; + int[] idsArray = new int[idsList.size()]; float[][] vectorsArray = new float[vectorsList.size()][vectorsList.get(0).length]; for (int i = 0; i < idsList.size(); i++) { idsArray[i] = idsList.get(i); diff --git a/src/test/java/org/opensearch/knn/index/KNNCreateIndexFromModelTests.java b/src/test/java/org/opensearch/knn/index/KNNCreateIndexFromModelTests.java index 3c7905126..e0b2c05cf 100644 --- a/src/test/java/org/opensearch/knn/index/KNNCreateIndexFromModelTests.java +++ b/src/test/java/org/opensearch/knn/index/KNNCreateIndexFromModelTests.java @@ -17,7 +17,6 @@ import org.opensearch.common.Strings; import org.opensearch.common.settings.Settings; import org.opensearch.common.xcontent.XContentFactory; -import org.opensearch.common.xcontent.json.JsonXContent; import org.opensearch.knn.KNNSingleNodeTestCase; import org.opensearch.knn.jni.JNIService; import org.opensearch.knn.index.util.KNNEngine; @@ -48,14 +47,23 @@ public void testCreateIndexFromModel() throws IOException, InterruptedException // "Train" a faiss flat index - this really just creates an empty index that does brute force k-NN long vectorsPointer = JNIService.transferVectors(0, new float[0][0]); - byte [] modelBlob = JNIService.trainIndex(ImmutableMap.of( - INDEX_DESCRIPTION_PARAMETER, "Flat", - SPACE_TYPE, spaceType.getValue()), dimension, vectorsPointer, - KNNEngine.FAISS.getName()); + byte[] modelBlob = JNIService.trainIndex( + ImmutableMap.of(INDEX_DESCRIPTION_PARAMETER, "Flat", SPACE_TYPE, spaceType.getValue()), + dimension, + vectorsPointer, + KNNEngine.FAISS.getName() + ); // Setup model - ModelMetadata modelMetadata = new ModelMetadata(knnEngine, spaceType, dimension, ModelState.CREATED, - ZonedDateTime.now(ZoneOffset.UTC).toString(), "", ""); + ModelMetadata modelMetadata = new ModelMetadata( + knnEngine, + spaceType, + dimension, + ModelState.CREATED, + ZonedDateTime.now(ZoneOffset.UTC).toString(), + "", + "" + ); Model model = new Model(modelMetadata, modelBlob, modelId); @@ -68,35 +76,30 @@ public void testCreateIndexFromModel() throws IOException, InterruptedException String fieldName = "test-field"; final String mapping = Strings.toString( - XContentFactory.jsonBuilder() - .startObject() - .startObject("properties") - .startObject(fieldName) - .field("type", "knn_vector") - .field("model_id", modelId) - .endObject() - .endObject() - .endObject()); + XContentFactory.jsonBuilder() + .startObject() + .startObject("properties") + .startObject(fieldName) + .field("type", "knn_vector") + .field("model_id", modelId) + .endObject() + .endObject() + .endObject() + ); modelDao.put(model, ActionListener.wrap(indexResponse -> { - CreateIndexRequestBuilder createIndexRequestBuilder = client().admin().indices().prepareCreate(indexName) - .setSettings(Settings.builder() - .put("number_of_shards", 1) - .put("number_of_replicas", 0) - .put("index.knn", true) - .build() - ).setMapping(mapping); - - client().admin().indices().create(createIndexRequestBuilder.request(), - ActionListener.wrap( - createIndexResponse -> { - assertTrue(createIndexResponse.isAcknowledged()); - inProgressLatch.countDown(); - }, e -> fail("Unable to create index: " + e.getMessage()) - ) - ); - - }, e ->fail("Unable to put model: " + e.getMessage()))); + CreateIndexRequestBuilder createIndexRequestBuilder = client().admin() + .indices() + .prepareCreate(indexName) + .setSettings(Settings.builder().put("number_of_shards", 1).put("number_of_replicas", 0).put("index.knn", true).build()) + .setMapping(mapping); + + client().admin().indices().create(createIndexRequestBuilder.request(), ActionListener.wrap(createIndexResponse -> { + assertTrue(createIndexResponse.isAcknowledged()); + inProgressLatch.countDown(); + }, e -> fail("Unable to create index: " + e.getMessage()))); + + }, e -> fail("Unable to put model: " + e.getMessage()))); assertTrue(inProgressLatch.await(20, TimeUnit.SECONDS)); } diff --git a/src/test/java/org/opensearch/knn/index/KNNVectorDVLeafFieldDataTests.java b/src/test/java/org/opensearch/knn/index/KNNVectorDVLeafFieldDataTests.java index e1d1889f4..8bda1aefc 100644 --- a/src/test/java/org/opensearch/knn/index/KNNVectorDVLeafFieldDataTests.java +++ b/src/test/java/org/opensearch/knn/index/KNNVectorDVLeafFieldDataTests.java @@ -43,9 +43,11 @@ private void createKNNVectorDocument(Directory directory) throws IOException { IndexWriter writer = new IndexWriter(directory, conf); Document knnDocument = new Document(); knnDocument.add( - new BinaryDocValuesField( - MOCK_INDEX_FIELD_NAME, - new VectorField(MOCK_INDEX_FIELD_NAME, new float[]{1.0f, 2.0f}, new FieldType()).binaryValue())); + new BinaryDocValuesField( + MOCK_INDEX_FIELD_NAME, + new VectorField(MOCK_INDEX_FIELD_NAME, new float[] { 1.0f, 2.0f }, new FieldType()).binaryValue() + ) + ); knnDocument.add(new NumericDocValuesField(MOCK_NUMERIC_INDEX_FIELD_NAME, 1000)); writer.addDocument(knnDocument); writer.commit(); @@ -67,16 +69,14 @@ public void testGetScriptValues() { } public void testGetScriptValuesWrongFieldName() { - KNNVectorDVLeafFieldData leafFieldData = new KNNVectorDVLeafFieldData( - leafReaderContext.reader(), "invalid"); + KNNVectorDVLeafFieldData leafFieldData = new KNNVectorDVLeafFieldData(leafReaderContext.reader(), "invalid"); ScriptDocValues scriptValues = leafFieldData.getScriptValues(); assertNotNull(scriptValues); } public void testGetScriptValuesWrongFieldType() { - KNNVectorDVLeafFieldData leafFieldData = new KNNVectorDVLeafFieldData( - leafReaderContext.reader(), MOCK_NUMERIC_INDEX_FIELD_NAME); - expectThrows(IllegalStateException.class, ()->leafFieldData.getScriptValues()); + KNNVectorDVLeafFieldData leafFieldData = new KNNVectorDVLeafFieldData(leafReaderContext.reader(), MOCK_NUMERIC_INDEX_FIELD_NAME); + expectThrows(IllegalStateException.class, () -> leafFieldData.getScriptValues()); } public void testRamBytesUsed() { @@ -86,7 +86,6 @@ public void testRamBytesUsed() { public void testGetBytesValues() { KNNVectorDVLeafFieldData leafFieldData = new KNNVectorDVLeafFieldData(leafReaderContext.reader(), ""); - expectThrows(UnsupportedOperationException.class, - () -> leafFieldData.getBytesValues()); + expectThrows(UnsupportedOperationException.class, () -> leafFieldData.getBytesValues()); } } diff --git a/src/test/java/org/opensearch/knn/index/KNNVectorIndexFieldDataTests.java b/src/test/java/org/opensearch/knn/index/KNNVectorIndexFieldDataTests.java index 346052618..8523c4146 100644 --- a/src/test/java/org/opensearch/knn/index/KNNVectorIndexFieldDataTests.java +++ b/src/test/java/org/opensearch/knn/index/KNNVectorIndexFieldDataTests.java @@ -72,13 +72,14 @@ public void testLoadDirect() throws IOException { public void testSortField() { - expectThrows(UnsupportedOperationException.class, - () -> indexFieldData.sortField(null, null, null, false)); + expectThrows(UnsupportedOperationException.class, () -> indexFieldData.sortField(null, null, null, false)); } public void testNewBucketedSort() { - expectThrows(UnsupportedOperationException.class, - () -> indexFieldData.newBucketedSort(null, null, null, null, null, null, 0, null)); + expectThrows( + UnsupportedOperationException.class, + () -> indexFieldData.newBucketedSort(null, null, null, null, null, null, 0, null) + ); } } diff --git a/src/test/java/org/opensearch/knn/index/KNNVectorScriptDocValuesTests.java b/src/test/java/org/opensearch/knn/index/KNNVectorScriptDocValuesTests.java index e54e19141..876117940 100644 --- a/src/test/java/org/opensearch/knn/index/KNNVectorScriptDocValuesTests.java +++ b/src/test/java/org/opensearch/knn/index/KNNVectorScriptDocValuesTests.java @@ -23,7 +23,7 @@ public class KNNVectorScriptDocValuesTests extends KNNTestCase { private static final String MOCK_INDEX_FIELD_NAME = "test-index-field-name"; - private static final float[] SAMPLE_VECTOR_DATA = new float[]{1.0f, 2.0f}; + private static final float[] SAMPLE_VECTOR_DATA = new float[] { 1.0f, 2.0f }; private KNNVectorScriptDocValues scriptDocValues; private Directory directory; private DirectoryReader reader; @@ -36,7 +36,9 @@ public void setUp() throws Exception { reader = DirectoryReader.open(directory); LeafReaderContext leafReaderContext = reader.getContext().leaves().get(0); scriptDocValues = new KNNVectorScriptDocValues( - leafReaderContext.reader().getBinaryDocValues(MOCK_INDEX_FIELD_NAME), MOCK_INDEX_FIELD_NAME); + leafReaderContext.reader().getBinaryDocValues(MOCK_INDEX_FIELD_NAME), + MOCK_INDEX_FIELD_NAME + ); } private void createKNNVectorDocument(Directory directory) throws IOException { @@ -44,9 +46,11 @@ private void createKNNVectorDocument(Directory directory) throws IOException { IndexWriter writer = new IndexWriter(directory, conf); Document knnDocument = new Document(); knnDocument.add( - new BinaryDocValuesField( - MOCK_INDEX_FIELD_NAME, - new VectorField(MOCK_INDEX_FIELD_NAME, SAMPLE_VECTOR_DATA, new FieldType()).binaryValue())); + new BinaryDocValuesField( + MOCK_INDEX_FIELD_NAME, + new VectorField(MOCK_INDEX_FIELD_NAME, SAMPLE_VECTOR_DATA, new FieldType()).binaryValue() + ) + ); writer.addDocument(knnDocument); writer.commit(); writer.close(); @@ -64,8 +68,7 @@ public void testGetValue() throws IOException { Assert.assertArrayEquals(SAMPLE_VECTOR_DATA, scriptDocValues.getValue(), 0.1f); } - - //Test getValue without calling setNextDocId + // Test getValue without calling setNextDocId public void testGetValueFails() throws IOException { expectThrows(IllegalStateException.class, () -> scriptDocValues.getValue()); } diff --git a/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java b/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java index dd7b82cf3..47d13a87e 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java @@ -19,7 +19,6 @@ import org.opensearch.knn.index.KNNWeight; import org.opensearch.knn.index.SpaceType; import org.opensearch.knn.index.VectorField; -import org.opensearch.knn.index.codec.KNN87Codec.KNN87Codec; import org.apache.lucene.codecs.Codec; import org.apache.lucene.document.Document; import org.apache.lucene.document.FieldType; diff --git a/src/test/java/org/opensearch/knn/plugin/script/KNNScoringUtilTests.java b/src/test/java/org/opensearch/knn/plugin/script/KNNScoringUtilTests.java index 0375b6b48..7f66e909a 100644 --- a/src/test/java/org/opensearch/knn/plugin/script/KNNScoringUtilTests.java +++ b/src/test/java/org/opensearch/knn/plugin/script/KNNScoringUtilTests.java @@ -35,36 +35,35 @@ private List getTestQueryVector() { } public void testL2SquaredScoringFunction() { - float[] queryVector = {1.0f, 1.0f, 1.0f}; - float[] inputVector = {4.0f, 4.0f, 4.0f}; + float[] queryVector = { 1.0f, 1.0f, 1.0f }; + float[] inputVector = { 4.0f, 4.0f, 4.0f }; Float distance = KNNScoringUtil.l2Squared(queryVector, inputVector); assertTrue(distance == 27.0f); } public void testWrongDimensionL2SquaredScoringFunction() { - float[] queryVector = {1.0f, 1.0f}; - float[] inputVector = {4.0f, 4.0f, 4.0f}; + float[] queryVector = { 1.0f, 1.0f }; + float[] inputVector = { 4.0f, 4.0f, 4.0f }; expectThrows(IllegalArgumentException.class, () -> KNNScoringUtil.l2Squared(queryVector, inputVector)); } public void testCosineSimilScoringFunction() { - float[] queryVector = {1.0f, 1.0f, 1.0f}; - float[] inputVector = {4.0f, 4.0f, 4.0f}; + float[] queryVector = { 1.0f, 1.0f, 1.0f }; + float[] inputVector = { 4.0f, 4.0f, 4.0f }; float queryVectorMagnitude = KNNScoringSpaceUtil.getVectorMagnitudeSquared(queryVector); float inputVectorMagnitude = KNNScoringSpaceUtil.getVectorMagnitudeSquared(inputVector); float dotProduct = 12.0f; float expectedScore = (float) (dotProduct / (Math.sqrt(queryVectorMagnitude * inputVectorMagnitude))); - Float actualScore = KNNScoringUtil.cosinesimil(queryVector, inputVector); assertEquals(expectedScore, actualScore, 0.0001); } public void testCosineSimilOptimizedScoringFunction() { - float[] queryVector = {1.0f, 1.0f, 1.0f}; - float[] inputVector = {4.0f, 4.0f, 4.0f}; + float[] queryVector = { 1.0f, 1.0f, 1.0f }; + float[] inputVector = { 4.0f, 4.0f, 4.0f }; float queryVectorMagnitude = KNNScoringSpaceUtil.getVectorMagnitudeSquared(queryVector); float inputVectorMagnitude = KNNScoringSpaceUtil.getVectorMagnitudeSquared(inputVector); float dotProduct = 12.0f; @@ -86,26 +85,26 @@ public void testConvertInvalidVectorToPrimitive() { } public void testCosineSimilQueryVectorZeroMagnitude() { - float[] queryVector = {0, 0}; - float[] inputVector = {4.0f, 4.0f}; + float[] queryVector = { 0, 0 }; + float[] inputVector = { 4.0f, 4.0f }; assertEquals(0, KNNScoringUtil.cosinesimil(queryVector, inputVector), 0.00001); } public void testCosineSimilOptimizedQueryVectorZeroMagnitude() { - float[] inputVector = {4.0f, 4.0f}; - float[] queryVector = {0, 0}; + float[] inputVector = { 4.0f, 4.0f }; + float[] queryVector = { 0, 0 }; assertTrue(0 == KNNScoringUtil.cosinesimilOptimized(queryVector, inputVector, 0.0f)); } public void testWrongDimensionCosineSimilScoringFunction() { - float[] queryVector = {1.0f, 1.0f}; - float[] inputVector = {4.0f, 4.0f, 4.0f}; + float[] queryVector = { 1.0f, 1.0f }; + float[] inputVector = { 4.0f, 4.0f, 4.0f }; expectThrows(IllegalArgumentException.class, () -> KNNScoringUtil.cosinesimil(queryVector, inputVector)); } public void testWrongDimensionCosineSimilOPtimizedScoringFunction() { - float[] queryVector = {1.0f, 1.0f}; - float[] inputVector = {4.0f, 4.0f, 4.0f}; + float[] queryVector = { 1.0f, 1.0f }; + float[] inputVector = { 4.0f, 4.0f, 4.0f }; expectThrows(IllegalArgumentException.class, () -> KNNScoringUtil.cosinesimilOptimized(queryVector, inputVector, 1.0f)); } @@ -173,7 +172,7 @@ public void testBitHammingDistance_Long() { public void testL2SquaredWhitelistedScoringFunction() throws IOException { List queryVector = getTestQueryVector(); TestKNNScriptDocValues dataset = new TestKNNScriptDocValues(); - dataset.createKNNVectorDocument(new float[]{4.0f, 4.0f, 4.0f}, "test-index-field-name"); + dataset.createKNNVectorDocument(new float[] { 4.0f, 4.0f, 4.0f }, "test-index-field-name"); KNNVectorScriptDocValues scriptDocValues = dataset.getScriptDocValues("test-index-field-name"); scriptDocValues.setNextDocId(0); Float distance = KNNScoringUtil.l2Squared(queryVector, scriptDocValues); @@ -184,7 +183,7 @@ public void testL2SquaredWhitelistedScoringFunction() throws IOException { public void testScriptDocValuesFailsL2() throws IOException { List queryVector = getTestQueryVector(); TestKNNScriptDocValues dataset = new TestKNNScriptDocValues(); - dataset.createKNNVectorDocument(new float[]{4.0f, 4.0f, 4.0f}, "test-index-field-name"); + dataset.createKNNVectorDocument(new float[] { 4.0f, 4.0f, 4.0f }, "test-index-field-name"); KNNVectorScriptDocValues scriptDocValues = dataset.getScriptDocValues("test-index-field-name"); expectThrows(IllegalStateException.class, () -> KNNScoringUtil.l2Squared(queryVector, scriptDocValues)); dataset.close(); @@ -193,7 +192,7 @@ public void testScriptDocValuesFailsL2() throws IOException { public void testCosineSimilarityScoringFunction() throws IOException { List queryVector = getTestQueryVector(); TestKNNScriptDocValues dataset = new TestKNNScriptDocValues(); - dataset.createKNNVectorDocument(new float[]{4.0f, 4.0f, 4.0f}, "test-index-field-name"); + dataset.createKNNVectorDocument(new float[] { 4.0f, 4.0f, 4.0f }, "test-index-field-name"); KNNVectorScriptDocValues scriptDocValues = dataset.getScriptDocValues("test-index-field-name"); scriptDocValues.setNextDocId(0); @@ -205,7 +204,7 @@ public void testCosineSimilarityScoringFunction() throws IOException { public void testScriptDocValuesFailsCosineSimilarity() throws IOException { List queryVector = getTestQueryVector(); TestKNNScriptDocValues dataset = new TestKNNScriptDocValues(); - dataset.createKNNVectorDocument(new float[]{4.0f, 4.0f, 4.0f}, "test-index-field-name"); + dataset.createKNNVectorDocument(new float[] { 4.0f, 4.0f, 4.0f }, "test-index-field-name"); KNNVectorScriptDocValues scriptDocValues = dataset.getScriptDocValues("test-index-field-name"); expectThrows(IllegalStateException.class, () -> KNNScoringUtil.cosineSimilarity(queryVector, scriptDocValues)); dataset.close(); @@ -214,7 +213,7 @@ public void testScriptDocValuesFailsCosineSimilarity() throws IOException { public void testCosineSimilarityOptimizedScoringFunction() throws IOException { List queryVector = getTestQueryVector(); TestKNNScriptDocValues dataset = new TestKNNScriptDocValues(); - dataset.createKNNVectorDocument(new float[]{4.0f, 4.0f, 4.0f}, "test-index-field-name"); + dataset.createKNNVectorDocument(new float[] { 4.0f, 4.0f, 4.0f }, "test-index-field-name"); KNNVectorScriptDocValues scriptDocValues = dataset.getScriptDocValues("test-index-field-name"); scriptDocValues.setNextDocId(0); Float actualScore = KNNScoringUtil.cosineSimilarity(queryVector, scriptDocValues, 3.0f); @@ -225,7 +224,7 @@ public void testCosineSimilarityOptimizedScoringFunction() throws IOException { public void testScriptDocValuesFailsCosineSimilarityOptimized() throws IOException { List queryVector = getTestQueryVector(); TestKNNScriptDocValues dataset = new TestKNNScriptDocValues(); - dataset.createKNNVectorDocument(new float[]{4.0f, 4.0f, 4.0f}, "test-index-field-name"); + dataset.createKNNVectorDocument(new float[] { 4.0f, 4.0f, 4.0f }, "test-index-field-name"); KNNVectorScriptDocValues scriptDocValues = dataset.getScriptDocValues("test-index-field-name"); expectThrows(IllegalStateException.class, () -> KNNScoringUtil.cosineSimilarity(queryVector, scriptDocValues, 3.0f)); dataset.close(); @@ -244,16 +243,14 @@ public KNNVectorScriptDocValues getScriptDocValues(String fieldName) throws IOEx if (scriptDocValues == null) { reader = DirectoryReader.open(directory); LeafReaderContext leafReaderContext = reader.getContext().leaves().get(0); - scriptDocValues = new KNNVectorScriptDocValues(leafReaderContext.reader().getBinaryDocValues(fieldName),fieldName ); + scriptDocValues = new KNNVectorScriptDocValues(leafReaderContext.reader().getBinaryDocValues(fieldName), fieldName); } return scriptDocValues; } public void close() throws IOException { - if (reader != null) - reader.close(); - if (directory != null) - directory.close(); + if (reader != null) reader.close(); + if (directory != null) directory.close(); } public void createKNNVectorDocument(final float[] content, final String fieldName) throws IOException { @@ -261,10 +258,7 @@ public void createKNNVectorDocument(final float[] content, final String fieldNam IndexWriter writer = new IndexWriter(directory, conf); conf.setMergePolicy(NoMergePolicy.INSTANCE); // prevent merges for this test Document knnDocument = new Document(); - knnDocument.add( - new BinaryDocValuesField( - fieldName, - new VectorField(fieldName, content, new FieldType()).binaryValue())); + knnDocument.add(new BinaryDocValuesField(fieldName, new VectorField(fieldName, content, new FieldType()).binaryValue())); writer.addDocument(knnDocument); writer.commit(); writer.close(); From a9ddc3afacdccbbc811362e0f877d318523d1792 Mon Sep 17 00:00:00 2001 From: Martin Gaievski Date: Mon, 21 Mar 2022 17:48:11 -0700 Subject: [PATCH 3/9] Refactor codec and doc format classes for Lucene 9.1 Signed-off-by: Martin Gaievski --- .../index/codec/KNN91Codec/KNN91Codec.java | 13 +++---- .../codec/KNN91Codec/KNN91DocFormat.java | 37 +++++++++++++++++++ .../{ => docformat}/KNN91BinaryDocValues.java | 10 ++++- .../{ => docformat}/KNN91CompoundFormat.java | 10 ++++- .../KNN91DocValuesConsumer.java | 10 ++++- .../{ => docformat}/KNN91DocValuesFormat.java | 10 ++++- .../{ => docformat}/KNN91DocValuesReader.java | 10 ++++- .../knn/index/codec/KNNCodecFactory.java | 22 +++++++++++ .../knn/index/codec/KNNCodecService.java | 3 +- .../knn/index/codec/KNNDocFormatFacade.java | 22 +++++++++++ .../knn/index/codec/KNNDocFormatFactory.java | 28 ++++++++++++++ .../KNN91BinaryDocValuesTests.java | 10 ++++- .../KNN91CompoundFormatTests.java | 11 +++++- .../KNN91DocValuesConsumerTests.java | 10 ++++- 14 files changed, 181 insertions(+), 25 deletions(-) create mode 100644 src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocFormat.java rename src/main/java/org/opensearch/knn/index/codec/KNN91Codec/{ => docformat}/KNN91BinaryDocValues.java (82%) rename src/main/java/org/opensearch/knn/index/codec/KNN91Codec/{ => docformat}/KNN91CompoundFormat.java (88%) rename src/main/java/org/opensearch/knn/index/codec/KNN91Codec/{ => docformat}/KNN91DocValuesConsumer.java (97%) rename src/main/java/org/opensearch/knn/index/codec/KNN91Codec/{ => docformat}/KNN91DocValuesFormat.java (80%) rename src/main/java/org/opensearch/knn/index/codec/KNN91Codec/{ => docformat}/KNN91DocValuesReader.java (85%) create mode 100644 src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java create mode 100644 src/main/java/org/opensearch/knn/index/codec/KNNDocFormatFacade.java create mode 100644 src/main/java/org/opensearch/knn/index/codec/KNNDocFormatFactory.java rename src/test/java/org/opensearch/knn/index/codec/KNN91Codec/{ => docformat}/KNN91BinaryDocValuesTests.java (89%) rename src/test/java/org/opensearch/knn/index/codec/KNN91Codec/{ => docformat}/KNN91CompoundFormatTests.java (90%) rename src/test/java/org/opensearch/knn/index/codec/KNN91Codec/{ => docformat}/KNN91DocValuesConsumerTests.java (98%) diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91Codec.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91Codec.java index e63585573..449979664 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91Codec.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91Codec.java @@ -16,6 +16,8 @@ import org.apache.lucene.codecs.DocValuesFormat; import org.apache.lucene.codecs.FilterCodec; import org.apache.lucene.codecs.lucene91.Lucene91Codec; +import org.opensearch.knn.index.codec.KNNDocFormatFacade; +import org.opensearch.knn.index.codec.KNNDocFormatFactory; /** * Extends the Codec to support a new file format for KNN index @@ -24,10 +26,8 @@ */ public final class KNN91Codec extends FilterCodec { - private final DocValuesFormat docValuesFormat; - private final CompoundFormat compoundFormat; - public static final String KNN_91 = "KNN91Codec"; + private KNNDocFormatFacade docFormatFacade; /** * No arg constructor that uses Lucene91 as the delegate @@ -43,17 +43,16 @@ public KNN91Codec() { */ public KNN91Codec(Codec delegate) { super(KNN_91, delegate); - this.docValuesFormat = new KNN91DocValuesFormat(delegate.docValuesFormat()); - this.compoundFormat = new KNN91CompoundFormat(delegate.compoundFormat()); + docFormatFacade = KNNDocFormatFactory.createKNN91DocFormat(delegate); } @Override public DocValuesFormat docValuesFormat() { - return this.docValuesFormat; + return docFormatFacade.docValuesFormat(); } @Override public CompoundFormat compoundFormat() { - return this.compoundFormat; + return docFormatFacade.compoundFormat(); } } diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocFormat.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocFormat.java new file mode 100644 index 000000000..72b2d082f --- /dev/null +++ b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocFormat.java @@ -0,0 +1,37 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.knn.index.codec.KNN91Codec; + +import org.apache.lucene.codecs.CompoundFormat; +import org.apache.lucene.codecs.DocValuesFormat; +import org.opensearch.knn.index.codec.KNNDocFormatFacade; + +public class KNN91DocFormat implements KNNDocFormatFacade { + + private final DocValuesFormat docValuesFormat; + private final CompoundFormat compoundFormat; + + public KNN91DocFormat(final DocValuesFormat docValuesFormat, final CompoundFormat compoundFormat) { + this.docValuesFormat = docValuesFormat; + this.compoundFormat = compoundFormat; + } + + @Override + public DocValuesFormat docValuesFormat() { + return docValuesFormat; + } + + @Override + public CompoundFormat compoundFormat() { + return compoundFormat; + } +} diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91BinaryDocValues.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValues.java similarity index 82% rename from src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91BinaryDocValues.java rename to src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValues.java index f9f89e589..7c3baf5c5 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91BinaryDocValues.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValues.java @@ -1,9 +1,15 @@ /* - * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. */ -package org.opensearch.knn.index.codec.KNN91Codec; +package org.opensearch.knn.index.codec.KNN91Codec.docformat; import org.apache.lucene.index.BinaryDocValues; import org.apache.lucene.index.DocIDMerger; diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CompoundFormat.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormat.java similarity index 88% rename from src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CompoundFormat.java rename to src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormat.java index 6d8b0fdb7..360806c5c 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CompoundFormat.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormat.java @@ -1,9 +1,15 @@ /* - * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. */ -package org.opensearch.knn.index.codec.KNN91Codec; +package org.opensearch.knn.index.codec.KNN91Codec.docformat; import org.apache.lucene.codecs.CompoundDirectory; import org.apache.lucene.codecs.CompoundFormat; diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesConsumer.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumer.java similarity index 97% rename from src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesConsumer.java rename to src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumer.java index 058ab755b..a02af15e3 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesConsumer.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumer.java @@ -1,9 +1,15 @@ /* - * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. */ -package org.opensearch.knn.index.codec.KNN91Codec; +package org.opensearch.knn.index.codec.KNN91Codec.docformat; import com.google.common.collect.ImmutableMap; import org.apache.logging.log4j.LogManager; diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesFormat.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesFormat.java similarity index 80% rename from src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesFormat.java rename to src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesFormat.java index 2d5613ace..be8f0a672 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesFormat.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesFormat.java @@ -1,9 +1,15 @@ /* - * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. */ -package org.opensearch.knn.index.codec.KNN91Codec; +package org.opensearch.knn.index.codec.KNN91Codec.docformat; import org.apache.lucene.codecs.DocValuesConsumer; import org.apache.lucene.codecs.DocValuesFormat; diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesReader.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesReader.java similarity index 85% rename from src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesReader.java rename to src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesReader.java index b6a9622ef..c63bed0f9 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesReader.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesReader.java @@ -1,9 +1,15 @@ /* - * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. */ -package org.opensearch.knn.index.codec.KNN91Codec; +package org.opensearch.knn.index.codec.KNN91Codec.docformat; import org.apache.lucene.codecs.DocValuesProducer; import org.apache.lucene.index.BinaryDocValues; diff --git a/src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java b/src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java new file mode 100644 index 000000000..418e87c64 --- /dev/null +++ b/src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java @@ -0,0 +1,22 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.knn.index.codec; + +import org.apache.lucene.codecs.Codec; +import org.opensearch.knn.index.codec.KNN91Codec.KNN91Codec; + +public class KNNCodecFactory { + + public static Codec createKNN91Codec(Codec userCodec) { + return new KNN91Codec(userCodec); + } +} diff --git a/src/main/java/org/opensearch/knn/index/codec/KNNCodecService.java b/src/main/java/org/opensearch/knn/index/codec/KNNCodecService.java index 8bbf37b03..622f6ef17 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNNCodecService.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNNCodecService.java @@ -8,7 +8,6 @@ import org.opensearch.index.codec.CodecServiceConfig; import org.apache.lucene.codecs.Codec; import org.opensearch.index.codec.CodecService; -import org.opensearch.knn.index.codec.KNN91Codec.KNN91Codec; /** * KNNCodecService to inject the right KNNCodec version @@ -27,6 +26,6 @@ public KNNCodecService(CodecServiceConfig codecServiceConfig) { */ @Override public Codec codec(String name) { - return new KNN91Codec(super.codec(name)); + return KNNCodecFactory.createKNN91Codec(super.codec(name)); } } diff --git a/src/main/java/org/opensearch/knn/index/codec/KNNDocFormatFacade.java b/src/main/java/org/opensearch/knn/index/codec/KNNDocFormatFacade.java new file mode 100644 index 000000000..d5b6b5ff1 --- /dev/null +++ b/src/main/java/org/opensearch/knn/index/codec/KNNDocFormatFacade.java @@ -0,0 +1,22 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.knn.index.codec; + +import org.apache.lucene.codecs.CompoundFormat; +import org.apache.lucene.codecs.DocValuesFormat; + +public interface KNNDocFormatFacade { + + DocValuesFormat docValuesFormat(); + + CompoundFormat compoundFormat(); +} diff --git a/src/main/java/org/opensearch/knn/index/codec/KNNDocFormatFactory.java b/src/main/java/org/opensearch/knn/index/codec/KNNDocFormatFactory.java new file mode 100644 index 000000000..08b43985f --- /dev/null +++ b/src/main/java/org/opensearch/knn/index/codec/KNNDocFormatFactory.java @@ -0,0 +1,28 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.knn.index.codec; + +import org.apache.lucene.codecs.Codec; +import org.opensearch.knn.index.codec.KNN91Codec.docformat.KNN91CompoundFormat; +import org.opensearch.knn.index.codec.KNN91Codec.KNN91DocFormat; +import org.opensearch.knn.index.codec.KNN91Codec.docformat.KNN91DocValuesFormat; + +public class KNNDocFormatFactory { + + public static KNNDocFormatFacade createKNN91DocFormat(Codec delegate) { + final KNNDocFormatFacade knnDocFormatFacade = new KNN91DocFormat( + new KNN91DocValuesFormat(delegate.docValuesFormat()), + new KNN91CompoundFormat(delegate.compoundFormat()) + ); + return knnDocFormatFacade; + } +} diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91BinaryDocValuesTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValuesTests.java similarity index 89% rename from src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91BinaryDocValuesTests.java rename to src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValuesTests.java index a548b4544..885a60e5d 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91BinaryDocValuesTests.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValuesTests.java @@ -1,9 +1,15 @@ /* - * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. */ -package org.opensearch.knn.index.codec.KNN91Codec; +package org.opensearch.knn.index.codec.KNN91Codec.docformat; import com.google.common.collect.ImmutableList; import org.apache.lucene.index.BinaryDocValues; diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CompoundFormatTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormatTests.java similarity index 90% rename from src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CompoundFormatTests.java rename to src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormatTests.java index dd9a358d7..567d0da86 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CompoundFormatTests.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormatTests.java @@ -1,9 +1,15 @@ /* - * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. */ -package org.opensearch.knn.index.codec.KNN91Codec; +package org.opensearch.knn.index.codec.KNN91Codec.docformat; import org.apache.lucene.codecs.Codec; import org.apache.lucene.codecs.CompoundDirectory; @@ -16,6 +22,7 @@ import org.junit.BeforeClass; import org.opensearch.common.util.set.Sets; import org.opensearch.knn.KNNTestCase; +import org.opensearch.knn.index.codec.KNN91Codec.KNN91Codec; import org.opensearch.knn.index.codec.KNNCodecTestUtil; import org.opensearch.knn.index.util.KNNEngine; diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesConsumerTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumerTests.java similarity index 98% rename from src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesConsumerTests.java rename to src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumerTests.java index 994540cf5..a573f273a 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocValuesConsumerTests.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumerTests.java @@ -1,9 +1,15 @@ /* - * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. */ -package org.opensearch.knn.index.codec.KNN91Codec; +package org.opensearch.knn.index.codec.KNN91Codec.docformat; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; From 13f075d4426f34a4f3f6bec498e66941c38e842a Mon Sep 17 00:00:00 2001 From: Martin Gaievski Date: Tue, 22 Mar 2022 09:12:53 -0700 Subject: [PATCH 4/9] Adding java doc to some java classes to kickoff new build Signed-off-by: Martin Gaievski --- .../java/org/opensearch/knn/index/codec/KNNCodecFactory.java | 3 +++ .../org/opensearch/knn/index/codec/KNNDocFormatFactory.java | 3 +++ 2 files changed, 6 insertions(+) diff --git a/src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java b/src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java index 418e87c64..378e85672 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java @@ -14,6 +14,9 @@ import org.apache.lucene.codecs.Codec; import org.opensearch.knn.index.codec.KNN91Codec.KNN91Codec; +/** + * Factory abstraction for KNN codes + */ public class KNNCodecFactory { public static Codec createKNN91Codec(Codec userCodec) { diff --git a/src/main/java/org/opensearch/knn/index/codec/KNNDocFormatFactory.java b/src/main/java/org/opensearch/knn/index/codec/KNNDocFormatFactory.java index 08b43985f..e3a3d5018 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNNDocFormatFactory.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNNDocFormatFactory.java @@ -16,6 +16,9 @@ import org.opensearch.knn.index.codec.KNN91Codec.KNN91DocFormat; import org.opensearch.knn.index.codec.KNN91Codec.docformat.KNN91DocValuesFormat; +/** + * Factory abstraction for KNN document format facades + */ public class KNNDocFormatFactory { public static KNNDocFormatFacade createKNN91DocFormat(Codec delegate) { From 82a05ffd84f4a0f972568b128021b7eefb71e4a8 Mon Sep 17 00:00:00 2001 From: Martin Gaievski Date: Tue, 22 Mar 2022 14:57:09 -0700 Subject: [PATCH 5/9] Reorginize factory and facade code, minor review comments Signed-off-by: Martin Gaievski --- .../index/codec/KNN91Codec/KNN91Codec.java | 26 ++++------ .../codec/KNN91Codec/KNN91DocFormat.java | 37 --------------- .../docformat/KNN91BinaryDocValues.java | 9 +--- .../docformat/KNN91CompoundFormat.java | 9 +--- .../docformat/KNN91DocValuesConsumer.java | 9 +--- .../docformat/KNN91DocValuesFormat.java | 14 +----- .../docformat/KNN91DocValuesReader.java | 9 +--- .../knn/index/codec/KNNCodecFactory.java | 47 +++++++++++++++---- .../knn/index/codec/KNNCodecService.java | 2 +- .../knn/index/codec/KNNDocFormatFacade.java | 22 --------- .../knn/index/codec/KNNDocFormatFactory.java | 31 ------------ .../knn/index/codec/KNNFormatFacade.java | 30 ++++++++++++ .../knn/index/codec/KNNFormatFactory.java | 23 +++++++++ 13 files changed, 106 insertions(+), 162 deletions(-) delete mode 100644 src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocFormat.java delete mode 100644 src/main/java/org/opensearch/knn/index/codec/KNNDocFormatFacade.java delete mode 100644 src/main/java/org/opensearch/knn/index/codec/KNNDocFormatFactory.java create mode 100644 src/main/java/org/opensearch/knn/index/codec/KNNFormatFacade.java create mode 100644 src/main/java/org/opensearch/knn/index/codec/KNNFormatFactory.java diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91Codec.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91Codec.java index 449979664..5894ee5a1 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91Codec.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91Codec.java @@ -1,23 +1,17 @@ /* + * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - * - * Modifications Copyright OpenSearch Contributors. See - * GitHub history for details. */ - package org.opensearch.knn.index.codec.KNN91Codec; import org.apache.lucene.codecs.Codec; import org.apache.lucene.codecs.CompoundFormat; import org.apache.lucene.codecs.DocValuesFormat; import org.apache.lucene.codecs.FilterCodec; -import org.apache.lucene.codecs.lucene91.Lucene91Codec; -import org.opensearch.knn.index.codec.KNNDocFormatFacade; -import org.opensearch.knn.index.codec.KNNDocFormatFactory; +import org.opensearch.knn.index.codec.KNNFormatFacade; +import org.opensearch.knn.index.codec.KNNFormatFactory; + +import static org.opensearch.knn.index.codec.KNNCodecFactory.CodecDelegateFactory.createKNN91DefaultDelegate; /** * Extends the Codec to support a new file format for KNN index @@ -27,13 +21,13 @@ public final class KNN91Codec extends FilterCodec { public static final String KNN_91 = "KNN91Codec"; - private KNNDocFormatFacade docFormatFacade; + private KNNFormatFacade knnFormatFacade; /** * No arg constructor that uses Lucene91 as the delegate */ public KNN91Codec() { - this(new Lucene91Codec()); + this(createKNN91DefaultDelegate()); } /** @@ -43,16 +37,16 @@ public KNN91Codec() { */ public KNN91Codec(Codec delegate) { super(KNN_91, delegate); - docFormatFacade = KNNDocFormatFactory.createKNN91DocFormat(delegate); + knnFormatFacade = KNNFormatFactory.createKNN91Format(delegate); } @Override public DocValuesFormat docValuesFormat() { - return docFormatFacade.docValuesFormat(); + return knnFormatFacade.docValuesFormat(); } @Override public CompoundFormat compoundFormat() { - return docFormatFacade.compoundFormat(); + return knnFormatFacade.compoundFormat(); } } diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocFormat.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocFormat.java deleted file mode 100644 index 72b2d082f..000000000 --- a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91DocFormat.java +++ /dev/null @@ -1,37 +0,0 @@ -/* - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - * - * Modifications Copyright OpenSearch Contributors. See - * GitHub history for details. - */ - -package org.opensearch.knn.index.codec.KNN91Codec; - -import org.apache.lucene.codecs.CompoundFormat; -import org.apache.lucene.codecs.DocValuesFormat; -import org.opensearch.knn.index.codec.KNNDocFormatFacade; - -public class KNN91DocFormat implements KNNDocFormatFacade { - - private final DocValuesFormat docValuesFormat; - private final CompoundFormat compoundFormat; - - public KNN91DocFormat(final DocValuesFormat docValuesFormat, final CompoundFormat compoundFormat) { - this.docValuesFormat = docValuesFormat; - this.compoundFormat = compoundFormat; - } - - @Override - public DocValuesFormat docValuesFormat() { - return docValuesFormat; - } - - @Override - public CompoundFormat compoundFormat() { - return compoundFormat; - } -} diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValues.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValues.java index 7c3baf5c5..b0d015e64 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValues.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValues.java @@ -1,14 +1,7 @@ /* + * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - * - * Modifications Copyright OpenSearch Contributors. See - * GitHub history for details. */ - package org.opensearch.knn.index.codec.KNN91Codec.docformat; import org.apache.lucene.index.BinaryDocValues; diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormat.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormat.java index 360806c5c..ac3e59db2 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormat.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormat.java @@ -1,14 +1,7 @@ /* + * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - * - * Modifications Copyright OpenSearch Contributors. See - * GitHub history for details. */ - package org.opensearch.knn.index.codec.KNN91Codec.docformat; import org.apache.lucene.codecs.CompoundDirectory; diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumer.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumer.java index a02af15e3..5dc187346 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumer.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumer.java @@ -1,14 +1,7 @@ /* + * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - * - * Modifications Copyright OpenSearch Contributors. See - * GitHub history for details. */ - package org.opensearch.knn.index.codec.KNN91Codec.docformat; import com.google.common.collect.ImmutableMap; diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesFormat.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesFormat.java index be8f0a672..3e24af87b 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesFormat.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesFormat.java @@ -1,20 +1,12 @@ /* + * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - * - * Modifications Copyright OpenSearch Contributors. See - * GitHub history for details. */ - package org.opensearch.knn.index.codec.KNN91Codec.docformat; import org.apache.lucene.codecs.DocValuesConsumer; import org.apache.lucene.codecs.DocValuesFormat; import org.apache.lucene.codecs.DocValuesProducer; -import org.apache.lucene.codecs.lucene90.Lucene90DocValuesFormat; import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.SegmentWriteState; @@ -26,10 +18,6 @@ public class KNN91DocValuesFormat extends DocValuesFormat { private final DocValuesFormat delegate; - public KNN91DocValuesFormat() { - this(new Lucene90DocValuesFormat()); - } - /** * Constructor that takes delegate in order to handle non-overridden methods * diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesReader.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesReader.java index c63bed0f9..d870710df 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesReader.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesReader.java @@ -1,14 +1,7 @@ /* + * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - * - * Modifications Copyright OpenSearch Contributors. See - * GitHub history for details. */ - package org.opensearch.knn.index.codec.KNN91Codec.docformat; import org.apache.lucene.codecs.DocValuesProducer; diff --git a/src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java b/src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java index 378e85672..68c4b879d 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java @@ -1,25 +1,52 @@ /* + * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - * - * Modifications Copyright OpenSearch Contributors. See - * GitHub history for details. */ - package org.opensearch.knn.index.codec; +import com.google.common.collect.ImmutableMap; import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.lucene91.Lucene91Codec; import org.opensearch.knn.index.codec.KNN91Codec.KNN91Codec; +import java.lang.reflect.Constructor; +import java.util.Map; + /** * Factory abstraction for KNN codes */ public class KNNCodecFactory { - public static Codec createKNN91Codec(Codec userCodec) { - return new KNN91Codec(userCodec); + private static Map CODEC_BY_VERSION = ImmutableMap.of(KNNCodecVersion.KNN91, KNN91Codec.class); + + private static KNNCodecVersion LATEST_KNN_CODEC_VERSION = KNNCodecVersion.KNN91; + + public static Codec createKNNCodec(final Codec userCodec) { + return getCodec(LATEST_KNN_CODEC_VERSION, userCodec); + } + + public static Codec createKNNCodec(final KNNCodecVersion knnCodecVersion, final Codec userCodec) { + return getCodec(knnCodecVersion, userCodec); + } + + private static Codec getCodec(final KNNCodecVersion knnCodecVersion, final Codec userCodec) { + try { + Constructor constructor = CODEC_BY_VERSION.getOrDefault(knnCodecVersion, CODEC_BY_VERSION.get(LATEST_KNN_CODEC_VERSION)) + .getConstructor(Codec.class); + return (Codec) constructor.newInstance(userCodec); + } catch (Exception ex) { + throw new RuntimeException("Cannot create instance of KNN codec", ex); + } + } + + public static class CodecDelegateFactory { + + public static Codec createKNN91DefaultDelegate() { + return new Lucene91Codec(); + } + } + + enum KNNCodecVersion { + KNN91 } } diff --git a/src/main/java/org/opensearch/knn/index/codec/KNNCodecService.java b/src/main/java/org/opensearch/knn/index/codec/KNNCodecService.java index 622f6ef17..cae6f7fb8 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNNCodecService.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNNCodecService.java @@ -26,6 +26,6 @@ public KNNCodecService(CodecServiceConfig codecServiceConfig) { */ @Override public Codec codec(String name) { - return KNNCodecFactory.createKNN91Codec(super.codec(name)); + return KNNCodecFactory.createKNNCodec(super.codec(name)); } } diff --git a/src/main/java/org/opensearch/knn/index/codec/KNNDocFormatFacade.java b/src/main/java/org/opensearch/knn/index/codec/KNNDocFormatFacade.java deleted file mode 100644 index d5b6b5ff1..000000000 --- a/src/main/java/org/opensearch/knn/index/codec/KNNDocFormatFacade.java +++ /dev/null @@ -1,22 +0,0 @@ -/* - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - * - * Modifications Copyright OpenSearch Contributors. See - * GitHub history for details. - */ - -package org.opensearch.knn.index.codec; - -import org.apache.lucene.codecs.CompoundFormat; -import org.apache.lucene.codecs.DocValuesFormat; - -public interface KNNDocFormatFacade { - - DocValuesFormat docValuesFormat(); - - CompoundFormat compoundFormat(); -} diff --git a/src/main/java/org/opensearch/knn/index/codec/KNNDocFormatFactory.java b/src/main/java/org/opensearch/knn/index/codec/KNNDocFormatFactory.java deleted file mode 100644 index e3a3d5018..000000000 --- a/src/main/java/org/opensearch/knn/index/codec/KNNDocFormatFactory.java +++ /dev/null @@ -1,31 +0,0 @@ -/* - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - * - * Modifications Copyright OpenSearch Contributors. See - * GitHub history for details. - */ - -package org.opensearch.knn.index.codec; - -import org.apache.lucene.codecs.Codec; -import org.opensearch.knn.index.codec.KNN91Codec.docformat.KNN91CompoundFormat; -import org.opensearch.knn.index.codec.KNN91Codec.KNN91DocFormat; -import org.opensearch.knn.index.codec.KNN91Codec.docformat.KNN91DocValuesFormat; - -/** - * Factory abstraction for KNN document format facades - */ -public class KNNDocFormatFactory { - - public static KNNDocFormatFacade createKNN91DocFormat(Codec delegate) { - final KNNDocFormatFacade knnDocFormatFacade = new KNN91DocFormat( - new KNN91DocValuesFormat(delegate.docValuesFormat()), - new KNN91CompoundFormat(delegate.compoundFormat()) - ); - return knnDocFormatFacade; - } -} diff --git a/src/main/java/org/opensearch/knn/index/codec/KNNFormatFacade.java b/src/main/java/org/opensearch/knn/index/codec/KNNFormatFacade.java new file mode 100644 index 000000000..bf9b1ad7f --- /dev/null +++ b/src/main/java/org/opensearch/knn/index/codec/KNNFormatFacade.java @@ -0,0 +1,30 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.knn.index.codec; + +import org.apache.lucene.codecs.CompoundFormat; +import org.apache.lucene.codecs.DocValuesFormat; + +/** + * Class abstracts facade for plugin formats. + */ +public class KNNFormatFacade { + + private final DocValuesFormat docValuesFormat; + private final CompoundFormat compoundFormat; + + public KNNFormatFacade(final DocValuesFormat docValuesFormat, final CompoundFormat compoundFormat) { + this.docValuesFormat = docValuesFormat; + this.compoundFormat = compoundFormat; + } + + public DocValuesFormat docValuesFormat() { + return docValuesFormat; + } + + public CompoundFormat compoundFormat() { + return compoundFormat; + } +} diff --git a/src/main/java/org/opensearch/knn/index/codec/KNNFormatFactory.java b/src/main/java/org/opensearch/knn/index/codec/KNNFormatFactory.java new file mode 100644 index 000000000..5f5a447c6 --- /dev/null +++ b/src/main/java/org/opensearch/knn/index/codec/KNNFormatFactory.java @@ -0,0 +1,23 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.knn.index.codec; + +import org.apache.lucene.codecs.Codec; +import org.opensearch.knn.index.codec.KNN91Codec.docformat.KNN91CompoundFormat; +import org.opensearch.knn.index.codec.KNN91Codec.docformat.KNN91DocValuesFormat; + +/** + * Factory abstraction for KNN document format facades + */ +public class KNNFormatFactory { + + public static KNNFormatFacade createKNN91Format(final Codec delegate) { + final KNNFormatFacade knnFormatFacade = new KNNFormatFacade( + new KNN91DocValuesFormat(delegate.docValuesFormat()), + new KNN91CompoundFormat(delegate.compoundFormat()) + ); + return knnFormatFacade; + } +} From ee411d5e95432a82b258c747f4863ee51653b08e Mon Sep 17 00:00:00 2001 From: Martin Gaievski Date: Wed, 23 Mar 2022 10:17:02 -0700 Subject: [PATCH 6/9] Rebase on main, switch to base 2.0 alpha1 snapshot using new approach Signed-off-by: Martin Gaievski --- .github/workflows/CI.yml | 2 +- build.gradle | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index f2595ee2b..9e838792b 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -35,7 +35,7 @@ jobs: - name: Run build run: | - ./gradlew build -Dopensearch.version=2.0.0-SNAPSHOT + ./gradlew build -Dopensearch.version=2.0.0-alpha1-SNAPSHOT - name: Run k-NN Backwards Compatibility Tests run: | diff --git a/build.gradle b/build.gradle index eb22c1e52..dc65c8c87 100644 --- a/build.gradle +++ b/build.gradle @@ -11,7 +11,7 @@ buildscript { ext { // build.version_qualifier parameter applies to knn plugin artifacts only. OpenSearch version must be set // explicitly as 'opensearch.version' property, for instance opensearch.version=2.0.0-alpha1-SNAPSHOT - opensearch_version = System.getProperty("opensearch.version", "2.0.0-SNAPSHOT") + opensearch_version = System.getProperty("opensearch.version", "2.0.0-alpha1-SNAPSHOT") knn_bwc_version = System.getProperty("bwc.version", "1.2.0.0-SNAPSHOT") version_qualifier = System.getProperty("build.version_qualifier", "alpha1") opensearch_bwc_version = "${knn_bwc_version}" - ".0-SNAPSHOT" From c62a0a9f6b8c65d37d97064c5055a5e0a7bb3c4d Mon Sep 17 00:00:00 2001 From: Martin Gaievski Date: Wed, 23 Mar 2022 11:57:59 -0700 Subject: [PATCH 7/9] Reverting logic of setting plugin mappings Signed-off-by: Martin Gaievski --- src/main/java/org/opensearch/knn/indices/ModelDao.java | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/src/main/java/org/opensearch/knn/indices/ModelDao.java b/src/main/java/org/opensearch/knn/indices/ModelDao.java index 6ab1b2cca..937e5f811 100644 --- a/src/main/java/org/opensearch/knn/indices/ModelDao.java +++ b/src/main/java/org/opensearch/knn/indices/ModelDao.java @@ -39,11 +39,8 @@ import org.opensearch.cluster.health.ClusterIndexHealth; import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.cluster.service.ClusterService; -import org.opensearch.common.Strings; import org.opensearch.common.settings.Settings; -import org.opensearch.common.xcontent.json.JsonXContent; import org.opensearch.index.IndexNotFoundException; -import org.opensearch.index.mapper.MapperService; import org.opensearch.knn.common.KNNConstants; import org.opensearch.knn.plugin.transport.DeleteModelResponse; import org.opensearch.knn.plugin.transport.GetModelResponse; @@ -201,10 +198,7 @@ public void create(ActionListener actionListener) throws IO if (isCreated()) { return; } - String mapping = Strings.toString( - JsonXContent.contentBuilder().startObject().startObject(MapperService.SINGLE_MAPPING_NAME).endObject().endObject() - ); - CreateIndexRequest request = new CreateIndexRequest(MODEL_INDEX_NAME).mapping(mapping) + CreateIndexRequest request = new CreateIndexRequest(MODEL_INDEX_NAME).mapping(getMapping()) .settings( Settings.builder() .put("index.hidden", true) From 7550e972666d74b98ec81fd726533eac4076701f Mon Sep 17 00:00:00 2001 From: Martin Gaievski Date: Wed, 23 Mar 2022 14:49:22 -0700 Subject: [PATCH 8/9] Adding unit tests, fix license headers Signed-off-by: Martin Gaievski --- .../codec/KNN87Codec/KNN87CodecTests.java | 13 ++----- .../docformat/KNN91BinaryDocValuesTests.java | 8 +---- .../docformat/KNN91CompoundFormatTests.java | 8 +---- .../KNN91DocValuesConsumerTests.java | 8 +---- .../knn/index/codec/KNNCodecFactoryTests.java | 34 +++++++++++++++++++ .../knn/index/codec/KNNCodecTestCase.java | 22 ++++++++++++ .../index/codec/KNNFormatFactoryTests.java | 24 +++++++++++++ 7 files changed, 85 insertions(+), 32 deletions(-) create mode 100644 src/test/java/org/opensearch/knn/index/codec/KNNCodecFactoryTests.java create mode 100644 src/test/java/org/opensearch/knn/index/codec/KNNFormatFactoryTests.java diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN87Codec/KNN87CodecTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN87Codec/KNN87CodecTests.java index 48c6e6c0d..ff93fd3a8 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNN87Codec/KNN87CodecTests.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNN87Codec/KNN87CodecTests.java @@ -5,20 +5,11 @@ package org.opensearch.knn.index.codec.KNN87Codec; -import org.junit.Ignore; import org.opensearch.knn.index.codec.KNNCodecTestCase; -import java.io.IOException; -import java.util.concurrent.ExecutionException; - -@Ignore public class KNN87CodecTests extends KNNCodecTestCase { - public void testMultiFieldsKnnIndex() throws Exception { - testMultiFieldsKnnIndex(new KNN87Codec()); - } - - public void testBuildFromModelTemplate() throws InterruptedException, ExecutionException, IOException { - testBuildFromModelTemplate(new KNN87Codec()); + public void testWriteByOldCodec() throws Exception { + testWriteByOldCodec(new KNN87Codec()); } } diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValuesTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValuesTests.java index 885a60e5d..8cb1934c6 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValuesTests.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValuesTests.java @@ -1,12 +1,6 @@ /* + * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - * - * Modifications Copyright OpenSearch Contributors. See - * GitHub history for details. */ package org.opensearch.knn.index.codec.KNN91Codec.docformat; diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormatTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormatTests.java index 567d0da86..04939603d 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormatTests.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormatTests.java @@ -1,12 +1,6 @@ /* + * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - * - * Modifications Copyright OpenSearch Contributors. See - * GitHub history for details. */ package org.opensearch.knn.index.codec.KNN91Codec.docformat; diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumerTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumerTests.java index a573f273a..3ba940997 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumerTests.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumerTests.java @@ -1,12 +1,6 @@ /* + * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - * - * Modifications Copyright OpenSearch Contributors. See - * GitHub history for details. */ package org.opensearch.knn.index.codec.KNN91Codec.docformat; diff --git a/src/test/java/org/opensearch/knn/index/codec/KNNCodecFactoryTests.java b/src/test/java/org/opensearch/knn/index/codec/KNNCodecFactoryTests.java new file mode 100644 index 000000000..241ae0e56 --- /dev/null +++ b/src/test/java/org/opensearch/knn/index/codec/KNNCodecFactoryTests.java @@ -0,0 +1,34 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.index.codec; + +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.lucene91.Lucene91Codec; +import org.opensearch.knn.KNNTestCase; +import org.opensearch.knn.index.codec.KNN91Codec.KNN91Codec; + +public class KNNCodecFactoryTests extends KNNTestCase { + + public void testKNN91DefaultDelegate() { + Codec knn91DefaultDelegate = KNNCodecFactory.CodecDelegateFactory.createKNN91DefaultDelegate(); + assertNotNull(knn91DefaultDelegate); + assertTrue(knn91DefaultDelegate instanceof Lucene91Codec); + } + + public void testKNN91DefaultCodec() { + Lucene91Codec lucene91CodecDelegate = new Lucene91Codec(); + Codec knnCodec = KNNCodecFactory.createKNNCodec(lucene91CodecDelegate); + assertNotNull(knnCodec); + assertTrue(knnCodec instanceof KNN91Codec); + } + + public void testKNN91CodecByVersion() { + Lucene91Codec lucene91CodecDelegate = new Lucene91Codec(); + Codec knnCodec = KNNCodecFactory.createKNNCodec(KNNCodecFactory.KNNCodecVersion.KNN91, lucene91CodecDelegate); + assertNotNull(knnCodec); + assertTrue(knnCodec instanceof KNN91Codec); + } +} diff --git a/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java b/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java index 47d13a87e..ad4558e0a 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java @@ -231,4 +231,26 @@ public void testBuildFromModelTemplate(Codec codec) throws IOException, Executio resourceWatcherService.close(); NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance().close(); } + + public void testWriteByOldCodec(Codec codec) throws IOException { + setUpMockClusterService(); + Directory dir = newFSDirectory(createTempDir()); + IndexWriterConfig iwc = newIndexWriterConfig(); + iwc.setMergeScheduler(new SerialMergeScheduler()); + iwc.setCodec(codec); + + /** + * Add doc with field "test_vector", expect it to fail + */ + float[] array = { 1.0f, 3.0f, 4.0f }; + VectorField vectorField = new VectorField("test_vector", array, sampleFieldType); + try (RandomIndexWriter writer = new RandomIndexWriter(random(), dir, iwc)) { + Document doc = new Document(); + doc.add(vectorField); + expectThrows(UnsupportedOperationException.class, () -> writer.addDocument(doc)); + } + + dir.close(); + NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance().close(); + } } diff --git a/src/test/java/org/opensearch/knn/index/codec/KNNFormatFactoryTests.java b/src/test/java/org/opensearch/knn/index/codec/KNNFormatFactoryTests.java new file mode 100644 index 000000000..778c3f951 --- /dev/null +++ b/src/test/java/org/opensearch/knn/index/codec/KNNFormatFactoryTests.java @@ -0,0 +1,24 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.index.codec; + +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.lucene91.Lucene91Codec; +import org.opensearch.knn.KNNTestCase; +import org.opensearch.knn.index.codec.KNN91Codec.KNN91Codec; + +public class KNNFormatFactoryTests extends KNNTestCase { + + public void testKNN91Format() { + final Codec lucene91CodecDelegate = KNNCodecFactory.CodecDelegateFactory.createKNN91DefaultDelegate(); + final Codec knnCodec = KNNCodecFactory.createKNNCodec(lucene91CodecDelegate); + KNNFormatFacade knnFormatFacade = KNNFormatFactory.createKNN91Format(knnCodec); + + assertNotNull(knnFormatFacade); + assertNotNull(knnFormatFacade.compoundFormat()); + assertNotNull(knnFormatFacade.docValuesFormat()); + } +} From 98f126473e3825db93480f653d6fa496c84126b4 Mon Sep 17 00:00:00 2001 From: Martin Gaievski Date: Wed, 23 Mar 2022 15:52:56 -0700 Subject: [PATCH 9/9] Switch to x.x.x versioning for codecs and doc value format Signed-off-by: Martin Gaievski --- .../KNN910Codec.java} | 16 +- .../docformat/KNN91BinaryDocValues.java | 63 --- .../docformat/KNN91CompoundFormat.java | 67 --- .../docformat/KNN91DocValuesConsumer.java | 256 ----------- .../docformat/KNN91DocValuesFormat.java | 40 -- .../docformat/KNN91DocValuesReader.java | 52 --- .../knn/index/codec/KNNCodecFactory.java | 16 +- .../knn/index/codec/KNNFormatFactory.java | 12 +- .../services/org.apache.lucene.codecs.Codec | 2 +- .../KNN910CodecTests.java} | 8 +- .../docformat/KNN91BinaryDocValuesTests.java | 69 --- .../docformat/KNN91CompoundFormatTests.java | 93 ---- .../KNN91DocValuesConsumerTests.java | 412 ------------------ .../knn/index/codec/KNNCodecFactoryTests.java | 8 +- .../knn/index/codec/KNNCodecTestCase.java | 4 +- .../index/codec/KNNFormatFactoryTests.java | 4 +- 16 files changed, 37 insertions(+), 1085 deletions(-) rename src/main/java/org/opensearch/knn/index/codec/{KNN91Codec/KNN91Codec.java => KNN910Codec/KNN910Codec.java} (75%) delete mode 100644 src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValues.java delete mode 100644 src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormat.java delete mode 100644 src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumer.java delete mode 100644 src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesFormat.java delete mode 100644 src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesReader.java rename src/test/java/org/opensearch/knn/index/codec/{KNN91Codec/KNN91CodecTests.java => KNN910Codec/KNN910CodecTests.java} (65%) delete mode 100644 src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValuesTests.java delete mode 100644 src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormatTests.java delete mode 100644 src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumerTests.java diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91Codec.java b/src/main/java/org/opensearch/knn/index/codec/KNN910Codec/KNN910Codec.java similarity index 75% rename from src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91Codec.java rename to src/main/java/org/opensearch/knn/index/codec/KNN910Codec/KNN910Codec.java index 5894ee5a1..0acaccfbf 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91Codec.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN910Codec/KNN910Codec.java @@ -2,7 +2,7 @@ * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.knn.index.codec.KNN91Codec; +package org.opensearch.knn.index.codec.KNN910Codec; import org.apache.lucene.codecs.Codec; import org.apache.lucene.codecs.CompoundFormat; @@ -18,15 +18,15 @@ * based on the mappings. * */ -public final class KNN91Codec extends FilterCodec { +public final class KNN910Codec extends FilterCodec { - public static final String KNN_91 = "KNN91Codec"; - private KNNFormatFacade knnFormatFacade; + private static final String KNN910 = "KNN910Codec"; + private final KNNFormatFacade knnFormatFacade; /** * No arg constructor that uses Lucene91 as the delegate */ - public KNN91Codec() { + public KNN910Codec() { this(createKNN91DefaultDelegate()); } @@ -35,9 +35,9 @@ public KNN91Codec() { * * @param delegate codec that will perform all operations this codec does not override */ - public KNN91Codec(Codec delegate) { - super(KNN_91, delegate); - knnFormatFacade = KNNFormatFactory.createKNN91Format(delegate); + public KNN910Codec(Codec delegate) { + super(KNN910, delegate); + knnFormatFacade = KNNFormatFactory.createKNN910Format(delegate); } @Override diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValues.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValues.java deleted file mode 100644 index b0d015e64..000000000 --- a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValues.java +++ /dev/null @@ -1,63 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ -package org.opensearch.knn.index.codec.KNN91Codec.docformat; - -import org.apache.lucene.index.BinaryDocValues; -import org.apache.lucene.index.DocIDMerger; -import org.apache.lucene.util.BytesRef; -import org.opensearch.knn.index.codec.util.BinaryDocValuesSub; - -import java.io.IOException; - -/** - * A per-document kNN numeric value. - */ -class KNN91BinaryDocValues extends BinaryDocValues { - - private DocIDMerger docIDMerger; - - KNN91BinaryDocValues(DocIDMerger docIdMerger) { - this.docIDMerger = docIdMerger; - } - - private BinaryDocValuesSub current; - private int docID = -1; - - @Override - public int docID() { - return docID; - } - - @Override - public int nextDoc() throws IOException { - current = docIDMerger.next(); - if (current == null) { - docID = NO_MORE_DOCS; - } else { - docID = current.mappedDocID; - } - return docID; - } - - @Override - public int advance(int target) throws IOException { - throw new UnsupportedOperationException(); - } - - @Override - public boolean advanceExact(int target) throws IOException { - throw new UnsupportedOperationException(); - } - - @Override - public long cost() { - throw new UnsupportedOperationException(); - } - - @Override - public BytesRef binaryValue() throws IOException { - return current.getValues().binaryValue(); - } -}; diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormat.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormat.java deleted file mode 100644 index ac3e59db2..000000000 --- a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormat.java +++ /dev/null @@ -1,67 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ -package org.opensearch.knn.index.codec.KNN91Codec.docformat; - -import org.apache.lucene.codecs.CompoundDirectory; -import org.apache.lucene.codecs.CompoundFormat; -import org.apache.lucene.index.SegmentInfo; -import org.apache.lucene.store.Directory; -import org.apache.lucene.store.IOContext; -import org.opensearch.knn.common.KNNConstants; -import org.opensearch.knn.index.util.KNNEngine; - -import java.io.IOException; -import java.util.HashSet; -import java.util.Set; -import java.util.stream.Collectors; - -/** - * Class to encode/decode compound file - */ -public class KNN91CompoundFormat extends CompoundFormat { - - private final CompoundFormat delegate; - - /** - * Constructor that takes a delegate to handle non-overridden methods - * - * @param delegate CompoundFormat that will handle non-overridden methods - */ - public KNN91CompoundFormat(CompoundFormat delegate) { - this.delegate = delegate; - } - - @Override - public CompoundDirectory getCompoundReader(Directory dir, SegmentInfo si, IOContext context) throws IOException { - return delegate.getCompoundReader(dir, si, context); - } - - @Override - public void write(Directory dir, SegmentInfo si, IOContext context) throws IOException { - for (KNNEngine knnEngine : KNNEngine.values()) { - writeEngineFiles(dir, si, context, knnEngine.getExtension()); - } - delegate.write(dir, si, context); - } - - private void writeEngineFiles(Directory dir, SegmentInfo si, IOContext context, String engineExtension) throws IOException { - /* - * If engine file present, remove it from the compounding file list to avoid header/footer checks - * and create a new compounding file format with extension engine + c. - */ - Set engineFiles = si.files().stream().filter(file -> file.endsWith(engineExtension)).collect(Collectors.toSet()); - - Set segmentFiles = new HashSet<>(si.files()); - - if (!engineFiles.isEmpty()) { - for (String engineFile : engineFiles) { - String engineCompoundFile = engineFile + KNNConstants.COMPOUND_EXTENSION; - dir.copyFrom(dir, engineFile, engineCompoundFile, context); - } - segmentFiles.removeAll(engineFiles); - si.setFiles(segmentFiles); - } - } -} diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumer.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumer.java deleted file mode 100644 index 5dc187346..000000000 --- a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumer.java +++ /dev/null @@ -1,256 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ -package org.opensearch.knn.index.codec.KNN91Codec.docformat; - -import com.google.common.collect.ImmutableMap; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.apache.lucene.codecs.CodecUtil; -import org.apache.lucene.codecs.DocValuesConsumer; -import org.apache.lucene.codecs.DocValuesProducer; -import org.apache.lucene.index.BinaryDocValues; -import org.apache.lucene.index.DocValuesType; -import org.apache.lucene.index.FieldInfo; -import org.apache.lucene.index.MergeState; -import org.apache.lucene.index.SegmentWriteState; -import org.apache.lucene.store.FSDirectory; -import org.apache.lucene.store.FilterDirectory; -import org.apache.lucene.store.IndexInput; -import org.apache.lucene.store.IndexOutput; -import org.apache.lucene.util.IOUtils; -import org.opensearch.common.xcontent.DeprecationHandler; -import org.opensearch.common.xcontent.NamedXContentRegistry; -import org.opensearch.common.xcontent.XContentFactory; -import org.opensearch.common.xcontent.XContentType; -import org.opensearch.knn.common.KNNConstants; -import org.opensearch.knn.index.KNNSettings; -import org.opensearch.knn.index.KNNVectorFieldMapper; -import org.opensearch.knn.index.SpaceType; -import org.opensearch.knn.index.codec.util.KNNCodecUtil; -import org.opensearch.knn.index.util.KNNEngine; -import org.opensearch.knn.indices.Model; -import org.opensearch.knn.indices.ModelCache; -import org.opensearch.knn.jni.JNIService; -import org.opensearch.knn.plugin.stats.KNNCounter; - -import java.io.Closeable; -import java.io.IOException; -import java.nio.file.Paths; -import java.security.AccessController; -import java.security.PrivilegedAction; -import java.util.HashMap; -import java.util.Map; - -import static org.opensearch.knn.common.KNNConstants.MODEL_ID; -import static org.opensearch.knn.common.KNNConstants.PARAMETERS; -import static org.opensearch.knn.index.codec.util.KNNCodecUtil.buildEngineFileName; - -/** - * This class writes the KNN docvalues to the segments - */ -class KNN91DocValuesConsumer extends DocValuesConsumer implements Closeable { - - private final Logger logger = LogManager.getLogger(KNN91DocValuesConsumer.class); - - private final String TEMP_SUFFIX = "tmp"; - private DocValuesConsumer delegatee; - private SegmentWriteState state; - - KNN91DocValuesConsumer(DocValuesConsumer delegatee, SegmentWriteState state) throws IOException { - this.delegatee = delegatee; - this.state = state; - } - - @Override - public void addBinaryField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException { - delegatee.addBinaryField(field, valuesProducer); - if (field.attributes().containsKey(KNNVectorFieldMapper.KNN_FIELD)) { - addKNNBinaryField(field, valuesProducer); - } - } - - public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException { - - // Get values to be indexed - BinaryDocValues values = valuesProducer.getBinary(field); - KNNCodecUtil.Pair pair = KNNCodecUtil.getFloats(values); - if (pair.vectors.length == 0 || pair.docs.length == 0) { - logger.info("Skipping engine index creation as there are no vectors or docs in the documents"); - return; - } - - // Increment counter for number of graph index requests - KNNCounter.GRAPH_INDEX_REQUESTS.increment(); - - // Create library index either from model or from scratch - String engineFileName; - String indexPath; - String tmpEngineFileName; - - if (field.attributes().containsKey(MODEL_ID)) { - - String modelId = field.attributes().get(MODEL_ID); - Model model = ModelCache.getInstance().get(modelId); - - KNNEngine knnEngine = model.getModelMetadata().getKnnEngine(); - - engineFileName = buildEngineFileName( - state.segmentInfo.name, - knnEngine.getLatestBuildVersion(), - field.name, - knnEngine.getExtension() - ); - indexPath = Paths.get(((FSDirectory) (FilterDirectory.unwrap(state.directory))).getDirectory().toString(), engineFileName) - .toString(); - tmpEngineFileName = engineFileName + TEMP_SUFFIX; - String tempIndexPath = indexPath + TEMP_SUFFIX; - - if (model.getModelBlob() == null) { - throw new RuntimeException("There is no trained model with id \"" + modelId + "\""); - } - - createKNNIndexFromTemplate(model.getModelBlob(), pair, knnEngine, tempIndexPath); - } else { - - // Get engine to be used for indexing - String engineName = field.attributes().getOrDefault(KNNConstants.KNN_ENGINE, KNNEngine.DEFAULT.getName()); - KNNEngine knnEngine = KNNEngine.getEngine(engineName); - - engineFileName = buildEngineFileName( - state.segmentInfo.name, - knnEngine.getLatestBuildVersion(), - field.name, - knnEngine.getExtension() - ); - indexPath = Paths.get(((FSDirectory) (FilterDirectory.unwrap(state.directory))).getDirectory().toString(), engineFileName) - .toString(); - tmpEngineFileName = engineFileName + TEMP_SUFFIX; - String tempIndexPath = indexPath + TEMP_SUFFIX; - - createKNNIndexFromScratch(field, pair, knnEngine, tempIndexPath); - } - - /* - * Adds Footer to the serialized graph - * 1. Copies the serialized graph to new file. - * 2. Adds Footer to the new file. - * - * We had to create new file here because adding footer directly to the - * existing file will miss calculating checksum for the serialized graph - * bytes and result in index corruption issues. - */ - // TODO: I think this can be refactored to avoid this copy and then write - // https://github.com/opendistro-for-elasticsearch/k-NN/issues/330 - try ( - IndexInput is = state.directory.openInput(tmpEngineFileName, state.context); - IndexOutput os = state.directory.createOutput(engineFileName, state.context) - ) { - os.copyBytes(is, is.length()); - CodecUtil.writeFooter(os); - } catch (Exception ex) { - KNNCounter.GRAPH_INDEX_ERRORS.increment(); - throw new RuntimeException("[KNN] Adding footer to serialized graph failed: " + ex); - } finally { - IOUtils.deleteFilesIgnoringExceptions(state.directory, tmpEngineFileName); - } - } - - private void createKNNIndexFromTemplate(byte[] model, KNNCodecUtil.Pair pair, KNNEngine knnEngine, String indexPath) { - Map parameters = ImmutableMap.of( - KNNConstants.INDEX_THREAD_QTY, - KNNSettings.state().getSettingValue(KNNSettings.KNN_ALGO_PARAM_INDEX_THREAD_QTY) - ); - AccessController.doPrivileged((PrivilegedAction) () -> { - JNIService.createIndexFromTemplate(pair.docs, pair.vectors, indexPath, model, parameters, knnEngine.getName()); - return null; - }); - } - - private void createKNNIndexFromScratch(FieldInfo fieldInfo, KNNCodecUtil.Pair pair, KNNEngine knnEngine, String indexPath) - throws IOException { - Map parameters = new HashMap<>(); - Map fieldAttributes = fieldInfo.attributes(); - String parametersString = fieldAttributes.get(KNNConstants.PARAMETERS); - - // parametersString will be null when legacy mapper is used - if (parametersString == null) { - parameters.put(KNNConstants.SPACE_TYPE, fieldAttributes.getOrDefault(KNNConstants.SPACE_TYPE, SpaceType.DEFAULT.getValue())); - - String efConstruction = fieldAttributes.get(KNNConstants.HNSW_ALGO_EF_CONSTRUCTION); - Map algoParams = new HashMap<>(); - if (efConstruction != null) { - algoParams.put(KNNConstants.METHOD_PARAMETER_EF_CONSTRUCTION, Integer.parseInt(efConstruction)); - } - - String m = fieldAttributes.get(KNNConstants.HNSW_ALGO_M); - if (m != null) { - algoParams.put(KNNConstants.METHOD_PARAMETER_M, Integer.parseInt(m)); - } - parameters.put(PARAMETERS, algoParams); - } else { - parameters.putAll( - XContentFactory.xContent(XContentType.JSON) - .createParser(NamedXContentRegistry.EMPTY, DeprecationHandler.THROW_UNSUPPORTED_OPERATION, parametersString) - .map() - ); - } - - // Used to determine how many threads to use when indexing - parameters.put(KNNConstants.INDEX_THREAD_QTY, KNNSettings.state().getSettingValue(KNNSettings.KNN_ALGO_PARAM_INDEX_THREAD_QTY)); - - // Pass the path for the nms library to save the file - AccessController.doPrivileged((PrivilegedAction) () -> { - JNIService.createIndex(pair.docs, pair.vectors, indexPath, parameters, knnEngine.getName()); - return null; - }); - } - - /** - * Merges in the fields from the readers in mergeState - * - * @param mergeState Holds common state used during segment merging - */ - @Override - public void merge(MergeState mergeState) { - try { - delegatee.merge(mergeState); - assert mergeState != null; - assert mergeState.mergeFieldInfos != null; - for (FieldInfo fieldInfo : mergeState.mergeFieldInfos) { - DocValuesType type = fieldInfo.getDocValuesType(); - if (type == DocValuesType.BINARY && fieldInfo.attributes().containsKey(KNNVectorFieldMapper.KNN_FIELD)) { - addKNNBinaryField(fieldInfo, new KNN91DocValuesReader(mergeState)); - } - } - } catch (Exception e) { - throw new RuntimeException(e); - } - } - - @Override - public void addSortedSetField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException { - delegatee.addSortedSetField(field, valuesProducer); - } - - @Override - public void addSortedNumericField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException { - delegatee.addSortedNumericField(field, valuesProducer); - } - - @Override - public void addSortedField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException { - delegatee.addSortedField(field, valuesProducer); - } - - @Override - public void addNumericField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException { - delegatee.addNumericField(field, valuesProducer); - } - - @Override - public void close() throws IOException { - delegatee.close(); - } -} diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesFormat.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesFormat.java deleted file mode 100644 index 3e24af87b..000000000 --- a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesFormat.java +++ /dev/null @@ -1,40 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ -package org.opensearch.knn.index.codec.KNN91Codec.docformat; - -import org.apache.lucene.codecs.DocValuesConsumer; -import org.apache.lucene.codecs.DocValuesFormat; -import org.apache.lucene.codecs.DocValuesProducer; -import org.apache.lucene.index.SegmentReadState; -import org.apache.lucene.index.SegmentWriteState; - -import java.io.IOException; - -/** - * Encodes/Decodes per document values - */ -public class KNN91DocValuesFormat extends DocValuesFormat { - private final DocValuesFormat delegate; - - /** - * Constructor that takes delegate in order to handle non-overridden methods - * - * @param delegate DocValuesFormat to handle non-overridden methods - */ - public KNN91DocValuesFormat(DocValuesFormat delegate) { - super(delegate.getName()); - this.delegate = delegate; - } - - @Override - public DocValuesConsumer fieldsConsumer(SegmentWriteState state) throws IOException { - return new KNN91DocValuesConsumer(delegate.fieldsConsumer(state), state); - } - - @Override - public DocValuesProducer fieldsProducer(SegmentReadState state) throws IOException { - return delegate.fieldsProducer(state); - } -} diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesReader.java b/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesReader.java deleted file mode 100644 index d870710df..000000000 --- a/src/main/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesReader.java +++ /dev/null @@ -1,52 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ -package org.opensearch.knn.index.codec.KNN91Codec.docformat; - -import org.apache.lucene.codecs.DocValuesProducer; -import org.apache.lucene.index.BinaryDocValues; -import org.apache.lucene.index.DocIDMerger; -import org.apache.lucene.index.DocValuesType; -import org.apache.lucene.index.EmptyDocValuesProducer; -import org.apache.lucene.index.FieldInfo; -import org.apache.lucene.index.MergeState; -import org.opensearch.knn.index.codec.util.BinaryDocValuesSub; - -import java.util.ArrayList; -import java.util.List; - -/** - * Reader for KNNDocValues from the segments - */ -class KNN91DocValuesReader extends EmptyDocValuesProducer { - - private final MergeState mergeState; - - KNN91DocValuesReader(MergeState mergeState) { - this.mergeState = mergeState; - } - - @Override - public BinaryDocValues getBinary(FieldInfo field) { - try { - List subs = new ArrayList<>(this.mergeState.docValuesProducers.length); - for (int i = 0; i < this.mergeState.docValuesProducers.length; i++) { - BinaryDocValues values = null; - DocValuesProducer docValuesProducer = mergeState.docValuesProducers[i]; - if (docValuesProducer != null) { - FieldInfo readerFieldInfo = mergeState.fieldInfos[i].fieldInfo(field.name); - if (readerFieldInfo != null && readerFieldInfo.getDocValuesType() == DocValuesType.BINARY) { - values = docValuesProducer.getBinary(readerFieldInfo); - } - if (values != null) { - subs.add(new BinaryDocValuesSub(mergeState.docMaps[i], values)); - } - } - } - return new KNN91BinaryDocValues(DocIDMerger.of(subs, mergeState.needsIndexSort)); - } catch (Exception e) { - throw new RuntimeException(e); - } - } -} diff --git a/src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java b/src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java index 68c4b879d..9662e9edb 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java @@ -7,19 +7,19 @@ import com.google.common.collect.ImmutableMap; import org.apache.lucene.codecs.Codec; import org.apache.lucene.codecs.lucene91.Lucene91Codec; -import org.opensearch.knn.index.codec.KNN91Codec.KNN91Codec; +import org.opensearch.knn.index.codec.KNN910Codec.KNN910Codec; import java.lang.reflect.Constructor; import java.util.Map; /** - * Factory abstraction for KNN codes + * Factory abstraction for KNN codec */ public class KNNCodecFactory { - private static Map CODEC_BY_VERSION = ImmutableMap.of(KNNCodecVersion.KNN91, KNN91Codec.class); + private static Map CODEC_BY_VERSION = ImmutableMap.of(KNNCodecVersion.KNN910, KNN910Codec.class); - private static KNNCodecVersion LATEST_KNN_CODEC_VERSION = KNNCodecVersion.KNN91; + private static KNNCodecVersion LATEST_KNN_CODEC_VERSION = KNNCodecVersion.KNN910; public static Codec createKNNCodec(final Codec userCodec) { return getCodec(LATEST_KNN_CODEC_VERSION, userCodec); @@ -39,6 +39,9 @@ private static Codec getCodec(final KNNCodecVersion knnCodecVersion, final Codec } } + /** + * Factory abstraction for codec delegate + */ public static class CodecDelegateFactory { public static Codec createKNN91DefaultDelegate() { @@ -46,7 +49,10 @@ public static Codec createKNN91DefaultDelegate() { } } + /** + * Collection of supported coded versions + */ enum KNNCodecVersion { - KNN91 + KNN910 } } diff --git a/src/main/java/org/opensearch/knn/index/codec/KNNFormatFactory.java b/src/main/java/org/opensearch/knn/index/codec/KNNFormatFactory.java index 5f5a447c6..6742dc3f4 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNNFormatFactory.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNNFormatFactory.java @@ -5,18 +5,18 @@ package org.opensearch.knn.index.codec; import org.apache.lucene.codecs.Codec; -import org.opensearch.knn.index.codec.KNN91Codec.docformat.KNN91CompoundFormat; -import org.opensearch.knn.index.codec.KNN91Codec.docformat.KNN91DocValuesFormat; +import org.opensearch.knn.index.codec.KNN80Codec.KNN80CompoundFormat; +import org.opensearch.knn.index.codec.KNN80Codec.KNN80DocValuesFormat; /** - * Factory abstraction for KNN document format facades + * Factory abstraction for KNN format facade creation */ public class KNNFormatFactory { - public static KNNFormatFacade createKNN91Format(final Codec delegate) { + public static KNNFormatFacade createKNN910Format(final Codec delegate) { final KNNFormatFacade knnFormatFacade = new KNNFormatFacade( - new KNN91DocValuesFormat(delegate.docValuesFormat()), - new KNN91CompoundFormat(delegate.compoundFormat()) + new KNN80DocValuesFormat(delegate.docValuesFormat()), + new KNN80CompoundFormat(delegate.compoundFormat()) ); return knnFormatFacade; } diff --git a/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec b/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec index 98f8a6a13..b897dc36a 100644 --- a/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec +++ b/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec @@ -2,4 +2,4 @@ org.opensearch.knn.index.codec.KNN80Codec.KNN80Codec org.opensearch.knn.index.codec.KNN84Codec.KNN84Codec org.opensearch.knn.index.codec.KNN86Codec.KNN86Codec org.opensearch.knn.index.codec.KNN87Codec.KNN87Codec -org.opensearch.knn.index.codec.KNN91Codec.KNN91Codec \ No newline at end of file +org.opensearch.knn.index.codec.KNN910Codec.KNN910Codec \ No newline at end of file diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CodecTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN910Codec/KNN910CodecTests.java similarity index 65% rename from src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CodecTests.java rename to src/test/java/org/opensearch/knn/index/codec/KNN910Codec/KNN910CodecTests.java index 39dc641dc..1ec28d6a4 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CodecTests.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNN910Codec/KNN910CodecTests.java @@ -3,20 +3,20 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.knn.index.codec.KNN91Codec; +package org.opensearch.knn.index.codec.KNN910Codec; import org.opensearch.knn.index.codec.KNNCodecTestCase; import java.io.IOException; import java.util.concurrent.ExecutionException; -public class KNN91CodecTests extends KNNCodecTestCase { +public class KNN910CodecTests extends KNNCodecTestCase { public void testMultiFieldsKnnIndex() throws Exception { - testMultiFieldsKnnIndex(new KNN91Codec()); + testMultiFieldsKnnIndex(new KNN910Codec()); } public void testBuildFromModelTemplate() throws InterruptedException, ExecutionException, IOException { - testBuildFromModelTemplate(new KNN91Codec()); + testBuildFromModelTemplate(new KNN910Codec()); } } diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValuesTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValuesTests.java deleted file mode 100644 index 8cb1934c6..000000000 --- a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91BinaryDocValuesTests.java +++ /dev/null @@ -1,69 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.knn.index.codec.KNN91Codec.docformat; - -import com.google.common.collect.ImmutableList; -import org.apache.lucene.index.BinaryDocValues; -import org.apache.lucene.index.DocIDMerger; -import org.apache.lucene.index.MergeState; -import org.opensearch.knn.KNNTestCase; -import org.opensearch.knn.index.codec.KNNCodecTestUtil; -import org.opensearch.knn.index.codec.util.BinaryDocValuesSub; - -import java.io.IOException; - -public class KNN91BinaryDocValuesTests extends KNNTestCase { - - public void testDocId() { - KNN91BinaryDocValues knn91BinaryDocValues = new KNN91BinaryDocValues(null); - assertEquals(-1, knn91BinaryDocValues.docID()); - } - - public void testNextDoc() throws IOException { - final int expectedDoc = 12; - - BinaryDocValuesSub sub = new BinaryDocValuesSub(new MergeState.DocMap() { - @Override - public int get(int docID) { - return expectedDoc; - } - }, new KNNCodecTestUtil.ConstantVectorBinaryDocValues(10, 128, 1.0f)); - - DocIDMerger docIDMerger = DocIDMerger.of(ImmutableList.of(sub), false); - KNN91BinaryDocValues knn91BinaryDocValues = new KNN91BinaryDocValues(docIDMerger); - assertEquals(expectedDoc, knn91BinaryDocValues.nextDoc()); - } - - public void testAdvance() { - KNN91BinaryDocValues knn91BinaryDocValues = new KNN91BinaryDocValues(null); - expectThrows(UnsupportedOperationException.class, () -> knn91BinaryDocValues.advance(0)); - } - - public void testAdvanceExact() { - KNN91BinaryDocValues knn91BinaryDocValues = new KNN91BinaryDocValues(null); - expectThrows(UnsupportedOperationException.class, () -> knn91BinaryDocValues.advanceExact(0)); - } - - public void testCost() { - KNN91BinaryDocValues knn91BinaryDocValues = new KNN91BinaryDocValues(null); - expectThrows(UnsupportedOperationException.class, knn91BinaryDocValues::cost); - } - - public void testBinaryValue() throws IOException { - BinaryDocValues binaryDocValues = new KNNCodecTestUtil.ConstantVectorBinaryDocValues(10, 128, 1.0f); - BinaryDocValuesSub sub = new BinaryDocValuesSub(new MergeState.DocMap() { - @Override - public int get(int docID) { - return docID; - } - }, binaryDocValues); - - DocIDMerger docIDMerger = DocIDMerger.of(ImmutableList.of(sub), false); - KNN91BinaryDocValues knn91BinaryDocValues = new KNN91BinaryDocValues(docIDMerger); - knn91BinaryDocValues.nextDoc(); - assertEquals(binaryDocValues.binaryValue(), knn91BinaryDocValues.binaryValue()); - } -} diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormatTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormatTests.java deleted file mode 100644 index 04939603d..000000000 --- a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91CompoundFormatTests.java +++ /dev/null @@ -1,93 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.knn.index.codec.KNN91Codec.docformat; - -import org.apache.lucene.codecs.Codec; -import org.apache.lucene.codecs.CompoundDirectory; -import org.apache.lucene.codecs.CompoundFormat; -import org.apache.lucene.index.SegmentInfo; -import org.apache.lucene.store.Directory; -import org.apache.lucene.store.IOContext; -import org.apache.lucene.store.IndexOutput; -import org.junit.AfterClass; -import org.junit.BeforeClass; -import org.opensearch.common.util.set.Sets; -import org.opensearch.knn.KNNTestCase; -import org.opensearch.knn.index.codec.KNN91Codec.KNN91Codec; -import org.opensearch.knn.index.codec.KNNCodecTestUtil; -import org.opensearch.knn.index.util.KNNEngine; - -import java.io.IOException; -import java.util.Arrays; -import java.util.Set; - -import static org.mockito.Mockito.doNothing; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - -public class KNN91CompoundFormatTests extends KNNTestCase { - - private static Directory directory; - private static Codec codec; - - @BeforeClass - public static void setStaticVariables() { - directory = newFSDirectory(createTempDir()); - codec = new KNN91Codec(); - } - - @AfterClass - public static void closeStaticVariables() throws IOException { - directory.close(); - } - - public void testGetCompoundReader() throws IOException { - CompoundDirectory dir = mock(CompoundDirectory.class); - CompoundFormat delegate = mock(CompoundFormat.class); - when(delegate.getCompoundReader(null, null, null)).thenReturn(dir); - KNN91CompoundFormat knn91CompoundFormat = new KNN91CompoundFormat(delegate); - assertEquals(dir, knn91CompoundFormat.getCompoundReader(null, null, null)); - } - - public void testWrite() throws IOException { - // Check that all normal engine files correctly get set to compound extension files after write - String segmentName = "_test"; - - Set segmentFiles = Sets.newHashSet( - String.format("%s_nmslib1%s", segmentName, KNNEngine.NMSLIB.getExtension()), - String.format("%s_nmslib2%s", segmentName, KNNEngine.NMSLIB.getExtension()), - String.format("%s_nmslib3%s", segmentName, KNNEngine.NMSLIB.getExtension()), - String.format("%s_faiss1%s", segmentName, KNNEngine.FAISS.getExtension()), - String.format("%s_faiss2%s", segmentName, KNNEngine.FAISS.getExtension()), - String.format("%s_faiss3%s", segmentName, KNNEngine.FAISS.getExtension()) - ); - - SegmentInfo segmentInfo = KNNCodecTestUtil.SegmentInfoBuilder.builder(directory, segmentName, segmentFiles.size(), codec).build(); - - for (String name : segmentFiles) { - IndexOutput indexOutput = directory.createOutput(name, IOContext.DEFAULT); - indexOutput.close(); - } - segmentInfo.setFiles(segmentFiles); - - CompoundFormat delegate = mock(CompoundFormat.class); - doNothing().when(delegate).write(directory, segmentInfo, IOContext.DEFAULT); - - KNN91CompoundFormat knn91CompoundFormat = new KNN91CompoundFormat(delegate); - knn91CompoundFormat.write(directory, segmentInfo, IOContext.DEFAULT); - - assertTrue(segmentInfo.files().isEmpty()); - - Arrays.stream(directory.listAll()).forEach(filename -> { - try { - directory.deleteFile(filename); - } catch (IOException e) { - fail(String.format("Failed to delete: %s", filename)); - } - }); - } - -} diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumerTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumerTests.java deleted file mode 100644 index 3ba940997..000000000 --- a/src/test/java/org/opensearch/knn/index/codec/KNN91Codec/docformat/KNN91DocValuesConsumerTests.java +++ /dev/null @@ -1,412 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.knn.index.codec.KNN91Codec.docformat; - -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.ImmutableSet; -import org.apache.lucene.codecs.Codec; -import org.apache.lucene.codecs.DocValuesConsumer; -import org.apache.lucene.codecs.DocValuesProducer; -import org.apache.lucene.index.FieldInfo; -import org.apache.lucene.index.FieldInfos; -import org.apache.lucene.index.SegmentInfo; -import org.apache.lucene.index.SegmentWriteState; -import org.apache.lucene.store.Directory; -import org.apache.lucene.store.IOContext; -import org.junit.AfterClass; -import org.junit.BeforeClass; -import org.opensearch.cluster.service.ClusterService; -import org.opensearch.common.Strings; -import org.opensearch.common.settings.ClusterSettings; -import org.opensearch.common.settings.Settings; -import org.opensearch.common.xcontent.XContentFactory; -import org.opensearch.knn.KNNTestCase; -import org.opensearch.knn.common.KNNConstants; -import org.opensearch.knn.index.KNNMethodContext; -import org.opensearch.knn.index.KNNVectorFieldMapper; -import org.opensearch.knn.index.MethodComponentContext; -import org.opensearch.knn.index.SpaceType; -import org.opensearch.knn.index.codec.KNN87Codec.KNN87Codec; -import org.opensearch.knn.index.codec.KNNCodecTestUtil; -import org.opensearch.knn.index.codec.util.KNNCodecUtil; -import org.opensearch.knn.index.util.KNNEngine; -import org.opensearch.knn.indices.Model; -import org.opensearch.knn.indices.ModelCache; -import org.opensearch.knn.indices.ModelDao; -import org.opensearch.knn.indices.ModelMetadata; -import org.opensearch.knn.indices.ModelState; -import org.opensearch.knn.jni.JNIService; -import org.opensearch.knn.plugin.stats.KNNCounter; - -import java.io.IOException; -import java.util.Map; -import java.util.concurrent.ExecutionException; - -import static org.mockito.Mockito.doNothing; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; -import static org.opensearch.knn.common.KNNConstants.INDEX_DESCRIPTION_PARAMETER; -import static org.opensearch.knn.common.KNNConstants.METHOD_HNSW; -import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_EF_CONSTRUCTION; -import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_M; -import static org.opensearch.knn.common.KNNConstants.MODEL_ID; -import static org.opensearch.knn.index.KNNSettings.MODEL_CACHE_SIZE_LIMIT_SETTING; -import static org.opensearch.knn.index.codec.KNNCodecTestUtil.RandomVectorDocValuesProducer; -import static org.opensearch.knn.index.codec.KNNCodecTestUtil.assertFileInCorrectLocation; -import static org.opensearch.knn.index.codec.KNNCodecTestUtil.assertLoadableByEngine; -import static org.opensearch.knn.index.codec.KNNCodecTestUtil.assertValidFooter; -import static org.opensearch.knn.index.codec.KNNCodecTestUtil.getRandomVectors; - -public class KNN91DocValuesConsumerTests extends KNNTestCase { - - private static Directory directory; - private static Codec codec; - - @BeforeClass - public static void setStaticVariables() { - directory = newFSDirectory(createTempDir()); - codec = new KNN87Codec(); - } - - @AfterClass - public static void closeStaticVariables() throws IOException { - directory.close(); - } - - public void testAddBinaryField_withKNN() throws IOException { - // Confirm that addKNNBinaryField will get called if the k-NN parameter is true - FieldInfo fieldInfo = KNNCodecTestUtil.FieldInfoBuilder.builder("test-field") - .addAttribute(KNNVectorFieldMapper.KNN_FIELD, "true") - .build(); - DocValuesProducer docValuesProducer = mock(DocValuesProducer.class); - - DocValuesConsumer delegate = mock(DocValuesConsumer.class); - doNothing().when(delegate).addBinaryField(fieldInfo, docValuesProducer); - - final boolean[] called = { false }; - KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(delegate, null) { - - @Override - public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer) { - called[0] = true; - } - }; - - knn91DocValuesConsumer.addBinaryField(fieldInfo, docValuesProducer); - - verify(delegate, times(1)).addBinaryField(fieldInfo, docValuesProducer); - assertTrue(called[0]); - } - - public void testAddBinaryField_withoutKNN() throws IOException { - // Confirm that the KNN91DocValuesConsumer will just call delegate AddBinaryField when k-NN parameter is - // not set - FieldInfo fieldInfo = KNNCodecTestUtil.FieldInfoBuilder.builder("test-field").build(); - DocValuesProducer docValuesProducer = mock(DocValuesProducer.class); - - DocValuesConsumer delegate = mock(DocValuesConsumer.class); - doNothing().when(delegate).addBinaryField(fieldInfo, docValuesProducer); - - final boolean[] called = { false }; - KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(delegate, null) { - - @Override - public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer) { - called[0] = true; - } - }; - - knn91DocValuesConsumer.addBinaryField(fieldInfo, docValuesProducer); - - verify(delegate, times(1)).addBinaryField(fieldInfo, docValuesProducer); - assertFalse(called[0]); - } - - public void testAddKNNBinaryField_noVectors() throws IOException { - // When there are no new vectors, no more graph index requests should be added - RandomVectorDocValuesProducer randomVectorDocValuesProducer = new RandomVectorDocValuesProducer(0, 128); - Long initialGraphIndexRequests = KNNCounter.GRAPH_INDEX_REQUESTS.getCount(); - KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(null, null); - knn91DocValuesConsumer.addKNNBinaryField(null, randomVectorDocValuesProducer); - assertEquals(initialGraphIndexRequests, KNNCounter.GRAPH_INDEX_REQUESTS.getCount()); - } - - public void testAddKNNBinaryField_fromScratch_nmslibCurrent() throws IOException { - // Set information about the segment and the fields - String segmentName = String.format("test_segment%s", randomAlphaOfLength(4)); - int docsInSegment = 100; - String fieldName = String.format("test_field%s", randomAlphaOfLength(4)); - - KNNEngine knnEngine = KNNEngine.NMSLIB; - SpaceType spaceType = SpaceType.COSINESIMIL; - int dimension = 16; - - SegmentInfo segmentInfo = KNNCodecTestUtil.SegmentInfoBuilder.builder(directory, segmentName, docsInSegment, codec).build(); - - KNNMethodContext knnMethodContext = new KNNMethodContext( - knnEngine, - spaceType, - new MethodComponentContext(METHOD_HNSW, ImmutableMap.of(METHOD_PARAMETER_M, 16, METHOD_PARAMETER_EF_CONSTRUCTION, 512)) - ); - - String parameterString = Strings.toString(XContentFactory.jsonBuilder().map(knnEngine.getMethodAsMap(knnMethodContext))); - - FieldInfo[] fieldInfoArray = new FieldInfo[] { - KNNCodecTestUtil.FieldInfoBuilder.builder(fieldName) - .addAttribute(KNNVectorFieldMapper.KNN_FIELD, "true") - .addAttribute(KNNConstants.KNN_ENGINE, knnEngine.getName()) - .addAttribute(KNNConstants.SPACE_TYPE, spaceType.getValue()) - .addAttribute(KNNConstants.PARAMETERS, parameterString) - .build() }; - - FieldInfos fieldInfos = new FieldInfos(fieldInfoArray); - SegmentWriteState state = new SegmentWriteState(null, directory, segmentInfo, fieldInfos, null, IOContext.DEFAULT); - - // Add documents to the field - KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(null, state); - RandomVectorDocValuesProducer randomVectorDocValuesProducer = new RandomVectorDocValuesProducer(docsInSegment, dimension); - knn91DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer); - - // The document should be created in the correct location - String expectedFile = KNNCodecUtil.buildEngineFileName( - segmentName, - knnEngine.getLatestBuildVersion(), - fieldName, - knnEngine.getExtension() - ); - assertFileInCorrectLocation(state, expectedFile); - - // The footer should be valid - assertValidFooter(state.directory, expectedFile); - - // The document should be readable by nmslib - assertLoadableByEngine(state, expectedFile, knnEngine, spaceType, dimension); - } - - public void testAddKNNBinaryField_fromScratch_nmslibLegacy() throws IOException { - // Set information about the segment and the fields - String segmentName = String.format("test_segment%s", randomAlphaOfLength(4)); - int docsInSegment = 100; - String fieldName = String.format("test_field%s", randomAlphaOfLength(4)); - - KNNEngine knnEngine = KNNEngine.NMSLIB; - SpaceType spaceType = SpaceType.COSINESIMIL; - int dimension = 16; - - SegmentInfo segmentInfo = KNNCodecTestUtil.SegmentInfoBuilder.builder(directory, segmentName, docsInSegment, codec).build(); - - FieldInfo[] fieldInfoArray = new FieldInfo[] { - KNNCodecTestUtil.FieldInfoBuilder.builder(fieldName) - .addAttribute(KNNVectorFieldMapper.KNN_FIELD, "true") - .addAttribute(KNNConstants.HNSW_ALGO_EF_CONSTRUCTION, "512") - .addAttribute(KNNConstants.HNSW_ALGO_M, "16") - .addAttribute(KNNConstants.SPACE_TYPE, spaceType.getValue()) - .build() }; - - FieldInfos fieldInfos = new FieldInfos(fieldInfoArray); - SegmentWriteState state = new SegmentWriteState(null, directory, segmentInfo, fieldInfos, null, IOContext.DEFAULT); - - // Add documents to the field - KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(null, state); - RandomVectorDocValuesProducer randomVectorDocValuesProducer = new RandomVectorDocValuesProducer(docsInSegment, dimension); - knn91DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer); - - // The document should be created in the correct location - String expectedFile = KNNCodecUtil.buildEngineFileName( - segmentName, - knnEngine.getLatestBuildVersion(), - fieldName, - knnEngine.getExtension() - ); - assertFileInCorrectLocation(state, expectedFile); - - // The footer should be valid - assertValidFooter(state.directory, expectedFile); - - // The document should be readable by nmslib - assertLoadableByEngine(state, expectedFile, knnEngine, spaceType, dimension); - } - - public void testAddKNNBinaryField_fromScratch_faissCurrent() throws IOException { - String segmentName = String.format("test_segment%s", randomAlphaOfLength(4)); - int docsInSegment = 100; - String fieldName = String.format("test_field%s", randomAlphaOfLength(4)); - - KNNEngine knnEngine = KNNEngine.FAISS; - SpaceType spaceType = SpaceType.INNER_PRODUCT; - int dimension = 16; - - SegmentInfo segmentInfo = KNNCodecTestUtil.SegmentInfoBuilder.builder(directory, segmentName, docsInSegment, codec).build(); - - KNNMethodContext knnMethodContext = new KNNMethodContext( - knnEngine, - spaceType, - new MethodComponentContext(METHOD_HNSW, ImmutableMap.of(METHOD_PARAMETER_M, 16, METHOD_PARAMETER_EF_CONSTRUCTION, 512)) - ); - - String parameterString = Strings.toString(XContentFactory.jsonBuilder().map(knnEngine.getMethodAsMap(knnMethodContext))); - - FieldInfo[] fieldInfoArray = new FieldInfo[] { - KNNCodecTestUtil.FieldInfoBuilder.builder(fieldName) - .addAttribute(KNNVectorFieldMapper.KNN_FIELD, "true") - .addAttribute(KNNConstants.KNN_ENGINE, knnEngine.getName()) - .addAttribute(KNNConstants.SPACE_TYPE, spaceType.getValue()) - .addAttribute(KNNConstants.PARAMETERS, parameterString) - .build() }; - - FieldInfos fieldInfos = new FieldInfos(fieldInfoArray); - SegmentWriteState state = new SegmentWriteState(null, directory, segmentInfo, fieldInfos, null, IOContext.DEFAULT); - - // Add documents to the field - KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(null, state); - RandomVectorDocValuesProducer randomVectorDocValuesProducer = new RandomVectorDocValuesProducer(docsInSegment, dimension); - knn91DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer); - - // The document should be created in the correct location - String expectedFile = KNNCodecUtil.buildEngineFileName( - segmentName, - knnEngine.getLatestBuildVersion(), - fieldName, - knnEngine.getExtension() - ); - assertFileInCorrectLocation(state, expectedFile); - - // The footer should be valid - assertValidFooter(state.directory, expectedFile); - - // The document should be readable by faiss - assertLoadableByEngine(state, expectedFile, knnEngine, spaceType, dimension); - } - - public void testAddKNNBinaryField_fromModel_faiss() throws IOException, ExecutionException, InterruptedException { - // Generate a trained faiss model - KNNEngine knnEngine = KNNEngine.FAISS; - SpaceType spaceType = SpaceType.INNER_PRODUCT; - int dimension = 16; - String modelId = "test-model-id"; - - float[][] trainingData = getRandomVectors(200, dimension); - long trainingPtr = JNIService.transferVectors(0, trainingData); - - Map parameters = ImmutableMap.of( - INDEX_DESCRIPTION_PARAMETER, - "IVF4,Flat", - KNNConstants.SPACE_TYPE, - SpaceType.L2.getValue() - ); - - byte[] modelBytes = JNIService.trainIndex(parameters, dimension, trainingPtr, knnEngine.getName()); - Model model = new Model( - new ModelMetadata(knnEngine, spaceType, dimension, ModelState.CREATED, "timestamp", "Empty description", ""), - modelBytes, - modelId - ); - JNIService.freeVectors(trainingPtr); - - // Setup the model cache to return the correct model - ModelDao modelDao = mock(ModelDao.class); - when(modelDao.get(modelId)).thenReturn(model); - ClusterService clusterService = mock(ClusterService.class); - when(clusterService.getSettings()).thenReturn(Settings.EMPTY); - - ClusterSettings clusterSettings = new ClusterSettings( - Settings.builder().put(MODEL_CACHE_SIZE_LIMIT_SETTING.getKey(), "10kb").build(), - ImmutableSet.of(MODEL_CACHE_SIZE_LIMIT_SETTING) - ); - - when(clusterService.getClusterSettings()).thenReturn(clusterSettings); - ModelCache.initialize(modelDao, clusterService); - - // Build the segment and field info - String segmentName = String.format("test_segment%s", randomAlphaOfLength(4)); - int docsInSegment = 100; - String fieldName = String.format("test_field%s", randomAlphaOfLength(4)); - - SegmentInfo segmentInfo = KNNCodecTestUtil.SegmentInfoBuilder.builder(directory, segmentName, docsInSegment, codec).build(); - - FieldInfo[] fieldInfoArray = new FieldInfo[] { - KNNCodecTestUtil.FieldInfoBuilder.builder(fieldName) - .addAttribute(KNNVectorFieldMapper.KNN_FIELD, "true") - .addAttribute(MODEL_ID, modelId) - .build() }; - - FieldInfos fieldInfos = new FieldInfos(fieldInfoArray); - SegmentWriteState state = new SegmentWriteState(null, directory, segmentInfo, fieldInfos, null, IOContext.DEFAULT); - - // Add documents to the field - KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(null, state); - RandomVectorDocValuesProducer randomVectorDocValuesProducer = new RandomVectorDocValuesProducer(docsInSegment, dimension); - knn91DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer); - - // The document should be created in the correct location - String expectedFile = KNNCodecUtil.buildEngineFileName( - segmentName, - knnEngine.getLatestBuildVersion(), - fieldName, - knnEngine.getExtension() - ); - assertFileInCorrectLocation(state, expectedFile); - - // The footer should be valid - assertValidFooter(state.directory, expectedFile); - - // The document should be readable by faiss - assertLoadableByEngine(state, expectedFile, knnEngine, spaceType, dimension); - } - - public void testMerge_exception() throws IOException { - KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(null, null); - expectThrows(RuntimeException.class, () -> knn91DocValuesConsumer.merge(null)); - } - - public void testAddSortedSetField() throws IOException { - // Verify that the delegate will be called - DocValuesConsumer delegate = mock(DocValuesConsumer.class); - doNothing().when(delegate).addSortedSetField(null, null); - KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(delegate, null); - knn91DocValuesConsumer.addSortedSetField(null, null); - verify(delegate, times(1)).addSortedSetField(null, null); - } - - public void testAddSortedNumericField() throws IOException { - // Verify that the delegate will be called - DocValuesConsumer delegate = mock(DocValuesConsumer.class); - doNothing().when(delegate).addSortedNumericField(null, null); - KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(delegate, null); - knn91DocValuesConsumer.addSortedNumericField(null, null); - verify(delegate, times(1)).addSortedNumericField(null, null); - } - - public void testAddSortedField() throws IOException { - // Verify that the delegate will be called - DocValuesConsumer delegate = mock(DocValuesConsumer.class); - doNothing().when(delegate).addSortedField(null, null); - KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(delegate, null); - knn91DocValuesConsumer.addSortedField(null, null); - verify(delegate, times(1)).addSortedField(null, null); - } - - public void testAddNumericField() throws IOException { - // Verify that the delegate will be called - DocValuesConsumer delegate = mock(DocValuesConsumer.class); - doNothing().when(delegate).addNumericField(null, null); - KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(delegate, null); - knn91DocValuesConsumer.addNumericField(null, null); - verify(delegate, times(1)).addNumericField(null, null); - } - - public void testClose() throws IOException { - // Verify that the delegate will be called - DocValuesConsumer delegate = mock(DocValuesConsumer.class); - doNothing().when(delegate).close(); - KNN91DocValuesConsumer knn91DocValuesConsumer = new KNN91DocValuesConsumer(delegate, null); - knn91DocValuesConsumer.close(); - verify(delegate, times(1)).close(); - } - -} diff --git a/src/test/java/org/opensearch/knn/index/codec/KNNCodecFactoryTests.java b/src/test/java/org/opensearch/knn/index/codec/KNNCodecFactoryTests.java index 241ae0e56..c20931253 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNNCodecFactoryTests.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNNCodecFactoryTests.java @@ -8,7 +8,7 @@ import org.apache.lucene.codecs.Codec; import org.apache.lucene.codecs.lucene91.Lucene91Codec; import org.opensearch.knn.KNNTestCase; -import org.opensearch.knn.index.codec.KNN91Codec.KNN91Codec; +import org.opensearch.knn.index.codec.KNN910Codec.KNN910Codec; public class KNNCodecFactoryTests extends KNNTestCase { @@ -22,13 +22,13 @@ public void testKNN91DefaultCodec() { Lucene91Codec lucene91CodecDelegate = new Lucene91Codec(); Codec knnCodec = KNNCodecFactory.createKNNCodec(lucene91CodecDelegate); assertNotNull(knnCodec); - assertTrue(knnCodec instanceof KNN91Codec); + assertTrue(knnCodec instanceof KNN910Codec); } public void testKNN91CodecByVersion() { Lucene91Codec lucene91CodecDelegate = new Lucene91Codec(); - Codec knnCodec = KNNCodecFactory.createKNNCodec(KNNCodecFactory.KNNCodecVersion.KNN91, lucene91CodecDelegate); + Codec knnCodec = KNNCodecFactory.createKNNCodec(KNNCodecFactory.KNNCodecVersion.KNN910, lucene91CodecDelegate); assertNotNull(knnCodec); - assertTrue(knnCodec instanceof KNN91Codec); + assertTrue(knnCodec instanceof KNN910Codec); } } diff --git a/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java b/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java index ad4558e0a..2251308a0 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java @@ -11,7 +11,7 @@ import org.opensearch.common.settings.ClusterSettings; import org.opensearch.knn.KNNTestCase; import org.opensearch.knn.common.KNNConstants; -import org.opensearch.knn.index.codec.KNN91Codec.KNN91Codec; +import org.opensearch.knn.index.codec.KNN910Codec.KNN910Codec; import org.opensearch.knn.jni.JNIService; import org.opensearch.knn.index.KNNQuery; import org.opensearch.knn.index.KNNSettings; @@ -61,7 +61,7 @@ */ public class KNNCodecTestCase extends KNNTestCase { - private static final KNN91Codec ACTUAL_CODEC = new KNN91Codec(); + private static final KNN910Codec ACTUAL_CODEC = new KNN910Codec(); private static FieldType sampleFieldType; static { sampleFieldType = new FieldType(KNNVectorFieldMapper.Defaults.FIELD_TYPE); diff --git a/src/test/java/org/opensearch/knn/index/codec/KNNFormatFactoryTests.java b/src/test/java/org/opensearch/knn/index/codec/KNNFormatFactoryTests.java index 778c3f951..545e2a141 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNNFormatFactoryTests.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNNFormatFactoryTests.java @@ -6,16 +6,14 @@ package org.opensearch.knn.index.codec; import org.apache.lucene.codecs.Codec; -import org.apache.lucene.codecs.lucene91.Lucene91Codec; import org.opensearch.knn.KNNTestCase; -import org.opensearch.knn.index.codec.KNN91Codec.KNN91Codec; public class KNNFormatFactoryTests extends KNNTestCase { public void testKNN91Format() { final Codec lucene91CodecDelegate = KNNCodecFactory.CodecDelegateFactory.createKNN91DefaultDelegate(); final Codec knnCodec = KNNCodecFactory.createKNNCodec(lucene91CodecDelegate); - KNNFormatFacade knnFormatFacade = KNNFormatFactory.createKNN91Format(knnCodec); + KNNFormatFacade knnFormatFacade = KNNFormatFactory.createKNN910Format(knnCodec); assertNotNull(knnFormatFacade); assertNotNull(knnFormatFacade.compoundFormat());