Merge branch 'main' into main

Signed-off-by: Vikasht34 <viktari@amazon.com> Signed-off-by: VIKASH TIWARI <viktari@amazon.com>
Vikasht34 · Aug 12, 2024 · 5fd6d84 · 5fd6d84
2 parents b71c94e + f5ba771
commit 5fd6d84
Show file tree

Hide file tree

Showing 56 changed files with 2,777 additions and 1,092 deletions.
diff --git a/.github/workflows/backwards_compatibility_tests_workflow.yml b/.github/workflows/backwards_compatibility_tests_workflow.yml
@@ -35,7 +35,7 @@ jobs:
       matrix:
         java: [ 21 ]
         os: [ubuntu-latest]
-        bwc_version : [ "2.0.1", "2.1.0", "2.2.1", "2.3.0", "2.4.1", "2.5.0", "2.6.0", "2.7.0", "2.8.0", "2.9.0", "2.10.0", "2.11.0", "2.12.0", "2.13.0", "2.14.0", "2.15.0", "2.16.0-SNAPSHOT"]
+        bwc_version : [ "2.0.1", "2.1.0", "2.2.1", "2.3.0", "2.4.1", "2.5.0", "2.6.0", "2.7.0", "2.8.0", "2.9.0", "2.10.0", "2.11.0", "2.12.0", "2.13.0", "2.14.0", "2.15.0", "2.16.0", "2.17.0-SNAPSHOT"]
         opensearch_version : [ "3.0.0-SNAPSHOT" ]
         exclude:
           - os: windows-latest
@@ -114,7 +114,7 @@ jobs:
       matrix:
         java: [ 21 ]
         os: [ubuntu-latest]
