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 78% 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 5894ee5a11..00d655803d 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"; + public static final String KNN910 = "KNN910Codec"; private 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 b0d015e64a..0000000000 --- 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 ac3e59db20..0000000000 --- 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 5dc1873466..0000000000 --- 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 3e24af87be..0000000000 --- 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 d870710dfa..0000000000 --- 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 68c4b879dd..b8a75ef130 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java @@ -7,7 +7,7 @@ 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; @@ -17,9 +17,9 @@ */ 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); @@ -47,6 +47,6 @@ public static Codec createKNN91DefaultDelegate() { } 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 5f5a447c6c..54b2fbb337 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 */ 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 98f8a6a139..b897dc36aa 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 39dc641dcd..1ec28d6a44 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 8cb1934c68..0000000000 --- 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 04939603d7..0000000000 --- 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 3ba9409973..0000000000 --- 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 241ae0e56f..c209312532 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 ad4558e0a6..2251308a08 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 778c3f9512..545e2a1419 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());