Add new cluster setting for keyword indexordocvalues query

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>

server/src/main/java/org/opensearch/common/settings/ClusterSettings.java

@@ -549,6 +549,7 @@ public void apply(Settings value, Settings current, Settings previous) {
                 SearchService.MAX_AGGREGATION_REWRITE_FILTERS,
                 SearchService.INDICES_MAX_CLAUSE_COUNT_SETTING,
                 SearchService.CARDINALITY_AGGREGATION_PRUNING_THRESHOLD,
+                SearchService.KEYWORD_INDEX_OR_DOC_VALUES_ENABLED,
                 CreatePitController.PIT_INIT_KEEP_ALIVE,
                 Node.WRITE_PORTS_FILE_SETTING,
                 Node.NODE_NAME_SETTING,

server/src/main/java/org/opensearch/index/IndexService.java

@@ -853,7 +853,7 @@ public IndexSettings getIndexSettings() {
      * {@link IndexReader}-specific optimizations, such as rewriting containing range queries.
      */
     public QueryShardContext newQueryShardContext(int shardId, IndexSearcher searcher, LongSupplier nowInMillis, String clusterAlias) {
-        return newQueryShardContext(shardId, searcher, nowInMillis, clusterAlias, false);
+        return newQueryShardContext(shardId, searcher, nowInMillis, clusterAlias, false, false);
     }
 
     /**
@@ -867,7 +867,8 @@ public QueryShardContext newQueryShardContext(
         IndexSearcher searcher,
         LongSupplier nowInMillis,
         String clusterAlias,
-        boolean validate
+        boolean validate,
+        boolean keywordIndexOrDocValuesEnabled
     ) {
         final SearchIndexNameMatcher indexNameMatcher = new SearchIndexNameMatcher(
             index().getName(),
@@ -893,7 +894,8 @@ public QueryShardContext newQueryShardContext(
             indexNameMatcher,
             allowExpensiveQueries,
             valuesSourceRegistry,
-            validate
+            validate,
+            keywordIndexOrDocValuesEnabled
         );
     }
 

server/src/main/java/org/opensearch/index/mapper/KeywordFieldMapper.java

@@ -392,6 +392,9 @@ public Query termsQuery(List<?> values, QueryShardContext context) {
             failIfNotIndexedAndNoDocValues();
             // has index and doc_values enabled
             if (isSearchable() && hasDocValues()) {
+                if (!context.keywordFieldIndexOrDocValuesEnabled()) {
+                    return super.termsQuery(values, context);
+                }
                 BytesRef[] bytesRefs = new BytesRef[values.size()];
                 for (int i = 0; i < bytesRefs.length; i++) {
                     bytesRefs[i] = indexedValueForSearch(values.get(i));
@@ -429,6 +432,9 @@ public Query prefixQuery(
             }
             failIfNotIndexedAndNoDocValues();
             if (isSearchable() && hasDocValues()) {
+                if (!context.keywordFieldIndexOrDocValuesEnabled()) {
+                    return super.prefixQuery(value, method, caseInsensitive, context);
+                }
                 Query indexQuery = super.prefixQuery(value, method, caseInsensitive, context);
                 Query dvQuery = super.prefixQuery(value, MultiTermQuery.DOC_VALUES_REWRITE, caseInsensitive, context);
                 return new IndexOrDocValuesQuery(indexQuery, dvQuery);
@@ -461,6 +467,9 @@ public Query regexpQuery(
             }
             failIfNotIndexedAndNoDocValues();
             if (isSearchable() && hasDocValues()) {
+                if (!context.keywordFieldIndexOrDocValuesEnabled()) {
+                    return super.regexpQuery(value, syntaxFlags, matchFlags, maxDeterminizedStates, method, context);
+                }
                 Query indexQuery = super.regexpQuery(value, syntaxFlags, matchFlags, maxDeterminizedStates, method, context);
                 Query dvQuery = super.regexpQuery(
                     value,
@@ -549,6 +558,9 @@ public Query fuzzyQuery(
                 );
             }
             if (isSearchable() && hasDocValues()) {
+                if (!context.keywordFieldIndexOrDocValuesEnabled()) {
+                    return super.fuzzyQuery(value, fuzziness, prefixLength, maxExpansions, transpositions, method, context);
+                }
                 Query indexQuery = super.fuzzyQuery(value, fuzziness, prefixLength, maxExpansions, transpositions, method, context);
                 Query dvQuery = super.fuzzyQuery(
                     value,
@@ -591,6 +603,9 @@ public Query wildcardQuery(
             // wildcard
             // query text
             if (isSearchable() && hasDocValues()) {
+                if (!context.keywordFieldIndexOrDocValuesEnabled()) {
+                    return super.wildcardQuery(value, method, caseInsensitive, true, context);
+                }
                 Query indexQuery = super.wildcardQuery(value, method, caseInsensitive, true, context);
                 Query dvQuery = super.wildcardQuery(value, MultiTermQuery.DOC_VALUES_REWRITE, caseInsensitive, true, context);
                 return new IndexOrDocValuesQuery(indexQuery, dvQuery);

server/src/main/java/org/opensearch/index/query/QueryShardContext.java

@@ -125,6 +125,7 @@ public class QueryShardContext extends QueryRewriteContext {
     private final ValuesSourceRegistry valuesSourceRegistry;
     private BitSetProducer parentFilter;
     private DerivedFieldResolver derivedFieldResolver;
+    private boolean keywordIndexOrDocValuesEnabled;
 
     public QueryShardContext(
         int shardId,
@@ -208,7 +209,55 @@ public QueryShardContext(
             ),
             allowExpensiveQueries,
             valuesSourceRegistry,
-            validate
+            validate,
+            false
+        );
+    }
+
+    public QueryShardContext(
+        int shardId,
+        IndexSettings indexSettings,
+        BigArrays bigArrays,
+        BitsetFilterCache bitsetFilterCache,
+        TriFunction<MappedFieldType, String, Supplier<SearchLookup>, IndexFieldData<?>> indexFieldDataLookup,
+        MapperService mapperService,
+        SimilarityService similarityService,
+        ScriptService scriptService,
+        NamedXContentRegistry xContentRegistry,
+        NamedWriteableRegistry namedWriteableRegistry,
+        Client client,
+        IndexSearcher searcher,
+        LongSupplier nowInMillis,
+        String clusterAlias,
+        Predicate<String> indexNameMatcher,
+        BooleanSupplier allowExpensiveQueries,
+        ValuesSourceRegistry valuesSourceRegistry,
+        boolean validate,
+        boolean keywordIndexOrDocValuesEnabled
+    ) {
+        this(
+            shardId,
+            indexSettings,
+            bigArrays,
+            bitsetFilterCache,
+            indexFieldDataLookup,
+            mapperService,
+            similarityService,
+            scriptService,
+            xContentRegistry,
+            namedWriteableRegistry,
+            client,
+            searcher,
+            nowInMillis,
+            indexNameMatcher,
+            new Index(
+                RemoteClusterAware.buildRemoteIndexName(clusterAlias, indexSettings.getIndex().getName()),
+                indexSettings.getIndex().getUUID()
+            ),
+            allowExpensiveQueries,
+            valuesSourceRegistry,
+            validate,
+            keywordIndexOrDocValuesEnabled
         );
     }
 
@@ -231,7 +280,8 @@ public QueryShardContext(QueryShardContext source) {
             source.fullyQualifiedIndex,
             source.allowExpensiveQueries,
             source.valuesSourceRegistry,
-            source.validate()
+            source.validate(),
+            source.keywordIndexOrDocValuesEnabled
         );
     }
 
@@ -253,7 +303,8 @@ private QueryShardContext(
         Index fullyQualifiedIndex,
         BooleanSupplier allowExpensiveQueries,
         ValuesSourceRegistry valuesSourceRegistry,
-        boolean validate
+        boolean validate,
+        boolean keywordIndexOrDocValuesEnabled
     ) {
         super(xContentRegistry, namedWriteableRegistry, client, nowInMillis, validate);
         this.shardId = shardId;
@@ -277,6 +328,7 @@ private QueryShardContext(
             emptyList(),
             indexSettings.isDerivedFieldAllowed()
         );
+        this.keywordIndexOrDocValuesEnabled = keywordIndexOrDocValuesEnabled;
     }
 
     private void reset() {
@@ -414,6 +466,14 @@ public void setDerivedFieldResolver(DerivedFieldResolver derivedFieldResolver) {
         this.derivedFieldResolver = derivedFieldResolver;
     }
 
+    public boolean keywordFieldIndexOrDocValuesEnabled() {
+        return keywordIndexOrDocValuesEnabled;
+    }
+
+    public void setKeywordFieldIndexOrDocValuesEnabled(boolean keywordIndexOrDocValuesEnabled) {
+        this.keywordIndexOrDocValuesEnabled = keywordIndexOrDocValuesEnabled;
+    }
+
     public void setAllowUnmappedFields(boolean allowUnmappedFields) {
         this.allowUnmappedFields = allowUnmappedFields;
     }

server/src/main/java/org/opensearch/search/DefaultSearchContext.java

@@ -120,6 +120,7 @@
 import static org.opensearch.search.SearchService.CONCURRENT_SEGMENT_SEARCH_MODE_ALL;
 import static org.opensearch.search.SearchService.CONCURRENT_SEGMENT_SEARCH_MODE_AUTO;
 import static org.opensearch.search.SearchService.CONCURRENT_SEGMENT_SEARCH_MODE_NONE;
+import static org.opensearch.search.SearchService.KEYWORD_INDEX_OR_DOC_VALUES_ENABLED;
 import static org.opensearch.search.SearchService.MAX_AGGREGATION_REWRITE_FILTERS;
 
 /**
@@ -206,6 +207,7 @@ final class DefaultSearchContext extends SearchContext {
     private final SetOnce<Boolean> requestShouldUseConcurrentSearch = new SetOnce<>();
     private final int maxAggRewriteFilters;
     private final int cardinalityAggregationPruningThreshold;
+    private final boolean keywordIndexOrDocValuesEnabled;
 
     DefaultSearchContext(
         ReaderContext readerContext,
@@ -256,7 +258,8 @@ final class DefaultSearchContext extends SearchContext {
             this.searcher,
             request::nowInMillis,
             shardTarget.getClusterAlias(),
-            validate
+            validate,
+            evaluateKeywordIndexOrDocValuesEnabled()
         );
         queryBoost = request.indexBoost();
         this.lowLevelCancellation = lowLevelCancellation;
@@ -265,6 +268,7 @@ final class DefaultSearchContext extends SearchContext {
         this.maxAggRewriteFilters = evaluateFilterRewriteSetting();
         this.cardinalityAggregationPruningThreshold = evaluateCardinalityAggregationPruningThreshold();
         this.concurrentSearchDeciders = concurrentSearchDeciders;
+        this.keywordIndexOrDocValuesEnabled = evaluateKeywordIndexOrDocValuesEnabled();
     }
 
     @Override
@@ -1117,10 +1121,22 @@ public int cardinalityAggregationPruningThreshold() {
         return cardinalityAggregationPruningThreshold;
     }
 
+    @Override
+    public boolean keywordIndexOrDocValuesEnabled() {
+        return keywordIndexOrDocValuesEnabled;
+    }
+
     private int evaluateCardinalityAggregationPruningThreshold() {
         if (clusterService != null) {
             return clusterService.getClusterSettings().get(CARDINALITY_AGGREGATION_PRUNING_THRESHOLD);
         }
         return 0;
     }
+
+    public boolean evaluateKeywordIndexOrDocValuesEnabled() {
+        if (clusterService != null) {
+            return clusterService.getClusterSettings().get(KEYWORD_INDEX_OR_DOC_VALUES_ENABLED);
+        }
+        return false;
+    }
 }

server/src/main/java/org/opensearch/search/SearchService.java

@@ -337,6 +337,13 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
         Property.NodeScope
     );
 
+    public static final Setting<Boolean> KEYWORD_INDEX_OR_DOC_VALUES_ENABLED = Setting.boolSetting(
+        "search.keyword_index_or_doc_values_enabled",
+        false,
+        Property.Dynamic,
+        Property.NodeScope
+    );
+
     public static final int DEFAULT_SIZE = 10;
     public static final int DEFAULT_FROM = 0;
 
@@ -1173,6 +1180,7 @@ private DefaultSearchContext createSearchContext(ReaderContext reader, ShardSear
                 context.getIndexSettings().isDerivedFieldAllowed() && allowDerivedField
             );
             context.setDerivedFieldResolver(derivedFieldResolver);
+            context.setKeywordFieldIndexOrDocValuesEnabled(searchContext.keywordIndexOrDocValuesEnabled());
             searchContext.getQueryShardContext().setDerivedFieldResolver(derivedFieldResolver);
             Rewriteable.rewrite(request.getRewriteable(), context, true);
             assert searchContext.getQueryShardContext().isCacheable();

server/src/main/java/org/opensearch/search/internal/SearchContext.java

@@ -526,4 +526,9 @@ public int maxAggRewriteFilters() {
     public int cardinalityAggregationPruningThreshold() {
         return 0;
     }
+
+    public boolean keywordIndexOrDocValuesEnabled() {
+        return false;
+    }
+
 }

server/src/test/java/org/opensearch/index/mapper/KeywordFieldTypeTests.java

@@ -136,7 +136,7 @@ public void testTermsQuery() {
             new TermInSetQuery("field", terms),
             new TermInSetQuery(MultiTermQuery.DOC_VALUES_REWRITE, "field", terms)
         );
-        assertEquals(expected, ft.termsQuery(Arrays.asList("foo", "bar"), null));
+        assertEquals(expected, ft.termsQuery(Arrays.asList("foo", "bar"), MOCK_QSC_ENABLE_INDEX_DOC_VALUES));
 
         MappedFieldType onlyIndexed = new KeywordFieldType("field", true, false, Collections.emptyMap());
         Query expectedIndex = new TermInSetQuery("field", terms);
@@ -225,7 +225,7 @@ public void testRegexpQuery() {
                 new RegexpQuery(new Term("field", "foo.*")),
                 new RegexpQuery(new Term("field", "foo.*"), 0, 0, RegexpQuery.DEFAULT_PROVIDER, 10, MultiTermQuery.DOC_VALUES_REWRITE)
             ),
-            ft.regexpQuery("foo.*", 0, 0, 10, MultiTermQuery.CONSTANT_SCORE_BLENDED_REWRITE, MOCK_QSC)
+            ft.regexpQuery("foo.*", 0, 0, 10, MultiTermQuery.CONSTANT_SCORE_BLENDED_REWRITE, MOCK_QSC_ENABLE_INDEX_DOC_VALUES)
         );
 
         Query indexExpected = new RegexpQuery(new Term("field", "foo.*"));
@@ -267,7 +267,7 @@ public void testFuzzyQuery() {
                 new FuzzyQuery(new Term("field", "foo"), 2, 1, 50, true),
                 new FuzzyQuery(new Term("field", "foo"), 2, 1, 50, true, MultiTermQuery.DOC_VALUES_REWRITE)
             ),
-            ft.fuzzyQuery("foo", Fuzziness.fromEdits(2), 1, 50, true, null, MOCK_QSC)
+            ft.fuzzyQuery("foo", Fuzziness.fromEdits(2), 1, 50, true, null, MOCK_QSC_ENABLE_INDEX_DOC_VALUES)
         );
 
         Query indexExpected = new FuzzyQuery(new Term("field", "foo"), 2, 1, 50, true);
@@ -308,7 +308,7 @@ public void testWildCardQuery() {
                 MultiTermQuery.DOC_VALUES_REWRITE
             )
         );
-        assertEquals(expected, ft.wildcardQuery("foo*", MultiTermQuery.CONSTANT_SCORE_BLENDED_REWRITE, MOCK_QSC));
+        assertEquals(expected, ft.wildcardQuery("foo*", MultiTermQuery.CONSTANT_SCORE_BLENDED_REWRITE, MOCK_QSC_ENABLE_INDEX_DOC_VALUES));
 
         Query indexExpected = new WildcardQuery(new Term("field", new BytesRef("foo*")));
         MappedFieldType onlyIndexed = new KeywordFieldType("field", true, false, Collections.emptyMap());

server/src/test/java/org/opensearch/search/DefaultSearchContextTests.java

@@ -168,9 +168,8 @@ public void testPreProcess() throws Exception {
         when(indexCache.query()).thenReturn(queryCache);
         when(indexService.cache()).thenReturn(indexCache);
         QueryShardContext queryShardContext = mock(QueryShardContext.class);
-        when(indexService.newQueryShardContext(eq(shardId.id()), any(), any(), nullable(String.class), anyBoolean())).thenReturn(
-            queryShardContext
-        );
+        when(indexService.newQueryShardContext(eq(shardId.id()), any(), any(), nullable(String.class), anyBoolean(), anyBoolean()))
+            .thenReturn(queryShardContext);
         MapperService mapperService = mock(MapperService.class);
         when(mapperService.hasNested()).thenReturn(randomBoolean());
         when(indexService.mapperService()).thenReturn(mapperService);
@@ -501,9 +500,8 @@ public void testClearQueryCancellationsOnClose() throws IOException {
 
         IndexService indexService = mock(IndexService.class);
         QueryShardContext queryShardContext = mock(QueryShardContext.class);
-        when(indexService.newQueryShardContext(eq(shardId.id()), any(), any(), nullable(String.class), anyBoolean())).thenReturn(
-            queryShardContext
-        );
+        when(indexService.newQueryShardContext(eq(shardId.id()), any(), any(), nullable(String.class), anyBoolean(), anyBoolean()))
+            .thenReturn(queryShardContext);
         Settings settings = Settings.builder()
             .put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT)
             .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1)
@@ -598,9 +596,8 @@ public void testSearchPathEvaluation() throws Exception {
 
         IndexService indexService = mock(IndexService.class);
         QueryShardContext queryShardContext = mock(QueryShardContext.class);
-        when(indexService.newQueryShardContext(eq(shardId.id()), any(), any(), nullable(String.class), anyBoolean())).thenReturn(
-            queryShardContext
-        );
+        when(indexService.newQueryShardContext(eq(shardId.id()), any(), any(), nullable(String.class), anyBoolean(), anyBoolean()))
+            .thenReturn(queryShardContext);
 
         IndexMetadata indexMetadata = IndexMetadata.builder("index").settings(settings).build();
         IndexSettings indexSettings = new IndexSettings(indexMetadata, Settings.EMPTY);
@@ -830,7 +827,7 @@ public void testSearchPathEvaluationWithConcurrentSearchModeAsAuto() throws Exce
 
         IndexService indexService = mock(IndexService.class);
         QueryShardContext queryShardContext = mock(QueryShardContext.class);
-        when(indexService.newQueryShardContext(eq(shardId.id()), any(), any(), nullable(String.class), anyBoolean())).thenReturn(
+        when(indexService.newQueryShardContext(eq(shardId.id()), any(), any(), nullable(String.class), anyBoolean(), any())).thenReturn(
             queryShardContext
         );
 

test/framework/src/main/java/org/opensearch/index/mapper/FieldTypeTestCase.java

@@ -46,16 +46,18 @@
 /** Base test case for subclasses of MappedFieldType */
 public abstract class FieldTypeTestCase extends OpenSearchTestCase {
 
-    public static final QueryShardContext MOCK_QSC = createMockQueryShardContext(true);
-    public static final QueryShardContext MOCK_QSC_DISALLOW_EXPENSIVE = createMockQueryShardContext(false);
+    public static final QueryShardContext MOCK_QSC = createMockQueryShardContext(true, false);
+    public static final QueryShardContext MOCK_QSC_DISALLOW_EXPENSIVE = createMockQueryShardContext(false, false);
+    public static final QueryShardContext MOCK_QSC_ENABLE_INDEX_DOC_VALUES = createMockQueryShardContext(true, true);
 
     protected QueryShardContext randomMockShardContext() {
         return randomFrom(MOCK_QSC, MOCK_QSC_DISALLOW_EXPENSIVE);
     }
 
-    static QueryShardContext createMockQueryShardContext(boolean allowExpensiveQueries) {
+    static QueryShardContext createMockQueryShardContext(boolean allowExpensiveQueries, boolean keywordIndexOrDocValuesEnabled) {
         QueryShardContext queryShardContext = mock(QueryShardContext.class);
         when(queryShardContext.allowExpensiveQueries()).thenReturn(allowExpensiveQueries);
+        when(queryShardContext.keywordFieldIndexOrDocValuesEnabled()).thenReturn(keywordIndexOrDocValuesEnabled);
         return queryShardContext;
     }