-        bwc_version: [ "2.16.0-SNAPSHOT" ]
+        bwc_version: [ "2.17.0-SNAPSHOT" ]
         opensearch_version: [ "3.0.0-SNAPSHOT" ]
 
     name: k-NN Rolling-Upgrade BWC Tests

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -19,6 +19,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Bug Fixes
 * Corrected search logic for scenario with non-existent fields in filter [#1874](https://github.com/opensearch-project/k-NN/pull/1874)
 * Add script_fields context to KNNAllowlist [#1917] (https://github.com/opensearch-project/k-NN/pull/1917)
+* Fix graph merge stats size calculation [#1844](https://github.com/opensearch-project/k-NN/pull/1844)
+* Integrate Lucene Vector field with native engines to use KNNVectorFormat during segment creation [#1945](https://github.com/opensearch-project/k-NN/pull/1945)
+* Disallow a vector field to have an invalid character for a physical file name. [#1936] (https://github.com/opensearch-project/k-NN/pull/1936)
 ### Infrastructure
 ### Documentation
 ### Maintenance
@@ -30,4 +33,6 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 * Refactor method structure and definitions [#1920](https://github.com/opensearch-project/k-NN/pull/1920)
 * Refactor KNNVectorFieldType from KNNVectorFieldMapper to a separate class for better readability. [#1931](https://github.com/opensearch-project/k-NN/pull/1931)
 * Generalize lib interface to return context objects [#1925](https://github.com/opensearch-project/k-NN/pull/1925)
+* Move k search k-NN query to re-write phase of vector search query for Native Engines [#1877](https://github.com/opensearch-project/k-NN/pull/1877)
+* Restructure mappers to better handle null cases and avoid branching in parsing [#1939](https://github.com/opensearch-project/k-NN/pull/1939)
 * Added Quantization Framework and implemented 1Bit and multibit quantizer[#1889](https://github.com/opensearch-project/k-NN/issues/1889)
diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/IndexingIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/IndexingIT.java
@@ -53,6 +53,21 @@ public void testKNNIndexDefaultLegacyFieldMapping() throws Exception {
         }
     }
 
+    // Ensure that when segments created with old mapping are forcemerged in new cluster, they
+    // succeed
+    public void testKNNIndexDefaultLegacyFieldMappingForceMerge() throws Exception {
+        waitForClusterHealthGreen(NODES_BWC_CLUSTER);
+
+        if (isRunningAgainstOldCluster()) {
+            createKnnIndex(testIndex, getKNNDefaultIndexSettings(), createKnnIndexMapping(TEST_FIELD, DIMENSIONS));
+            addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, 100);
+            // Flush to ensure that index is not re-indexed when node comes back up
+            flush(testIndex, true);
+        } else {
+            forceMergeKnnIndex(testIndex);
+        }
+    }
+
     // Custom Legacy Field Mapping
     // space_type : "linf", engine : "nmslib", m : 2, ef_construction : 2
     public void testKNNIndexCustomLegacyFieldMapping() throws Exception {

diff --git a/src/main/java/org/opensearch/knn/common/featureflags/KNNFeatureFlags.java b/src/main/java/org/opensearch/knn/common/featureflags/KNNFeatureFlags.java
@@ -0,0 +1,46 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package org.opensearch.knn.common.featureflags;
+
+import com.google.common.annotations.VisibleForTesting;
+import lombok.experimental.UtilityClass;
+import org.opensearch.common.settings.Setting;
+import org.opensearch.knn.index.KNNSettings;
+
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import static org.opensearch.common.settings.Setting.Property.Dynamic;
+import static org.opensearch.common.settings.Setting.Property.NodeScope;
+
+/**
+ * Class to manage KNN feature flags
+ */
+@UtilityClass
+public class KNNFeatureFlags {
+
+    // Feature flags
+    private static final String KNN_LAUNCH_QUERY_REWRITE_ENABLED = "knn.feature.query.rewrite.enabled";
+    private static final boolean KNN_LAUNCH_QUERY_REWRITE_ENABLED_DEFAULT = true;
+
+    @VisibleForTesting
+    public static final Setting<Boolean> KNN_LAUNCH_QUERY_REWRITE_ENABLED_SETTING = Setting.boolSetting(
+        KNN_LAUNCH_QUERY_REWRITE_ENABLED,
+        KNN_LAUNCH_QUERY_REWRITE_ENABLED_DEFAULT,
+        NodeScope,
+        Dynamic
+    );
+
+    public static List<Setting<?>> getFeatureFlags() {
+        return Stream.of(KNN_LAUNCH_QUERY_REWRITE_ENABLED_SETTING).collect(Collectors.toUnmodifiableList());
+    }
+
+    public static boolean isKnnQueryRewriteEnabled() {
+        return Boolean.parseBoolean(KNNSettings.state().getSettingValue(KNN_LAUNCH_QUERY_REWRITE_ENABLED).toString());
+    }
+}
diff --git a/src/main/java/org/opensearch/knn/index/KNNSettings.java b/src/main/java/org/opensearch/knn/index/KNNSettings.java
@@ -9,17 +9,17 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.OpenSearchParseException;
-import org.opensearch.cluster.metadata.IndexMetadata;
-import org.opensearch.core.action.ActionListener;
 import org.opensearch.action.admin.cluster.settings.ClusterUpdateSettingsRequest;
 import org.opensearch.action.admin.cluster.settings.ClusterUpdateSettingsResponse;
 import org.opensearch.client.Client;
+import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.settings.Setting;
 import org.opensearch.common.settings.Settings;
+import org.opensearch.common.unit.TimeValue;
+import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.unit.ByteSizeUnit;
 import org.opensearch.core.common.unit.ByteSizeValue;
-import org.opensearch.common.unit.TimeValue;
 import org.opensearch.index.IndexModule;
 import org.opensearch.knn.index.memory.NativeMemoryCacheManager;
 import org.opensearch.knn.index.memory.NativeMemoryCacheManagerDto;
@@ -28,20 +28,22 @@
 import org.opensearch.monitor.os.OsProbe;
 
 import java.security.InvalidParameterException;
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
+import static java.util.stream.Collectors.toUnmodifiableMap;
 import static org.opensearch.common.settings.Setting.Property.Dynamic;
 import static org.opensearch.common.settings.Setting.Property.IndexScope;
 import static org.opensearch.common.settings.Setting.Property.NodeScope;
-import static org.opensearch.core.common.unit.ByteSizeValue.parseBytesSizeValue;
 import static org.opensearch.common.unit.MemorySizeValue.parseBytesSizeValueOrHeapRatio;
+import static org.opensearch.core.common.unit.ByteSizeValue.parseBytesSizeValue;
+import static org.opensearch.knn.common.featureflags.KNNFeatureFlags.getFeatureFlags;
 
 /**
  * This class defines
@@ -80,6 +82,12 @@ public class KNNSettings {
     public static final String MODEL_CACHE_SIZE_LIMIT = "knn.model.cache.size.limit";
     public static final String ADVANCED_FILTERED_EXACT_SEARCH_THRESHOLD = "index.knn.advanced.filtered_exact_search_threshold";
     public static final String KNN_FAISS_AVX2_DISABLED = "knn.faiss.avx2.disabled";
+    /**
+     * TODO: This setting is only added to ensure that main branch of k_NN plugin doesn't break till other parts of the
+     * code is getting ready. Will remove this setting once all changes related to integration of KNNVectorsFormat is added
+     * for native engines.
+     */
+    public static final String KNN_USE_LUCENE_VECTOR_FORMAT_ENABLED = "knn.use.format.enabled";
 
     /**
      * Default setting values
@@ -253,6 +261,17 @@ public class KNNSettings {
         NodeScope
     );
 
+    /**
+     * TODO: This setting is only added to ensure that main branch of k_NN plugin doesn't break till other parts of the
+     * code is getting ready. Will remove this setting once all changes related to integration of KNNVectorsFormat is added
+     * for native engines.
+     */
+    public static final Setting<Boolean> KNN_USE_LUCENE_VECTOR_FORMAT_ENABLED_SETTING = Setting.boolSetting(
+        KNN_USE_LUCENE_VECTOR_FORMAT_ENABLED,
+        false,
+        NodeScope
+    );
+
     /**
      * Dynamic settings
      */
@@ -289,6 +308,9 @@ public class KNNSettings {
         }
     };
 
+    private final static Map<String, Setting<?>> FEATURE_FLAGS = getFeatureFlags().stream()
+        .collect(toUnmodifiableMap(Setting::getKey, Function.identity()));
+
     private ClusterService clusterService;
     private Client client;
 
@@ -326,7 +348,7 @@ private void setSettingsUpdateConsumers() {
             );
 
             NativeMemoryCacheManager.getInstance().rebuildCache(builder.build());
-        }, new ArrayList<>(dynamicCacheSettings.values()));
+        }, Stream.concat(dynamicCacheSettings.values().stream(), FEATURE_FLAGS.values().stream()).collect(Collectors.toUnmodifiableList()));
     }
 
     /**
@@ -346,6 +368,10 @@ private Setting<?> getSetting(String key) {
             return dynamicCacheSettings.get(key);
         }
 
+        if (FEATURE_FLAGS.containsKey(key)) {
+            return FEATURE_FLAGS.get(key);
+        }
+
         if (KNN_CIRCUIT_BREAKER_TRIGGERED.equals(key)) {
             return KNN_CIRCUIT_BREAKER_TRIGGERED_SETTING;
         }
@@ -370,6 +396,10 @@ private Setting<?> getSetting(String key) {
             return KNN_VECTOR_STREAMING_MEMORY_LIMIT_PCT_SETTING;
         }
 
+        if (KNN_USE_LUCENE_VECTOR_FORMAT_ENABLED.equals(key)) {
+            return KNN_USE_LUCENE_VECTOR_FORMAT_ENABLED_SETTING;
+        }
+
         throw new IllegalArgumentException("Cannot find setting by key [" + key + "]");
     }
 
@@ -388,9 +418,11 @@ public List<Setting<?>> getSettings() {
             MODEL_CACHE_SIZE_LIMIT_SETTING,
             ADVANCED_FILTERED_EXACT_SEARCH_THRESHOLD_SETTING,
             KNN_FAISS_AVX2_DISABLED_SETTING,
-            KNN_VECTOR_STREAMING_MEMORY_LIMIT_PCT_SETTING
+            KNN_VECTOR_STREAMING_MEMORY_LIMIT_PCT_SETTING,
+            KNN_USE_LUCENE_VECTOR_FORMAT_ENABLED_SETTING
         );
-        return Stream.concat(settings.stream(), dynamicCacheSettings.values().stream()).collect(Collectors.toList());
+        return Stream.concat(settings.stream(), Stream.concat(getFeatureFlags().stream(), dynamicCacheSettings.values().stream()))
+            .collect(Collectors.toList());
     }
 
     public static boolean isKNNPluginEnabled() {
@@ -433,6 +465,15 @@ public static Integer getFilteredExactSearchThreshold(final String indexName) {
             .getAsInt(ADVANCED_FILTERED_EXACT_SEARCH_THRESHOLD, ADVANCED_FILTERED_EXACT_SEARCH_THRESHOLD_DEFAULT_VALUE);
     }
 
+    /**
+     * TODO: This setting is only added to ensure that main branch of k_NN plugin doesn't break till other parts of the
+     * code is getting ready. Will remove this setting once all changes related to integration of KNNVectorsFormat is added
+     * for native engines.
+     */
+    public static boolean getIsLuceneVectorFormatEnabled() {
+        return KNNSettings.state().getSettingValue(KNNSettings.KNN_USE_LUCENE_VECTOR_FORMAT_ENABLED);
+    }
+
     public void initialize(Client client, ClusterService clusterService) {
         this.client = client;
         this.clusterService = clusterService;

diff --git a/src/main/java/org/opensearch/knn/index/codec/BasePerFieldKnnVectorsFormat.java b/src/main/java/org/opensearch/knn/index/codec/BasePerFieldKnnVectorsFormat.java
@@ -13,6 +13,8 @@
 import org.opensearch.knn.index.codec.params.KNNScalarQuantizedVectorsFormatParams;
 import org.opensearch.knn.index.codec.params.KNNVectorsFormatParams;
 import org.opensearch.knn.index.engine.KNNEngine;
+import org.opensearch.knn.index.engine.KNNMethodContext;
+import org.opensearch.knn.index.mapper.KNNMappingConfig;
 import org.opensearch.knn.index.mapper.KNNVectorFieldType;
 
 import java.util.Optional;
@@ -66,16 +68,19 @@ public KnnVectorsFormat getKnnVectorsFormatForField(final String field) {
             );
             return defaultFormatSupplier.get();
         }
-        var type = (KNNVectorFieldType) mapperService.orElseThrow(
+        KNNVectorFieldType mappedFieldType = (KNNVectorFieldType) mapperService.orElseThrow(
             () -> new IllegalStateException(
                 String.format("Cannot read field type for field [%s] because mapper service is not available", field)
             )
         ).fieldType(field);
-        var params = type.getKnnMethodContext().getMethodComponentContext().getParameters();
 
-        if (type.getKnnMethodContext().getKnnEngine() == KNNEngine.LUCENE
-            && params != null
-            && params.containsKey(METHOD_ENCODER_PARAMETER)) {
+        KNNMappingConfig knnMappingConfig = mappedFieldType.getKnnMappingConfig();
+        KNNMethodContext knnMethodContext = knnMappingConfig.getKnnMethodContext()
+            .orElseThrow(() -> new IllegalArgumentException("KNN method context cannot be empty"));
+
+        var params = knnMethodContext.getMethodComponentContext().getParameters();
+
+        if (knnMethodContext.getKnnEngine() == KNNEngine.LUCENE && params != null && params.containsKey(METHOD_ENCODER_PARAMETER)) {
             KNNScalarQuantizedVectorsFormatParams knnScalarQuantizedVectorsFormatParams = new KNNScalarQuantizedVectorsFormatParams(
                 params,
                 defaultMaxConnections,

diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java
@@ -14,14 +14,14 @@
 import org.opensearch.core.xcontent.MediaTypeRegistry;
 import org.opensearch.core.xcontent.DeprecationHandler;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
+import org.opensearch.knn.index.SpaceType;
 import org.opensearch.knn.index.util.IndexUtil;
 import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.index.VectorDataType;
 import org.opensearch.knn.index.codec.transfer.VectorTransfer;
 import org.opensearch.knn.index.codec.transfer.VectorTransferByte;
 import org.opensearch.knn.index.codec.transfer.VectorTransferFloat;
 import org.opensearch.knn.jni.JNIService;
-import org.opensearch.knn.index.SpaceType;
 import org.opensearch.knn.index.codec.util.KNNCodecUtil;
 import org.opensearch.knn.index.engine.KNNEngine;
 import org.opensearch.knn.indices.Model;
@@ -131,19 +131,20 @@ public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer,
         NativeIndexCreator indexCreator;
         KNNCodecUtil.Pair pair;
         Map<String, String> fieldAttributes = field.attributes();
+        VectorDataType vectorDataType;
 
         if (fieldAttributes.containsKey(MODEL_ID)) {
             String modelId = fieldAttributes.get(MODEL_ID);
             Model model = ModelCache.getInstance().get(modelId);
             if (model.getModelBlob() == null) {
                 throw new RuntimeException(String.format("There is no trained model with id \"%s\"", modelId));
             }
-            VectorDataType vectorDataType = model.getModelMetadata().getVectorDataType();
+            vectorDataType = model.getModelMetadata().getVectorDataType();
             pair = KNNCodecUtil.getPair(values, getVectorTransfer(vectorDataType));
             indexCreator = () -> createKNNIndexFromTemplate(model, pair, knnEngine, indexPath);
         } else {
             // get vector data type from field attributes or provide default value
-            VectorDataType vectorDataType = VectorDataType.get(
+            vectorDataType = VectorDataType.get(
                 fieldAttributes.getOrDefault(KNNConstants.VECTOR_DATA_TYPE_FIELD, VectorDataType.DEFAULT.getValue())
             );
             pair = KNNCodecUtil.getPair(values, getVectorTransfer(vectorDataType));
@@ -156,7 +157,7 @@ public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer,
             return;
         }
 
-        long arraySize = calculateArraySize(pair.docs.length, pair.getDimension(), pair.serializationMode);
+        long arraySize = calculateArraySize(pair.docs.length, pair.getDimension(), vectorDataType);
 
         if (isMerge) {
             KNNGraphValue.MERGE_CURRENT_OPERATIONS.increment();
@@ -215,8 +216,7 @@ private void createKNNIndexFromScratch(FieldInfo fieldInfo, KNNCodecUtil.Pair pa
         throws IOException {
         Map<String, Object> parameters = new HashMap<>();
         Map<String, String> fieldAttributes = fieldInfo.attributes();
-        String parametersString = fieldAttributes.get(KNNConstants.PARAMETERS);
-
+        String parametersString = fieldAttributes.get(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()));