Adding new rewrite override parameter to terms query

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>
harshavamsi · Jul 29, 2024 · 6617c87 · 6617c87
1 parent f5b0eba
commit 6617c87
Show file tree

Hide file tree

Showing 6 changed files with 146 additions and 23 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -14,6 +14,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Bump `actions/github-script` from 6 to 7 ([#14997](https://github.com/opensearch-project/OpenSearch/pull/14997))
 
 ### Changed
+- Updated `termsQuery` in `KeywordField` to set custom rewrite_override before executing query
 
 ### Deprecated
 

diff --git a/server/src/main/java/org/opensearch/index/mapper/KeywordFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/KeywordFieldMapper.java
@@ -388,28 +388,50 @@ protected BytesRef indexedValueForSearch(Object value) {
         }
 
         @Override
-        public Query termsQuery(List<?> values, QueryShardContext context) {
+        public Query termsQuery(List<?> values, QueryShardContext context, RewriteOverride rewriteOverride) {
             failIfNotIndexedAndNoDocValues();
-            // has index and doc_values enabled
-            if (isSearchable() && hasDocValues()) {
-                BytesRef[] bytesRefs = new BytesRef[values.size()];
-                for (int i = 0; i < bytesRefs.length; i++) {
-                    bytesRefs[i] = indexedValueForSearch(values.get(i));
-                }
-                Query indexQuery = new TermInSetQuery(name(), bytesRefs);
-                Query dvQuery = new TermInSetQuery(MultiTermQuery.DOC_VALUES_REWRITE, name(), bytesRefs);
-                return new IndexOrDocValuesQuery(indexQuery, dvQuery);
+            if (rewriteOverride == null) {
+                rewriteOverride = RewriteOverride.DEFAULT;
             }
-            // if we only have doc_values enabled, we construct a new query with doc_values re-written
-            if (hasDocValues()) {
-                BytesRef[] bytesRefs = new BytesRef[values.size()];
-                for (int i = 0; i < bytesRefs.length; i++) {
-                    bytesRefs[i] = indexedValueForSearch(values.get(i));
-                }
-                return new TermInSetQuery(MultiTermQuery.DOC_VALUES_REWRITE, name(), bytesRefs);
+            Query query = null;
+            switch (rewriteOverride) {
+                case DEFAULT:
+                    // has index and doc_values enabled
+                    if (isSearchable() && hasDocValues()) {
+                        BytesRef[] bytesRefs = new BytesRef[values.size()];
+                        for (int i = 0; i < bytesRefs.length; i++) {
+                            bytesRefs[i] = indexedValueForSearch(values.get(i));
+                        }
+                        Query indexQuery = new TermInSetQuery(name(), bytesRefs);
+                        Query dvQuery = new TermInSetQuery(MultiTermQuery.DOC_VALUES_REWRITE, name(), bytesRefs);
+                        query = new IndexOrDocValuesQuery(indexQuery, dvQuery);
+                    }
+                    // if we only have doc_values enabled, we construct a new query with doc_values re-written
+                    else if (hasDocValues()) {
+                        BytesRef[] bytesRefs = new BytesRef[values.size()];
+                        for (int i = 0; i < bytesRefs.length; i++) {
+                            bytesRefs[i] = indexedValueForSearch(values.get(i));
+                        }
+                        query = new TermInSetQuery(MultiTermQuery.DOC_VALUES_REWRITE, name(), bytesRefs);
+                    } else {
+                        query = super.termsQuery(values, context);
+                    }
+                    break;
+                case INDEX_ONLY:
+                    failIfNotIndexed();
+                    query = super.termsQuery(values, context);
+                    break;
+                case DOC_VALUES_ONLY:
+                    failIfNoDocValues();
+                    BytesRef[] bytesRefs = new BytesRef[values.size()];
+                    for (int i = 0; i < bytesRefs.length; i++) {
+                        bytesRefs[i] = indexedValueForSearch(values.get(i));
+                    }
+                    query = new TermInSetQuery(MultiTermQuery.DOC_VALUES_REWRITE, name(), bytesRefs);
+                    break;
             }
-            // has index enabled, we're going to return the query as is
-            return super.termsQuery(values, context);
+
+            return query;
         }
 
         @Override

diff --git a/server/src/main/java/org/opensearch/index/mapper/MappedFieldType.java b/server/src/main/java/org/opensearch/index/mapper/MappedFieldType.java
@@ -230,6 +230,14 @@ public Query termQueryCaseInsensitive(Object value, @Nullable QueryShardContext
         );
     }
 
+    public Query termsQuery(List<?> values, @Nullable QueryShardContext context, @Nullable RewriteOverride rewriteOverride) {
+        BooleanQuery.Builder builder = new BooleanQuery.Builder();
+        for (Object value : values) {
+            builder.add(termQuery(value, context), Occur.SHOULD);
+        }
+        return new ConstantScoreQuery(builder.build());
+    }
+
     /** Build a constant-scoring query that matches all values. The default implementation uses a
      * {@link ConstantScoreQuery} around a {@link BooleanQuery} whose {@link Occur#SHOULD} clauses
      * are generated with {@link #termQuery}. */

diff --git a/server/src/main/java/org/opensearch/index/mapper/RewriteOverride.java b/server/src/main/java/org/opensearch/index/mapper/RewriteOverride.java
@@ -0,0 +1,29 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.index.mapper;
+
+import org.opensearch.common.annotation.PublicApi;
+
+/*
+* A custom rewrite override for a query. Default executes the query as is and other values determine which structure to use at run-time.
+*
+* @opensearch.internal
+* */
+@PublicApi(since = "2.17.0")
+public enum RewriteOverride {
+
+    // don't override the rewrite, use default
+    DEFAULT,
+
+    // use index structure
+    INDEX_ONLY,
+
+    // use doc_values structure
+    DOC_VALUES_ONLY
+}
diff --git a/server/src/main/java/org/opensearch/index/query/TermsQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/TermsQueryBuilder.java
@@ -41,7 +41,9 @@
 import org.opensearch.client.Client;
 import org.opensearch.common.SetOnce;
 import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.common.xcontent.LoggingDeprecationHandler;
 import org.opensearch.common.xcontent.support.XContentMapValues;
+import org.opensearch.core.ParseField;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.ParsingException;
 import org.opensearch.core.common.Strings;
@@ -53,6 +55,8 @@
 import org.opensearch.index.IndexSettings;
 import org.opensearch.index.mapper.ConstantFieldType;
 import org.opensearch.index.mapper.MappedFieldType;
+import org.opensearch.index.mapper.RewriteOverride;
+import org.opensearch.index.query.support.QueryParsers;
 import org.opensearch.indices.TermsLookup;
 
 import java.io.IOException;
@@ -82,6 +86,10 @@ public class TermsQueryBuilder extends AbstractQueryBuilder<TermsQueryBuilder> {
     private final TermsLookup termsLookup;
     private final Supplier<List<?>> supplier;
 
+    private static final ParseField REWRITE_OVERRIDE = new ParseField("rewrite_override");
+
+    private String rewrite_override;
+
     public TermsQueryBuilder(String fieldName, TermsLookup termsLookup) {
         this(fieldName, null, termsLookup);
     }
@@ -201,6 +209,7 @@ public TermsQueryBuilder(StreamInput in) throws IOException {
         super(in);
         fieldName = in.readString();
         termsLookup = in.readOptionalWriteable(TermsLookup::new);
+        rewrite_override = in.readOptionalString();
         values = (List<?>) in.readGenericValue();
         this.supplier = null;
     }
@@ -212,6 +221,7 @@ protected void doWriteTo(StreamOutput out) throws IOException {
         }
         out.writeString(fieldName);
         out.writeOptionalWriteable(termsLookup);
+        out.writeOptionalString(rewrite_override);
         out.writeGenericValue(values);
     }
 
@@ -227,6 +237,15 @@ public TermsLookup termsLookup() {
         return this.termsLookup;
     }
 
+    public TermsQueryBuilder rewrite_override(String rewrite_override) {
+        this.rewrite_override = rewrite_override;
+        return this;
+    }
+
+    public String rewrite_override() {
+        return this.rewrite_override;
+    }
+
     private static final Set<Class<? extends Number>> INTEGER_TYPES = new HashSet<>(
         Arrays.asList(Byte.class, Short.class, Integer.class, Long.class)
     );
@@ -352,6 +371,9 @@ public Object get(int index) {
     @Override
     protected void doXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject(NAME);
+        if (rewrite_override != null) {
+            builder.field(REWRITE_OVERRIDE.getPreferredName(), rewrite_override);
+        }
         if (this.termsLookup != null) {
             builder.startObject(fieldName);
             termsLookup.toXContent(builder, params);
@@ -368,6 +390,8 @@ public static TermsQueryBuilder fromXContent(XContentParser parser) throws IOExc
         List<Object> values = null;
         TermsLookup termsLookup = null;
 
+        String rewrite_override = null;
+
         String queryName = null;
         float boost = AbstractQueryBuilder.DEFAULT_BOOST;
 
@@ -401,6 +425,13 @@ public static TermsQueryBuilder fromXContent(XContentParser parser) throws IOExc
                 }
                 fieldName = currentFieldName;
                 termsLookup = TermsLookup.parseTermsLookup(parser);
+                while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
+                    if (token == XContentParser.Token.FIELD_NAME) {
+                        currentFieldName = parser.currentName();
+                    } else if (REWRITE_OVERRIDE.match(currentFieldName, parser.getDeprecationHandler())) {
+                        rewrite_override = parser.textOrNull();
+                    }
+                }
             } else if (token.isValue()) {
                 if (AbstractQueryBuilder.BOOST_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
                     boost = parser.floatValue();
@@ -430,7 +461,7 @@ public static TermsQueryBuilder fromXContent(XContentParser parser) throws IOExc
             );
         }
 
-        return new TermsQueryBuilder(fieldName, values, termsLookup).boost(boost).queryName(queryName);
+        return new TermsQueryBuilder(fieldName, values, termsLookup).boost(boost).queryName(queryName).rewrite_override(rewrite_override);
     }
 
     static List<Object> parseValues(XContentParser parser) throws IOException {
@@ -473,7 +504,12 @@ protected Query doToQuery(QueryShardContext context) throws IOException {
         if (fieldType == null) {
             throw new IllegalStateException("Rewrite first");
         }
-        return fieldType.termsQuery(values, context);
+        RewriteOverride rewriteOverride = QueryParsers.parseRewriteOverride(
+            rewrite_override,
+            RewriteOverride.DEFAULT,
+            LoggingDeprecationHandler.INSTANCE
+        );
+        return fieldType.termsQuery(values, context, rewriteOverride);
     }
 
     private void fetch(TermsLookup termsLookup, Client client, ActionListener<List<Object>> actionListener) {
@@ -491,14 +527,15 @@ private void fetch(TermsLookup termsLookup, Client client, ActionListener<List<O
 
     @Override
     protected int doHashCode() {
-        return Objects.hash(fieldName, values, termsLookup, supplier);
+        return Objects.hash(fieldName, values, termsLookup, rewrite_override, supplier);
     }
 
     @Override
     protected boolean doEquals(TermsQueryBuilder other) {
         return Objects.equals(fieldName, other.fieldName)
             && Objects.equals(values, other.values)
             && Objects.equals(termsLookup, other.termsLookup)
+            && Objects.equals(rewrite_override, other.rewrite_override)
             && Objects.equals(supplier, other.supplier);
     }
 
@@ -528,7 +565,7 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) {
                 // This logic is correct for all field types, but by only applying it to constant
                 // fields we also have the guarantee that it doesn't perform I/O, which is important
                 // since rewrites might happen on a network thread.
-                Query query = fieldType.termsQuery(values, context);
+                Query query = fieldType.termsQuery(values, context, null);
                 if (query instanceof MatchAllDocsQuery) {
                     return new MatchAllQueryBuilder();
                 } else if (query instanceof MatchNoDocsQuery) {

diff --git a/server/src/main/java/org/opensearch/index/query/support/QueryParsers.java b/server/src/main/java/org/opensearch/index/query/support/QueryParsers.java
@@ -36,6 +36,7 @@
 import org.opensearch.common.Nullable;
 import org.opensearch.core.ParseField;
 import org.opensearch.core.xcontent.DeprecationHandler;
+import org.opensearch.index.mapper.RewriteOverride;
 
 /**
  * Utility class for Query Parsers
@@ -51,6 +52,10 @@ public final class QueryParsers {
     public static final ParseField TOP_TERMS_BOOST = new ParseField("top_terms_boost_");
     public static final ParseField TOP_TERMS_BLENDED_FREQS = new ParseField("top_terms_blended_freqs_");
 
+    public static final ParseField DEFAULT = new ParseField("default");
+    public static final ParseField INDEX_ONLY = new ParseField("index_only");
+    public static final ParseField DOC_VALUES_ONLY = new ParseField("doc_values_only");
+
     private QueryParsers() {
 
     }
@@ -66,6 +71,27 @@ public static MultiTermQuery.RewriteMethod parseRewriteMethod(@Nullable String r
         return parseRewriteMethod(rewriteMethod, MultiTermQuery.CONSTANT_SCORE_REWRITE, deprecationHandler);
     }
 
+    public static RewriteOverride parseRewriteOverride(
+        @Nullable String rewrite_override,
+        @Nullable RewriteOverride defaultRewriteOverride,
+        DeprecationHandler deprecationHandler
+    ) {
+        if (rewrite_override == null) {
+            return defaultRewriteOverride;
+        }
+        if (DEFAULT.match(rewrite_override, deprecationHandler)) {
+            return RewriteOverride.DEFAULT;
+        }
+        if (INDEX_ONLY.match(rewrite_override, deprecationHandler)) {
+            return RewriteOverride.INDEX_ONLY;
+        }
+        if (DOC_VALUES_ONLY.match(rewrite_override, deprecationHandler)) {
+            return RewriteOverride.DOC_VALUES_ONLY;
+        }
+        throw new IllegalArgumentException("Failed to parse rewrite_override [" + rewrite_override + "]");
+
+    }
+
     public static MultiTermQuery.RewriteMethod parseRewriteMethod(
         @Nullable String rewriteMethod,
         @Nullable MultiTermQuery.RewriteMethod defaultRewriteMethod,