Adding multiple query types

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>
harshavamsi · Aug 8, 2024 · ce8a043 · ce8a043
1 parent de9e944
commit ce8a043
Show file tree

Hide file tree

Showing 34 changed files with 508 additions and 107 deletions.
diff --git a/modules/mapper-extras/src/main/java/org/opensearch/index/mapper/ScaledFloatFieldMapper.java b/modules/mapper-extras/src/main/java/org/opensearch/index/mapper/ScaledFloatFieldMapper.java
@@ -44,6 +44,7 @@
 import org.apache.lucene.search.Query;
 import org.apache.lucene.util.BytesRef;
 import org.opensearch.common.Explicit;
+import org.opensearch.common.Nullable;
 import org.opensearch.common.settings.Setting;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.xcontent.support.XContentMapValues;
@@ -225,7 +226,7 @@ public Query termQuery(Object value, QueryShardContext context) {
         }
 
         @Override
-        public Query termsQuery(List<?> values, QueryShardContext context) {
+        public Query termsQuery(List<?> values, @Nullable RewriteOverride rewriteOverride, QueryShardContext context) {
             failIfNotIndexedAndNoDocValues();
             List<Long> scaledValues = new ArrayList<>(values.size());
             for (Object value : values) {
@@ -245,7 +246,14 @@ public Query termsQuery(List<?> values, QueryShardContext context) {
         }
 
         @Override
-        public Query rangeQuery(Object lowerTerm, Object upperTerm, boolean includeLower, boolean includeUpper, QueryShardContext context) {
+        public Query rangeQuery(
+            Object lowerTerm,
+            Object upperTerm,
+            boolean includeLower,
+            boolean includeUpper,
+            @Nullable RewriteOverride rewriteOverride,
+            QueryShardContext context
+        ) {
             failIfNotIndexedAndNoDocValues();
             Long lo = null;
             if (lowerTerm != null) {

diff --git a/...s/mapper-extras/src/main/java/org/opensearch/index/mapper/SearchAsYouTypeFieldMapper.java b/...s/mapper-extras/src/main/java/org/opensearch/index/mapper/SearchAsYouTypeFieldMapper.java
@@ -59,6 +59,7 @@
 import org.apache.lucene.util.automaton.Automata;
 import org.apache.lucene.util.automaton.Automaton;
 import org.apache.lucene.util.automaton.Operations;
+import org.opensearch.common.Nullable;
 import org.opensearch.common.collect.Iterators;
 import org.opensearch.common.lucene.search.AutomatonQueries;
 import org.opensearch.index.analysis.AnalyzerScope;
@@ -335,7 +336,13 @@ public ValueFetcher valueFetcher(QueryShardContext context, SearchLookup searchL
         }
 
         @Override
-        public Query prefixQuery(String value, MultiTermQuery.RewriteMethod method, boolean caseInsensitive, QueryShardContext context) {
+        public Query prefixQuery(
+            String value,
+            MultiTermQuery.RewriteMethod method,
+            @Nullable RewriteOverride rewriteOverride,
+            boolean caseInsensitive,
+            QueryShardContext context
+        ) {
             if (prefixField == null || prefixField.termLengthWithinBounds(value.length()) == false) {
                 return super.prefixQuery(value, method, caseInsensitive, context);
             } else {
@@ -419,7 +426,13 @@ boolean termLengthWithinBounds(int length) {
         }
 
         @Override
-        public Query prefixQuery(String value, MultiTermQuery.RewriteMethod method, boolean caseInsensitive, QueryShardContext context) {
+        public Query prefixQuery(
+            String value,
+            MultiTermQuery.RewriteMethod method,
+            @Nullable RewriteOverride rewriteOverride,
+            boolean caseInsensitive,
+            QueryShardContext context
+        ) {
             if (value.length() >= minChars) {
                 if (caseInsensitive) {
                     return super.termQueryCaseInsensitive(value, context);
@@ -557,7 +570,13 @@ public String typeName() {
         }
 
         @Override
-        public Query prefixQuery(String value, MultiTermQuery.RewriteMethod method, boolean caseInsensitive, QueryShardContext context) {
+        public Query prefixQuery(
+            String value,
+            MultiTermQuery.RewriteMethod method,
+            @Nullable RewriteOverride rewriteOverride,
+            boolean caseInsensitive,
+            QueryShardContext context
+        ) {
             if (prefixFieldType == null || prefixFieldType.termLengthWithinBounds(value.length()) == false) {
                 return super.prefixQuery(value, method, caseInsensitive, context);
             } else {

diff --git a/rest-api-spec/src/main/resources/rest-api-spec/test/search/340_doc_values_field.yml b/rest-api-spec/src/main/resources/rest-api-spec/test/search/340_doc_values_field.yml
@@ -82,6 +82,21 @@
 
   - match: { hits.hits.0._source.some_keyword: "ingesting some random keyword data" }
 
+  - do:
+      search:
+        rest_total_hits_as_int: true
+        index: test-iodvq
+        body:
+          query:
+            prefix: {
+              some_keyword: {
+                              "value": "ing",
+                        "rewrite_override": "index_only"
+              }
+            }
+
+  - match: { hits.hits.0._source.some_keyword: "ingesting some random keyword data" }
+
   - do:
       search:
         rest_total_hits_as_int: true
@@ -95,6 +110,20 @@
 
   - match: { hits.total: 2 }
 
+  - do:
+      search:
+        rest_total_hits_as_int: true
+        index: test-iodvq
+        body:
+          query:
+            range: {
+              "some_keyword": {
+                "lt": 500,
+                "rewrite_override": "index_only"
+              } }
+
+  - match: { hits.total: 2 }
+
 
   - do:
       search:
@@ -619,6 +648,18 @@
 
   - match: { hits.hits.0._source.some_keyword: "ingesting some random keyword data" }
 
+  - do:
+      search:
+        rest_total_hits_as_int: true
+        index: test-index
+        body:
+          query:
+            prefix:
+              some_keyword: "ing"
+              rewrite_override: "index"
+
+  - match: { hits.hits.0._source.some_keyword: "ingesting some random keyword data" }
+
   - do:
       search:
         rest_total_hits_as_int: true

diff --git a/server/src/main/java/org/opensearch/index/mapper/BooleanFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/BooleanFieldMapper.java
@@ -280,7 +280,7 @@ public Query termQuery(Object value, QueryShardContext context) {
         }
 
         @Override
-        public Query termsQuery(List<?> values, QueryShardContext context) {
+        public Query termsQuery(List<?> values, @Nullable RewriteOverride rewriteOverride, QueryShardContext context) {
             failIfNotIndexedAndNoDocValues();
             int distinct = 0;
             Set<?> distinctValues = new HashSet<>(values);
@@ -305,7 +305,14 @@ public Query termsQuery(List<?> values, QueryShardContext context) {
         }
 
         @Override
-        public Query rangeQuery(Object lowerTerm, Object upperTerm, boolean includeLower, boolean includeUpper, QueryShardContext context) {
+        public Query rangeQuery(
+            Object lowerTerm,
+            Object upperTerm,
+            boolean includeLower,
+            boolean includeUpper,
+            @Nullable RewriteOverride rewriteOverride,
+            QueryShardContext context
+        ) {
             failIfNotIndexedAndNoDocValues();
             if (lowerTerm == null) {
                 lowerTerm = false;

diff --git a/server/src/main/java/org/opensearch/index/mapper/ConstantFieldType.java b/server/src/main/java/org/opensearch/index/mapper/ConstantFieldType.java
@@ -101,7 +101,7 @@ public final Query termQueryCaseInsensitive(Object value, QueryShardContext cont
     }
 
     @Override
-    public final Query termsQuery(List<?> values, QueryShardContext context) {
+    public final Query termsQuery(List<?> values, @Nullable RewriteOverride rewriteOverride, QueryShardContext context) {
         for (Object value : values) {
             String pattern = valueToString(value);
             if (matches(pattern, false, context)) {
@@ -116,6 +116,7 @@ public final Query termsQuery(List<?> values, QueryShardContext context) {
     public final Query prefixQuery(
         String prefix,
         @Nullable MultiTermQuery.RewriteMethod method,
+        @Nullable RewriteOverride rewriteOverride,
         boolean caseInsensitive,
         QueryShardContext context
     ) {
@@ -131,6 +132,7 @@ public final Query prefixQuery(
     public final Query wildcardQuery(
         String value,
         @Nullable MultiTermQuery.RewriteMethod method,
+        @Nullable RewriteOverride rewriteOverride,
         boolean caseInsensitive,
         QueryShardContext context
     ) {

diff --git a/server/src/main/java/org/opensearch/index/mapper/ConstantKeywordFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/ConstantKeywordFieldMapper.java
@@ -143,6 +143,7 @@ public Query rangeQuery(
             ShapeRelation relation,
             ZoneId timeZone,
             DateMathParser parser,
+            @Nullable RewriteOverride rewriteOverride,
             QueryShardContext context
         ) {
             if (lowerTerm != null) {
@@ -173,6 +174,7 @@ public Query regexpQuery(
             int matchFlags,
             int maxDeterminizedStates,
             @Nullable MultiTermQuery.RewriteMethod method,
+            @Nullable RewriteOverride rewriteOverride,
             QueryShardContext context
         ) {
             Automaton automaton = new RegExp(value, syntaxFlags, matchFlags).toAutomaton(

diff --git a/server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java
@@ -455,6 +455,7 @@ public Query rangeQuery(
             ShapeRelation relation,
             @Nullable ZoneId timeZone,
             @Nullable DateMathParser forcedDateParser,
+            @Nullable RewriteOverride rewriteOverride,
             QueryShardContext context
         ) {
             failIfNotIndexedAndNoDocValues();

diff --git a/server/src/main/java/org/opensearch/index/mapper/DerivedFieldType.java b/server/src/main/java/org/opensearch/index/mapper/DerivedFieldType.java
@@ -193,7 +193,7 @@ public Query termQueryCaseInsensitive(Object value, @Nullable QueryShardContext
     }
 
     @Override
-    public Query termsQuery(List<?> values, @Nullable QueryShardContext context) {
+    public Query termsQuery(List<?> values, @Nullable RewriteOverride rewriteOverride, @Nullable QueryShardContext context) {
         Query query = typeFieldMapper.mappedFieldType.termsQuery(values, context);
         DerivedFieldValueFetcher valueFetcher = valueFetcher(context, context.lookup(), null);
         DerivedFieldQuery derivedFieldQuery = new DerivedFieldQuery(
@@ -218,6 +218,7 @@ public Query rangeQuery(
         ShapeRelation relation,
         ZoneId timeZone,
         DateMathParser parser,
+        @Nullable RewriteOverride rewriteOverride,
         QueryShardContext context
     ) {
         Query query = typeFieldMapper.mappedFieldType.rangeQuery(
@@ -248,6 +249,7 @@ public Query fuzzyQuery(
         int prefixLength,
         int maxExpansions,
         boolean transpositions,
+        @Nullable RewriteOverride rewriteOverride,
         QueryShardContext context
     ) {
         Query query = typeFieldMapper.mappedFieldType.fuzzyQuery(value, fuzziness, prefixLength, maxExpansions, transpositions, context);
@@ -312,6 +314,7 @@ public Query fuzzyQuery(
     public Query prefixQuery(
         String value,
         @Nullable MultiTermQuery.RewriteMethod method,
+        @Nullable RewriteOverride rewriteOverride,
         boolean caseInsensitive,
         QueryShardContext context
     ) {
@@ -339,6 +342,7 @@ public Query prefixQuery(
     public Query wildcardQuery(
         String value,
         @Nullable MultiTermQuery.RewriteMethod method,
+        @Nullable RewriteOverride rewriteOverride,
         boolean caseInsensitive,
         QueryShardContext context
     ) {
@@ -391,6 +395,7 @@ public Query regexpQuery(
         int matchFlags,
         int maxDeterminizedStates,
         @Nullable MultiTermQuery.RewriteMethod method,
+        @Nullable RewriteOverride rewriteOverride,
         QueryShardContext context
     ) {
         Query query = typeFieldMapper.mappedFieldType.regexpQuery(value, syntaxFlags, matchFlags, maxDeterminizedStates, method, context);

diff --git a/server/src/main/java/org/opensearch/index/mapper/FlatObjectFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/FlatObjectFieldMapper.java
@@ -350,7 +350,7 @@ public Query termQuery(Object value, @Nullable QueryShardContext context) {
         }
 
         @Override
-        public Query termsQuery(List<?> values, QueryShardContext context) {
+        public Query termsQuery(List<?> values, @Nullable RewriteOverride rewriteOverride, QueryShardContext context) {
             failIfNotIndexed();
             String directedSearchFieldName = directSubfield();
             BytesRef[] bytesRefs = new BytesRef[values.size()];
@@ -446,7 +446,13 @@ private String inputToString(Object inputValue) {
         }
 
         @Override
-        public Query prefixQuery(String value, MultiTermQuery.RewriteMethod method, boolean caseInsensitive, QueryShardContext context) {
+        public Query prefixQuery(
+            String value,
+            MultiTermQuery.RewriteMethod method,
+            @Nullable RewriteOverride rewriteOverride,
+            boolean caseInsensitive,
+            QueryShardContext context
+        ) {
             String directSubfield = directSubfield();
             String rewriteValue = rewriteValue(value);
 
@@ -469,7 +475,14 @@ public Query prefixQuery(String value, MultiTermQuery.RewriteMethod method, bool
         }
 
         @Override
-        public Query rangeQuery(Object lowerTerm, Object upperTerm, boolean includeLower, boolean includeUpper, QueryShardContext context) {
+        public Query rangeQuery(
+            Object lowerTerm,
+            Object upperTerm,
+            boolean includeLower,
+            boolean includeUpper,
+            @Nullable RewriteOverride rewriteOverride,
+            QueryShardContext context
+        ) {
             String directSubfield = directSubfield();
             String rewriteUpperTerm = rewriteValue(inputToString(upperTerm));
             String rewriteLowerTerm = rewriteValue(inputToString(lowerTerm));
@@ -512,6 +525,7 @@ public Query existsQuery(QueryShardContext context) {
         public Query wildcardQuery(
             String value,
             @Nullable MultiTermQuery.RewriteMethod method,
+            @Nullable RewriteOverride rewriteOverride,
             boolean caseInsensitve,
             QueryShardContext context
         ) {

diff --git a/server/src/main/java/org/opensearch/index/mapper/IdFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/IdFieldMapper.java
@@ -41,6 +41,7 @@
 import org.apache.lucene.search.SortField;
 import org.apache.lucene.search.TermInSetQuery;
 import org.apache.lucene.util.BytesRef;
+import org.opensearch.common.Nullable;
 import org.opensearch.common.annotation.PublicApi;
 import org.opensearch.common.logging.DeprecationLogger;
 import org.opensearch.common.lucene.Lucene;
@@ -161,7 +162,7 @@ public Query existsQuery(QueryShardContext context) {
         }
 
         @Override
-        public Query termsQuery(List<?> values, QueryShardContext context) {
+        public Query termsQuery(List<?> values, @Nullable RewriteOverride rewriteOverride, QueryShardContext context) {
             failIfNotIndexed();
             BytesRef[] bytesRefs = new BytesRef[values.size()];
             for (int i = 0; i < bytesRefs.length; i++) {

diff --git a/server/src/main/java/org/opensearch/index/mapper/IpFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/IpFieldMapper.java
@@ -264,7 +264,7 @@ public Query termQuery(Object value, @Nullable QueryShardContext context) {
         }
 
         @Override
-        public Query termsQuery(List<?> values, QueryShardContext context) {
+        public Query termsQuery(List<?> values, @Nullable RewriteOverride rewriteOverride, QueryShardContext context) {
             failIfNotIndexedAndNoDocValues();
             InetAddress[] addresses = new InetAddress[values.size()];
             int i = 0;
@@ -289,7 +289,14 @@ public Query termsQuery(List<?> values, QueryShardContext context) {
         }
 
         @Override
-        public Query rangeQuery(Object lowerTerm, Object upperTerm, boolean includeLower, boolean includeUpper, QueryShardContext context) {
+        public Query rangeQuery(
+            Object lowerTerm,
+            Object upperTerm,
+            boolean includeLower,
+            boolean includeUpper,
+            @Nullable RewriteOverride rewriteOverride,
+            QueryShardContext context
+        ) {
             failIfNotIndexedAndNoDocValues();
             return rangeQuery(lowerTerm, upperTerm, includeLower, includeUpper, (lower, upper) -> {
                 Query query = InetAddressPoint.newRangeQuery(name(), lower, upper);

diff --git a/server/src/main/java/org/opensearch/index/mapper/KeywordFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/KeywordFieldMapper.java
@@ -564,7 +564,7 @@ public Query rangeQuery(
             Object upperTerm,
             boolean includeLower,
             boolean includeUpper,
-            RewriteOverride rewriteOverride,
+            @Nullable RewriteOverride rewriteOverride,
             QueryShardContext context
         ) {
             if (context.allowExpensiveQueries() == false) {
@@ -730,15 +730,16 @@ public Query wildcardQuery(
                         Query indexQuery = super.wildcardQuery(value, method, caseInsensitive, true, context);
                         Query dvQuery = super.wildcardQuery(value, MultiTermQuery.DOC_VALUES_REWRITE, caseInsensitive, true, context);
                         query = new IndexOrDocValuesQuery(indexQuery, dvQuery);
-                    }
-                    if (hasDocValues()) {
+                    } else if (hasDocValues()) {
                         Term term;
                         value = normalizeWildcardPattern(name(), value, getTextSearchInfo().getSearchAnalyzer());
                         term = new Term(name(), value);
                         if (caseInsensitive) {
                             return AutomatonQueries.caseInsensitiveWildcardQuery(term, method);
                         }
                         query = new WildcardQuery(term, Operations.DEFAULT_DETERMINIZE_WORK_LIMIT, MultiTermQuery.DOC_VALUES_REWRITE);
+                    } else {
+                        query = super.wildcardQuery(value, method, caseInsensitive, true, context);
                     }
                     break;
                 case INDEX_ONLY: