From 6617c87e1bd4517215b143c84137eb77b6597c66 Mon Sep 17 00:00:00 2001 From: Harsha Vamsi Kalluri Date: Mon, 29 Jul 2024 15:19:29 -0700 Subject: [PATCH] Adding new rewrite override parameter to terms query Signed-off-by: Harsha Vamsi Kalluri --- CHANGELOG.md | 1 + .../index/mapper/KeywordFieldMapper.java | 60 +++++++++++++------ .../index/mapper/MappedFieldType.java | 8 +++ .../index/mapper/RewriteOverride.java | 29 +++++++++ .../index/query/TermsQueryBuilder.java | 45 ++++++++++++-- .../index/query/support/QueryParsers.java | 26 ++++++++ 6 files changed, 146 insertions(+), 23 deletions(-) create mode 100644 server/src/main/java/org/opensearch/index/mapper/RewriteOverride.java diff --git a/CHANGELOG.md b/CHANGELOG.md index 7cc918b2ac089..a4fee25885e39 100644 --- 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 index 2116ac522b705..7df80670bea18 100644 --- 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 index 66d4654e543a2..291f1f26054fa 100644 --- 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 new file mode 100644 index 0000000000000..966d8feebab17 --- /dev/null +++ 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 index ac0ca3919ea38..b868a2946efcc 100644 --- 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 { private final TermsLookup termsLookup; private final Supplier> 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> 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 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 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> actionListener) { @@ -491,7 +527,7 @@ private void fetch(TermsLookup termsLookup, Client client, ActionListener