Merge branch 'main' into optimize-exact-search

Signed-off-by: Junqiu Lei <[email protected]>
junqiu-lei · Sep 27, 2024 · 05c4b2f · 05c4b2f
2 parents 7d1ad23 + eba9d98
commit 05c4b2f
Show file tree

Hide file tree

Showing 11 changed files with 255 additions and 15 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,7 +7,6 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ## [Unreleased 3.0](https://github.com/opensearch-project/k-NN/compare/2.x...HEAD)
 ### Features
 ### Enhancements
-* Adds concurrent segment search support for mode auto [#2111](https://github.com/opensearch-project/k-NN/pull/2111)
 ### Bug Fixes 
 * Add DocValuesProducers for releasing memory when close index [#1946](https://github.com/opensearch-project/k-NN/pull/1946)
 ### Infrastructure
@@ -23,7 +22,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Enhancements
 * Add short circuit if no live docs are in segments [#2059](https://github.com/opensearch-project/k-NN/pull/2059)
 * Optimize reduceToTopK in ResultUtil by removing pre-filling and reducing peek calls [#2146](https://github.com/opensearch-project/k-NN/pull/2146)
+* Update Default Rescore Context based on Dimension [#2149](https://github.com/opensearch-project/k-NN/pull/2149)
 ### Bug Fixes
+* KNN80DocValues should only be considered for BinaryDocValues fields [#2147](https://github.com/opensearch-project/k-NN/pull/2147)
 ### Infrastructure
 ### Documentation
 ### Maintenance

diff --git a/jni/cmake/init-nmslib.cmake b/jni/cmake/init-nmslib.cmake
@@ -19,6 +19,7 @@ if(NOT DEFINED APPLY_LIB_PATCHES OR "${APPLY_LIB_PATCHES}" STREQUAL true)
     set(PATCH_FILE_LIST)
     list(APPEND PATCH_FILE_LIST "${CMAKE_CURRENT_SOURCE_DIR}/patches/nmslib/0001-Initialize-maxlevel-during-add-from-enterpoint-level.patch")
     list(APPEND PATCH_FILE_LIST "${CMAKE_CURRENT_SOURCE_DIR}/patches/nmslib/0002-Adds-ability-to-pass-ef-parameter-in-the-query-for-h.patch")
+    list(APPEND PATCH_FILE_LIST "${CMAKE_CURRENT_SOURCE_DIR}/patches/nmslib/0003-Adding-two-apis-using-stream-to-load-save-in-Hnsw.patch")
 
     # Get patch id of the last commit
     execute_process(COMMAND sh -c "git --no-pager show HEAD | git patch-id --stable" OUTPUT_VARIABLE PATCH_ID_OUTPUT_FROM_COMMIT WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib)

diff --git a/jni/patches/nmslib/0003-Adding-two-apis-using-stream-to-load-save-in-Hnsw.patch b/jni/patches/nmslib/0003-Adding-two-apis-using-stream-to-load-save-in-Hnsw.patch
@@ -0,0 +1,93 @@
+From 7e099ec111e5c9db4b243da249c73f0ecc206281 Mon Sep 17 00:00:00 2001
+From: Dooyong Kim <[email protected]>
+Date: Thu, 26 Sep 2024 15:20:53 -0700
+Subject: [PATCH] Adding two apis using stream to load/save in Hnsw.
+
+Signed-off-by: Dooyong Kim <[email protected]>
+---
+ similarity_search/include/method/hnsw.h |  4 +++
+ similarity_search/src/method/hnsw.cc    | 44 +++++++++++++++++++++++++
+ 2 files changed, 48 insertions(+)
+
+diff --git a/similarity_search/include/method/hnsw.h b/similarity_search/include/method/hnsw.h
+index 57d99d0..7ff3f3d 100644
+--- a/similarity_search/include/method/hnsw.h
++++ b/similarity_search/include/method/hnsw.h
+@@ -455,8 +455,12 @@ namespace similarity {
+     public:
+         virtual void SaveIndex(const string &location) override;
+
++        void SaveIndexWithStream(std::ostream& output);
++
+         virtual void LoadIndex(const string &location) override;
+
++        void LoadIndexWithStream(std::istream& in);
++
+         Hnsw(bool PrintProgress, const Space<dist_t> &space, const ObjectVector &data);
+         void CreateIndex(const AnyParams &IndexParams) override;
+
+diff --git a/similarity_search/src/method/hnsw.cc b/similarity_search/src/method/hnsw.cc
+index 35b372c..e7a2c9e 100644
+--- a/similarity_search/src/method/hnsw.cc
++++ b/similarity_search/src/method/hnsw.cc
+@@ -771,6 +771,25 @@ namespace similarity {
+         output.close();
+     }
+
++    template <typename dist_t>
++    void Hnsw<dist_t>::SaveIndexWithStream(std::ostream &output) {
++        output.exceptions(ios::badbit | ios::failbit);
++
++        unsigned int optimIndexFlag = data_level0_memory_ != nullptr;
++
++        writeBinaryPOD(output, optimIndexFlag);
++
++        if (!optimIndexFlag) {
++#if USE_TEXT_REGULAR_INDEX
++            SaveRegularIndexText(output);
++#else
++            SaveRegularIndexBin(output);
++#endif
++        } else {
++            SaveOptimizedIndex(output);
++        }
++    }
++
+     template <typename dist_t>
+     void
+     Hnsw<dist_t>::SaveOptimizedIndex(std::ostream& output) {
+@@ -1021,6 +1040,31 @@ namespace similarity {
+
+     }
+
++    template <typename dist_t>
++    void Hnsw<dist_t>::LoadIndexWithStream(std::istream& input) {
++        LOG(LIB_INFO) << "Loading index from an input stream.";
++        CHECK_MSG(input, "Cannot open file for reading with an input stream");
++
++        input.exceptions(ios::badbit | ios::failbit);
++
++#if USE_TEXT_REGULAR_INDEX
++        LoadRegularIndexText(input);
++#else
++        unsigned int optimIndexFlag= 0;
++
++        readBinaryPOD(input, optimIndexFlag);
++
++        if (!optimIndexFlag) {
++            LoadRegularIndexBin(input);
++        } else {
++            LoadOptimizedIndex(input);
++        }
++#endif
++
++        LOG(LIB_INFO) << "Finished loading index";
++        visitedlistpool = new VisitedListPool(1, totalElementsStored_);
++    }
++
+
+     template <typename dist_t>
+     void
+-- 
+2.39.5 (Apple Git-154)
+
diff --git a/release-notes/opensearch-knn.release-notes-2.17.1.0.md b/release-notes/opensearch-knn.release-notes-2.17.1.0.md
@@ -0,0 +1,8 @@
+## Version 2.17.1.0 Release Notes
+
+Compatible with OpenSearch 2.17.1
+
+### Enhancements
+* Adds concurrent segment search support for mode auto [#2111](https://github.com/opensearch-project/k-NN/pull/2111)
+### Bug Fixes
+* Change min oversample to 1 [#2117](https://github.com/opensearch-project/k-NN/pull/2117)
diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesProducer.java b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesProducer.java
@@ -15,6 +15,7 @@
 import lombok.extern.log4j.Log4j2;
 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.NumericDocValues;
 import org.apache.lucene.index.SegmentReadState;
@@ -69,6 +70,13 @@ public KNN80DocValuesProducer(DocValuesProducer delegate, SegmentReadState state
             if (!field.attributes().containsKey(KNN_FIELD)) {
                 continue;
             }
+            // Only segments that contains BinaryDocValues and doesn't have vector values should be considered.
+            // By default, we don't create BinaryDocValues for knn field anymore. However, users can set doc_values = true
+            // to create binary doc values explicitly like any other field. Hence, we only want to include fields
+            // where approximate search is possible only by BinaryDocValues.
+            if (field.getDocValuesType() != DocValuesType.BINARY || field.hasVectorValues() == true) {
+                continue;
+            }
             // Only Native Engine put into indexPathMap
             KNNEngine knnEngine = getNativeKNNEngine(field);
             if (knnEngine == null) {
@@ -77,6 +85,7 @@ public KNN80DocValuesProducer(DocValuesProducer delegate, SegmentReadState state
             List<String> engineFiles = KNNCodecUtil.getEngineFiles(knnEngine.getExtension(), field.name, state.segmentInfo);
             Path indexPath = PathUtils.get(directoryPath, engineFiles.get(0));
             indexPathMap.putIfAbsent(field.getName(), indexPath.toString());
+
         }
     }
 

diff --git a/src/main/java/org/opensearch/knn/index/mapper/CompressionLevel.java b/src/main/java/org/opensearch/knn/index/mapper/CompressionLevel.java
@@ -94,9 +94,34 @@ public static boolean isConfigured(CompressionLevel compressionLevel) {
         return compressionLevel != null && compressionLevel != NOT_CONFIGURED;
     }
 
-    public RescoreContext getDefaultRescoreContext(Mode mode) {
+    /**
+     * Returns the appropriate {@link RescoreContext} based on the given {@code mode} and {@code dimension}.
+     *
+     * <p>If the {@code mode} is present in the valid {@code modesForRescore} set, the method checks the value of
+     * {@code dimension}:
+     * <ul>
+     *     <li>If {@code dimension} is less than or equal to 1000, it returns a {@link RescoreContext} with an
+     *         oversample factor of 5.0f.</li>
+     *     <li>If {@code dimension} is greater than 1000, it returns the default {@link RescoreContext} associated with
+     *         the {@link CompressionLevel}. If no default is set, it falls back to {@link RescoreContext#getDefault()}.</li>
+     * </ul>
+     * If the {@code mode} is not valid, the method returns {@code null}.
+     *
+     * @param mode      The {@link Mode} for which to retrieve the {@link RescoreContext}.
+     * @param dimension The dimensional value that determines the {@link RescoreContext} behavior.
+     * @return          A {@link RescoreContext} with an oversample factor of 5.0f if {@code dimension} is less than
+     *                  or equal to 1000, the default {@link RescoreContext} if greater, or {@code null} if the mode
+     *                  is invalid.
+     */
+    public RescoreContext getDefaultRescoreContext(Mode mode, int dimension) {
         if (modesForRescore.contains(mode)) {
-            return defaultRescoreContext;
+            // Adjust RescoreContext based on dimension
+            if (dimension <= RescoreContext.DIMENSION_THRESHOLD) {
+                // For dimensions <= 1000, return a RescoreContext with 5.0f oversample factor
+                return RescoreContext.builder().oversampleFactor(RescoreContext.OVERSAMPLE_FACTOR_BELOW_DIMENSION_THRESHOLD).build();
+            } else {
+                return defaultRescoreContext;
+            }
         }
         return null;
     }

diff --git a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldType.java b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldType.java
@@ -93,6 +93,10 @@ public RescoreContext resolveRescoreContext(RescoreContext userProvidedContext)
         if (userProvidedContext != null) {
             return userProvidedContext;
         }
-        return getKnnMappingConfig().getCompressionLevel().getDefaultRescoreContext(getKnnMappingConfig().getMode());
+        KNNMappingConfig knnMappingConfig = getKnnMappingConfig();
+        int dimension = knnMappingConfig.getDimension();
+        CompressionLevel compressionLevel = knnMappingConfig.getCompressionLevel();
+        Mode mode = knnMappingConfig.getMode();
+        return compressionLevel.getDefaultRescoreContext(mode, dimension);
     }
 }
diff --git a/src/main/java/org/opensearch/knn/index/query/rescore/RescoreContext.java b/src/main/java/org/opensearch/knn/index/query/rescore/RescoreContext.java
@@ -21,6 +21,8 @@ public final class RescoreContext {
     public static final float MIN_OVERSAMPLE_FACTOR = 1.0f;
 
     public static final int MAX_FIRST_PASS_RESULTS = 10000;
+    public static final int DIMENSION_THRESHOLD = 1000;
+    public static final float OVERSAMPLE_FACTOR_BELOW_DIMENSION_THRESHOLD = 5.0f;
 
     // Todo:- We will improve this in upcoming releases
     public static final int MIN_FIRST_PASS_RESULTS = 100;

diff --git a/src/test/java/org/opensearch/knn/index/OpenSearchIT.java b/src/test/java/org/opensearch/knn/index/OpenSearchIT.java
@@ -16,6 +16,7 @@
 import java.util.Locale;
 import lombok.SneakyThrows;
 import org.junit.BeforeClass;
+import org.junit.Ignore;
 import org.opensearch.knn.KNNRestTestCase;
 import org.opensearch.knn.KNNResult;
 import org.apache.hc.core5.http.io.entity.EntityUtils;
@@ -483,6 +484,9 @@ public void testIndexingVectorValidation_updateVectorWithNull() throws Exception
         assertArrayEquals(vectorForDocumentOne, vectorRestoreInitialValue);
     }
 
+    // This doesn't work since indices that are created post 2.17 don't evict by default when indices are closed or deleted.
+    // Enable this PR once https://github.com/opensearch-project/k-NN/issues/2148 is resolved.
+    @Ignore
     public void testCacheClear_whenCloseIndex() throws Exception {
         String indexName = "test-index-1";
         KNNEngine knnEngine1 = KNNEngine.NMSLIB;

diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesProducerTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesProducerTests.java
@@ -9,6 +9,7 @@
 import org.apache.lucene.codecs.Codec;
 import org.apache.lucene.codecs.DocValuesFormat;
 import org.apache.lucene.codecs.DocValuesProducer;
+import org.apache.lucene.index.DocValuesType;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.FieldInfos;
 import org.apache.lucene.index.SegmentInfo;
@@ -127,4 +128,57 @@ public void testProduceKNNBinaryField_fromCodec_nmslibCurrent() throws IOExcepti
         assertTrue(path.contains(segmentFiles.get(0)));
     }
 
+    public void testProduceKNNBinaryField_whenFieldHasNonBinaryDocValues_thenSkipThoseField() throws IOException {
+        // Set information about the segment and the fields
+        DocValuesFormat mockDocValuesFormat = mock(DocValuesFormat.class);
+        Codec mockDelegateCodec = mock(Codec.class);
+        DocValuesProducer mockDocValuesProducer = mock(DocValuesProducer.class);
+        when(mockDelegateCodec.docValuesFormat()).thenReturn(mockDocValuesFormat);
+        when(mockDocValuesFormat.fieldsProducer(any())).thenReturn(mockDocValuesProducer);
+        when(mockDocValuesFormat.getName()).thenReturn("mockDocValuesFormat");
+        Codec codec = new KNN87Codec(mockDelegateCodec);
+
+        String segmentName = "_test";
+        int docsInSegment = 100;
+        String fieldName1 = String.format("test_field1%s", randomAlphaOfLength(4));
+        String fieldName2 = String.format("test_field2%s", randomAlphaOfLength(4));
+        List<String> segmentFiles = Arrays.asList(
+            String.format("%s_2011_%s%s", segmentName, fieldName1, KNNEngine.NMSLIB.getExtension()),
+            String.format("%s_165_%s%s", segmentName, fieldName2, KNNEngine.FAISS.getExtension())
+        );
+
+        KNNEngine knnEngine = KNNEngine.NMSLIB;
+        SpaceType spaceType = SpaceType.COSINESIMIL;
+        SegmentInfo segmentInfo = KNNCodecTestUtil.segmentInfoBuilder()
+            .directory(directory)
+            .segmentName(segmentName)
+            .docsInSegment(docsInSegment)
+            .codec(codec)
+            .build();
+
+        for (String name : segmentFiles) {
+            IndexOutput indexOutput = directory.createOutput(name, IOContext.DEFAULT);
+            indexOutput.close();
+        }
+        segmentInfo.setFiles(segmentFiles);
+
+        FieldInfo[] fieldInfoArray = new FieldInfo[] {
+            KNNCodecTestUtil.FieldInfoBuilder.builder(fieldName1)
+                .addAttribute(KNNVectorFieldMapper.KNN_FIELD, "true")
+                .addAttribute(KNNConstants.KNN_ENGINE, knnEngine.getName())
+                .addAttribute(KNNConstants.SPACE_TYPE, spaceType.getValue())
+                .docValuesType(DocValuesType.NONE)
+                .dvGen(-1)
+                .build() };
+
+        FieldInfos fieldInfos = new FieldInfos(fieldInfoArray);
+        SegmentReadState state = new SegmentReadState(directory, segmentInfo, fieldInfos, IOContext.DEFAULT);
+
+        DocValuesFormat docValuesFormat = codec.docValuesFormat();
+        assertTrue(docValuesFormat instanceof KNN80DocValuesFormat);
+        DocValuesProducer producer = docValuesFormat.fieldsProducer(state);
+        assertTrue(producer instanceof KNN80DocValuesProducer);
+        assertEquals(0, ((KNN80DocValuesProducer) producer).getOpenedIndexPath().size());
+    }
+
 }
diff --git a/src/test/java/org/opensearch/knn/index/mapper/CompressionLevelTests.java b/src/test/java/org/opensearch/knn/index/mapper/CompressionLevelTests.java
@@ -44,26 +44,65 @@ public void testIsConfigured() {
     public void testGetDefaultRescoreContext() {
         // Test rescore context for ON_DISK mode
         Mode mode = Mode.ON_DISK;
+        int belowThresholdDimension = 500; // A dimension below the threshold
+        int aboveThresholdDimension = 1500; // A dimension above the threshold
 
-        // x32 should have RescoreContext with an oversample factor of 3.0f
-        RescoreContext rescoreContext = CompressionLevel.x32.getDefaultRescoreContext(mode);
+        // x32 with dimension <= 1000 should have an oversample factor of 5.0f
+        RescoreContext rescoreContext = CompressionLevel.x32.getDefaultRescoreContext(mode, belowThresholdDimension);
+        assertNotNull(rescoreContext);
+        assertEquals(5.0f, rescoreContext.getOversampleFactor(), 0.0f);
+
+        // x32 with dimension > 1000 should have an oversample factor of 3.0f
+        rescoreContext = CompressionLevel.x32.getDefaultRescoreContext(mode, aboveThresholdDimension);
         assertNotNull(rescoreContext);
         assertEquals(3.0f, rescoreContext.getOversampleFactor(), 0.0f);
 
-        // x16 should have RescoreContext with an oversample factor of 3.0f
-        rescoreContext = CompressionLevel.x16.getDefaultRescoreContext(mode);
+        // x16 with dimension <= 1000 should have an oversample factor of 5.0f
+        rescoreContext = CompressionLevel.x16.getDefaultRescoreContext(mode, belowThresholdDimension);
+        assertNotNull(rescoreContext);
+        assertEquals(5.0f, rescoreContext.getOversampleFactor(), 0.0f);
+
+        // x16 with dimension > 1000 should have an oversample factor of 3.0f
+        rescoreContext = CompressionLevel.x16.getDefaultRescoreContext(mode, aboveThresholdDimension);
         assertNotNull(rescoreContext);
         assertEquals(3.0f, rescoreContext.getOversampleFactor(), 0.0f);
 
-        // x8 should have RescoreContext with an oversample factor of 2.0f
-        rescoreContext = CompressionLevel.x8.getDefaultRescoreContext(mode);
+        // x8 with dimension <= 1000 should have an oversample factor of 5.0f
+        rescoreContext = CompressionLevel.x8.getDefaultRescoreContext(mode, belowThresholdDimension);
+        assertNotNull(rescoreContext);
+        assertEquals(5.0f, rescoreContext.getOversampleFactor(), 0.0f);
+
+        // x8 with dimension > 1000 should have an oversample factor of 2.0f
+        rescoreContext = CompressionLevel.x8.getDefaultRescoreContext(mode, aboveThresholdDimension);
         assertNotNull(rescoreContext);
         assertEquals(2.0f, rescoreContext.getOversampleFactor(), 0.0f);
 
-        // Other compression levels should not have a RescoreContext for ON_DISK mode
-        assertNull(CompressionLevel.x4.getDefaultRescoreContext(mode));
-        assertNull(CompressionLevel.x2.getDefaultRescoreContext(mode));
-        assertNull(CompressionLevel.x1.getDefaultRescoreContext(mode));
-        assertNull(CompressionLevel.NOT_CONFIGURED.getDefaultRescoreContext(mode));
+        // x4 with dimension <= 1000 should have an oversample factor of 5.0f (though it doesn't have its own RescoreContext)
+        rescoreContext = CompressionLevel.x4.getDefaultRescoreContext(mode, belowThresholdDimension);
+        assertNull(rescoreContext);
+        // x4 with dimension > 1000 should return null (no RescoreContext is configured for x4)
+        rescoreContext = CompressionLevel.x4.getDefaultRescoreContext(mode, aboveThresholdDimension);
+        assertNull(rescoreContext);
+
+        // Other compression levels should behave similarly with respect to dimension
+
+        rescoreContext = CompressionLevel.x2.getDefaultRescoreContext(mode, belowThresholdDimension);
+        assertNull(rescoreContext);
+
+        // x2 with dimension > 1000 should return null
+        rescoreContext = CompressionLevel.x2.getDefaultRescoreContext(mode, aboveThresholdDimension);
+        assertNull(rescoreContext);
+
+        rescoreContext = CompressionLevel.x1.getDefaultRescoreContext(mode, belowThresholdDimension);
+        assertNull(rescoreContext);
+
+        // x1 with dimension > 1000 should return null
+        rescoreContext = CompressionLevel.x1.getDefaultRescoreContext(mode, aboveThresholdDimension);
+        assertNull(rescoreContext);
+
+        // NOT_CONFIGURED with dimension <= 1000 should return a RescoreContext with an oversample factor of 5.0f
+        rescoreContext = CompressionLevel.NOT_CONFIGURED.getDefaultRescoreContext(mode, belowThresholdDimension);
+        assertNull(rescoreContext);
+
     }
 }