diff --git a/src/main/java/org/opensearch/knn/index/mapper/KNNMappingConfig.java b/src/main/java/org/opensearch/knn/index/mapper/KNNMappingConfig.java index 5b1955f23..cd77ebd9a 100644 --- a/src/main/java/org/opensearch/knn/index/mapper/KNNMappingConfig.java +++ b/src/main/java/org/opensearch/knn/index/mapper/KNNMappingConfig.java @@ -6,6 +6,7 @@ package org.opensearch.knn.index.mapper; import org.opensearch.knn.index.engine.KNNMethodContext; +import org.opensearch.knn.index.engine.qframe.QuantizationConfig; import java.util.Optional; @@ -48,6 +49,14 @@ default CompressionLevel getCompressionLevel() { return CompressionLevel.NOT_CONFIGURED; } + /** + * Returns quantization config + * @return + */ + default QuantizationConfig getQuantizationConfig() { + return QuantizationConfig.EMPTY; + } + /** * * @return the dimension of the index; for model based indices, it will be null diff --git a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java index 10ba0d94d..265876310 100644 --- a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java +++ b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java @@ -241,7 +241,8 @@ public KNNVectorFieldMapper build(BuilderContext context) { hasDocValues.get(), modelDao, indexCreatedVersion, - originalParameters + originalParameters, + knnMethodConfigContext ); } diff --git a/src/main/java/org/opensearch/knn/index/mapper/MethodFieldMapper.java b/src/main/java/org/opensearch/knn/index/mapper/MethodFieldMapper.java index d479da39c..3d11949fe 100644 --- a/src/main/java/org/opensearch/knn/index/mapper/MethodFieldMapper.java +++ b/src/main/java/org/opensearch/knn/index/mapper/MethodFieldMapper.java @@ -51,6 +51,12 @@ public static MethodFieldMapper createFieldMapper( boolean hasDocValues, OriginalMappingParameters originalMappingParameters ) { + + KNNMethodContext knnMethodContext = originalMappingParameters.getResolvedKnnMethodContext(); + QuantizationConfig quantizationConfig = knnMethodContext.getKnnEngine() + .getKNNLibraryIndexingContext(knnMethodContext, knnMethodConfigContext) + .getQuantizationConfig(); + final KNNVectorFieldType mappedFieldType = new KNNVectorFieldType( fullname, metaValue, @@ -75,6 +81,11 @@ public Mode getMode() { public CompressionLevel getCompressionLevel() { return knnMethodConfigContext.getCompressionLevel(); } + + @Override + public QuantizationConfig getQuantizationConfig() { + return quantizationConfig; + } } ); return new MethodFieldMapper( diff --git a/src/main/java/org/opensearch/knn/index/mapper/ModelFieldMapper.java b/src/main/java/org/opensearch/knn/index/mapper/ModelFieldMapper.java index 42a0c10c4..013cb0c53 100644 --- a/src/main/java/org/opensearch/knn/index/mapper/ModelFieldMapper.java +++ b/src/main/java/org/opensearch/knn/index/mapper/ModelFieldMapper.java @@ -56,9 +56,17 @@ public static ModelFieldMapper createFieldMapper( boolean hasDocValues, ModelDao modelDao, Version indexCreatedVersion, - OriginalMappingParameters originalMappingParameters + OriginalMappingParameters originalMappingParameters, + KNNMethodConfigContext knnMethodConfigContext ) { + final KNNMethodContext knnMethodContext = originalMappingParameters.getKnnMethodContext(); + final QuantizationConfig quantizationConfig = knnMethodContext == null + ? QuantizationConfig.EMPTY + : knnMethodContext.getKnnEngine() + .getKNNLibraryIndexingContext(knnMethodContext, knnMethodConfigContext) + .getQuantizationConfig(); + final KNNVectorFieldType mappedFieldType = new KNNVectorFieldType(fullname, metaValue, vectorDataType, new KNNMappingConfig() { private Integer dimension = null; private Mode mode = null; @@ -94,6 +102,11 @@ public CompressionLevel getCompressionLevel() { return compressionLevel; } + @Override + public QuantizationConfig getQuantizationConfig() { + return quantizationConfig; + } + // ModelMetadata relies on cluster state which may not be available during field mapper creation. Thus, // we lazily initialize it. private void initFromModelMetadata() { diff --git a/src/main/java/org/opensearch/knn/index/query/KNNQueryBuilder.java b/src/main/java/org/opensearch/knn/index/query/KNNQueryBuilder.java index b699a7705..8f7c5a3ff 100644 --- a/src/main/java/org/opensearch/knn/index/query/KNNQueryBuilder.java +++ b/src/main/java/org/opensearch/knn/index/query/KNNQueryBuilder.java @@ -25,6 +25,7 @@ import org.opensearch.index.query.QueryShardContext; import org.opensearch.knn.index.engine.KNNMethodConfigContext; import org.opensearch.knn.index.engine.model.QueryContext; +import org.opensearch.knn.index.engine.qframe.QuantizationConfig; import org.opensearch.knn.index.mapper.KNNMappingConfig; import org.opensearch.knn.index.mapper.KNNVectorFieldType; import org.opensearch.knn.index.query.parser.RescoreParser; @@ -451,6 +452,10 @@ protected Query doToQuery(QueryShardContext context) { if (vectorDataType == VectorDataType.BINARY) { throw new UnsupportedOperationException(String.format(Locale.ROOT, "Binary data type does not support radial search")); } + + if (knnMappingConfig.getQuantizationConfig() != QuantizationConfig.EMPTY) { + throw new UnsupportedOperationException("Radial search is not supported for indices which have quantization enabled"); + } } // Currently, k-NN supports distance and score types radial search diff --git a/src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java b/src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java index ac5a72945..c13d86554 100644 --- a/src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java +++ b/src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java @@ -7,6 +7,7 @@ import lombok.Getter; import lombok.RequiredArgsConstructor; +import lombok.extern.log4j.Log4j2; import org.apache.lucene.index.IndexReader; import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.search.IndexSearcher; @@ -18,6 +19,7 @@ import org.apache.lucene.search.Weight; import org.apache.lucene.util.BitSet; import org.apache.lucene.util.Bits; +import org.opensearch.common.StopWatch; import org.opensearch.knn.index.query.KNNQuery; import org.opensearch.knn.index.query.KNNWeight; import org.opensearch.knn.index.query.ResultUtil; @@ -39,6 +41,7 @@ * for k-NN query if required. This is done by overriding rewrite method to execute ANN on each leaf * {@link KNNQuery} does not give the ability to post process segment results. */ +@Log4j2 @Getter @RequiredArgsConstructor public class NativeEngineKnnVectorQuery extends Query { @@ -60,7 +63,11 @@ public Weight createWeight(IndexSearcher indexSearcher, ScoreMode scoreMode, flo int firstPassK = rescoreContext.getFirstPassK(finalK); perLeafResults = doSearch(indexSearcher, leafReaderContexts, knnWeight, firstPassK); ResultUtil.reduceToTopK(perLeafResults, firstPassK); + + StopWatch stopWatch = new StopWatch().start(); perLeafResults = doRescore(indexSearcher, leafReaderContexts, knnWeight, perLeafResults, finalK); + long rescoreTime = stopWatch.stop().totalTime().millis(); + log.debug("Rescoring results took {} ms. oversampled k:{}, segments:{}", rescoreTime, firstPassK, leafReaderContexts.size()); } ResultUtil.reduceToTopK(perLeafResults, finalK); TopDocs[] topDocs = new TopDocs[perLeafResults.size()]; diff --git a/src/test/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperTests.java b/src/test/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperTests.java index 2b034de37..496641339 100644 --- a/src/test/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperTests.java +++ b/src/test/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperTests.java @@ -1130,6 +1130,12 @@ public void testModelFieldMapperParseCreateField_validInput_thenDifferentFieldTy MockedStatic utilMockedStatic = Mockito.mockStatic(KNNVectorFieldMapperUtil.class); MockedStatic modelUtilMockedStatic = Mockito.mockStatic(ModelUtil.class) ) { + KNNMethodConfigContext knnMethodConfigContext = KNNMethodConfigContext.builder() + .vectorDataType(VectorDataType.FLOAT) + .versionCreated(CURRENT) + .dimension(TEST_DIMENSION) + .build(); + for (VectorDataType dataType : VectorDataType.values()) { log.info("Vector Data Type is : {}", dataType); SpaceType spaceType = VectorDataType.BINARY == dataType ? SpaceType.DEFAULT_BINARY : SpaceType.INNER_PRODUCT; @@ -1173,7 +1179,8 @@ public void testModelFieldMapperParseCreateField_validInput_thenDifferentFieldTy false, modelDao, CURRENT, - originalMappingParameters + originalMappingParameters, + knnMethodConfigContext ); modelFieldMapper.parseCreateField(parseContext); @@ -1214,7 +1221,8 @@ public void testModelFieldMapperParseCreateField_validInput_thenDifferentFieldTy false, modelDao, CURRENT, - originalMappingParameters + originalMappingParameters, + knnMethodConfigContext ); modelFieldMapper.parseCreateField(parseContext); diff --git a/src/test/java/org/opensearch/knn/index/query/KNNQueryBuilderTests.java b/src/test/java/org/opensearch/knn/index/query/KNNQueryBuilderTests.java index 1c1b6edd9..3db03085b 100644 --- a/src/test/java/org/opensearch/knn/index/query/KNNQueryBuilderTests.java +++ b/src/test/java/org/opensearch/knn/index/query/KNNQueryBuilderTests.java @@ -29,7 +29,10 @@ import org.opensearch.index.query.QueryShardContext; import org.opensearch.index.query.TermQueryBuilder; import org.opensearch.knn.KNNTestCase; +import org.opensearch.knn.index.engine.qframe.QuantizationConfig; +import org.opensearch.knn.index.mapper.KNNMappingConfig; import org.opensearch.knn.index.mapper.KNNVectorFieldType; +import org.opensearch.knn.index.mapper.Mode; import org.opensearch.knn.index.query.rescore.RescoreContext; import org.opensearch.knn.index.util.KNNClusterUtil; import org.opensearch.knn.index.engine.KNNMethodContext; @@ -41,6 +44,7 @@ import org.opensearch.knn.indices.ModelDao; import org.opensearch.knn.indices.ModelMetadata; import org.opensearch.knn.indices.ModelState; +import org.opensearch.knn.quantization.enums.ScalarQuantizationType; import java.io.IOException; import java.util.Arrays; @@ -440,6 +444,47 @@ public void testDoToQuery_whenRadialSearchOnBinaryIndex_thenException() { assertTrue(e.getMessage().contains("Binary data type does not support radial search")); } + public void testDoToQuery_whenRadialSearchOnDiskMode_thenException() { + float[] queryVector = { 1.0f }; + KNNQueryBuilder knnQueryBuilder = KNNQueryBuilder.builder() + .fieldName(FIELD_NAME) + .vector(queryVector) + .maxDistance(MAX_DISTANCE) + .build(); + Index dummyIndex = new Index("dummy", "dummy"); + QueryShardContext mockQueryShardContext = mock(QueryShardContext.class); + KNNVectorFieldType mockKNNVectorField = mock(KNNVectorFieldType.class); + when(mockQueryShardContext.index()).thenReturn(dummyIndex); + when(mockKNNVectorField.getVectorDataType()).thenReturn(VectorDataType.FLOAT); + when(mockQueryShardContext.fieldMapper(anyString())).thenReturn(mockKNNVectorField); + MethodComponentContext methodComponentContext = new MethodComponentContext( + org.opensearch.knn.common.KNNConstants.METHOD_HNSW, + ImmutableMap.of() + ); + KNNMethodContext knnMethodContext = new KNNMethodContext(KNNEngine.FAISS, SpaceType.L2, methodComponentContext); + when(mockKNNVectorField.getKnnMappingConfig()).thenReturn(new KNNMappingConfig() { + @Override + public Optional getKnnMethodContext() { + return Optional.of(knnMethodContext); + } + + @Override + public int getDimension() { + return 1; + } + + public Mode getMode() { + return Mode.ON_DISK; + } + + public QuantizationConfig getQuantizationConfig() { + return QuantizationConfig.builder().quantizationType(ScalarQuantizationType.ONE_BIT).build(); + } + }); + Exception e = expectThrows(UnsupportedOperationException.class, () -> knnQueryBuilder.doToQuery(mockQueryShardContext)); + assertEquals("Radial search is not supported for indices which have quantization enabled", e.getMessage()); + } + public void testDoToQuery_KnnQueryWithFilter_Lucene() throws Exception { // Given float[] queryVector = { 1.0f, 2.0f, 3.0f, 4.0f };