Skip to content

Commit

Permalink
Add rewrite_override to more queries
Browse files Browse the repository at this point in the history
Signed-off-by: Harsha Vamsi Kalluri <[email protected]>
  • Loading branch information
harshavamsi committed Aug 5, 2024
1 parent 07e3caa commit 67425be
Show file tree
Hide file tree
Showing 8 changed files with 403 additions and 122 deletions.
340 changes: 236 additions & 104 deletions server/src/main/java/org/opensearch/index/mapper/KeywordFieldMapper.java

Large diffs are not rendered by default.

Original file line number Diff line number Diff line change
Expand Up @@ -230,7 +230,7 @@ public Query termQueryCaseInsensitive(Object value, @Nullable QueryShardContext
);
}

public Query termsQuery(List<?> values, @Nullable QueryShardContext context, @Nullable RewriteOverride rewriteOverride) {
public Query termsQuery(List<?> values, @Nullable RewriteOverride rewriteOverride, @Nullable QueryShardContext context) {
BooleanQuery.Builder builder = new BooleanQuery.Builder();
for (Object value : values) {
builder.add(termQuery(value, context), Occur.SHOULD);
Expand Down Expand Up @@ -294,6 +294,34 @@ public Query fuzzyQuery(
);
}

public Query fuzzyQuery(
Object value,
Fuzziness fuzziness,
int prefixLength,
int maxExpansions,
boolean transpositions,
@Nullable MultiTermQuery.RewriteMethod method,
@Nullable RewriteOverride rewriteOverride,
QueryShardContext context
) {
throw new IllegalArgumentException(
"Can only use fuzzy queries on keyword and text fields - not on [" + name + "] which is of type [" + typeName() + "]"
);
}

public Query prefixQuery(
String value,
@Nullable MultiTermQuery.RewriteMethod method,
@Nullable RewriteOverride rewriteOverride,
boolean caseInsensitve,
QueryShardContext context
) {
throw new QueryShardException(
context,
"Can only use prefix queries on keyword and text fields - not on [" + name + "] which is of type [" + typeName() + "]"
);
}

// Case sensitive form of prefix query
public final Query prefixQuery(String value, @Nullable MultiTermQuery.RewriteMethod method, QueryShardContext context) {
return prefixQuery(value, method, false, context);
Expand Down Expand Up @@ -328,6 +356,18 @@ public Query wildcardQuery(
"Can only use wildcard queries on keyword and text fields - not on [" + name + "] which is of type [" + typeName() + "]"
);
}
public Query wildcardQuery(
String value,
@Nullable MultiTermQuery.RewriteMethod method,
@Nullable RewriteOverride rewriteOverride,
boolean caseInsensitve,
QueryShardContext context
) {
throw new QueryShardException(
context,
"Can only use wildcard queries on keyword and text fields - not on [" + name + "] which is of type [" + typeName() + "]"
);
}

/** always normalizes the wildcard pattern to lowercase */
public Query normalizedWildcardQuery(String value, @Nullable MultiTermQuery.RewriteMethod method, QueryShardContext context) {
Expand All @@ -351,6 +391,21 @@ public Query regexpQuery(
);
}

public Query regexpQuery(
String value,
int syntaxFlags,
int matchFlags,
int maxDeterminizedStates,
@Nullable MultiTermQuery.RewriteMethod method,
@Nullable RewriteOverride rewriteOverride,
QueryShardContext context
) {
throw new QueryShardException(
context,
"Can only use regexp queries on keyword and text fields - not on [" + name + "] which is of type [" + typeName() + "]"
);
}

public Query existsQuery(QueryShardContext context) {
if (hasDocValues()) {
return new DocValuesFieldExistsQuery(name());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -86,4 +86,8 @@ protected Query rangeQuery(Object lowerTerm, Object upperTerm, boolean includeLo
throw new IllegalArgumentException("Field [" + name() + "] of type [" + typeName() + "] does not support range queries");
}

protected Query rangeQuery(Object lowerTerm, Object upperTerm, boolean includeLower, boolean includeUpper,RewriteOverride rewriteOverride, QueryShardContext context) {
throw new IllegalArgumentException("Field [" + name() + "] of type [" + typeName() + "] does not support range queries");
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
import org.apache.lucene.search.MatchNoDocsQuery;
import org.apache.lucene.search.MultiTermQuery;
import org.apache.lucene.search.Query;
import org.opensearch.Version;
import org.opensearch.common.xcontent.LoggingDeprecationHandler;
import org.opensearch.core.ParseField;
import org.opensearch.core.common.ParsingException;
Expand All @@ -62,6 +63,10 @@ public class PrefixQueryBuilder extends AbstractQueryBuilder<PrefixQueryBuilder>
private static final ParseField PREFIX_FIELD = new ParseField("value");
private static final ParseField REWRITE_FIELD = new ParseField("rewrite");

private static final ParseField REWRITE_OVERRIDE = new ParseField("rewrite_override");

private String rewrite_override;

private final String fieldName;

private final String value;
Expand Down Expand Up @@ -98,6 +103,9 @@ public PrefixQueryBuilder(StreamInput in) throws IOException {
value = in.readString();
rewrite = in.readOptionalString();
caseInsensitive = in.readBoolean();
if (in.getVersion().after(Version.V_2_16_0)) {
rewrite_override = in.readOptionalString();
}
}

@Override
Expand All @@ -106,6 +114,9 @@ protected void doWriteTo(StreamOutput out) throws IOException {
out.writeString(value);
out.writeOptionalString(rewrite);
out.writeBoolean(caseInsensitive);
if (out.getVersion().after(Version.V_2_16_0)) {
out.writeOptionalString(rewrite_override);
}
}

@Override
Expand All @@ -131,6 +142,11 @@ public PrefixQueryBuilder rewrite(String rewrite) {
return this;
}

public PrefixQueryBuilder rewrite_override(String rewrite_override) {
this.rewrite_override = rewrite_override;
return this;
}

public String rewrite() {
return this.rewrite;
}
Expand All @@ -146,6 +162,9 @@ public void doXContent(XContentBuilder builder, Params params) throws IOExceptio
if (caseInsensitive != DEFAULT_CASE_INSENSITIVITY) {
builder.field(CASE_INSENSITIVE_FIELD.getPreferredName(), caseInsensitive);
}
if (rewrite_override != null) {
builder.field(REWRITE_OVERRIDE.getPreferredName(), rewrite_override);
}
printBoostAndQueryName(builder);
builder.endObject();
builder.endObject();
Expand All @@ -155,6 +174,8 @@ public static PrefixQueryBuilder fromXContent(XContentParser parser) throws IOEx
String fieldName = null;
String value = null;
String rewrite = null;
String rewrite_override = null;


String queryName = null;
float boost = AbstractQueryBuilder.DEFAULT_BOOST;
Expand All @@ -181,7 +202,9 @@ public static PrefixQueryBuilder fromXContent(XContentParser parser) throws IOEx
rewrite = parser.textOrNull();
} else if (CASE_INSENSITIVE_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
caseInsensitive = parser.booleanValue();
} else {
} else if (REWRITE_OVERRIDE.match(currentFieldName, parser.getDeprecationHandler())) {
rewrite_override = parser.textOrNull();
}else {
throw new ParsingException(
parser.getTokenLocation(),
"[prefix] query does not support [" + currentFieldName + "]"
Expand All @@ -196,7 +219,7 @@ public static PrefixQueryBuilder fromXContent(XContentParser parser) throws IOEx
}
}

return new PrefixQueryBuilder(fieldName, value).rewrite(rewrite).boost(boost).queryName(queryName).caseInsensitive(caseInsensitive);
return new PrefixQueryBuilder(fieldName, value).rewrite(rewrite).boost(boost).queryName(queryName).caseInsensitive(caseInsensitive).rewrite_override(rewrite_override);
}

@Override
Expand Down Expand Up @@ -242,14 +265,15 @@ protected Query doToQuery(QueryShardContext context) throws IOException {

@Override
protected final int doHashCode() {
return Objects.hash(fieldName, value, rewrite, caseInsensitive);
return Objects.hash(fieldName, value, rewrite, caseInsensitive, rewrite_override);
}

@Override
protected boolean doEquals(PrefixQueryBuilder other) {
return Objects.equals(fieldName, other.fieldName)
&& Objects.equals(value, other.value)
&& Objects.equals(rewrite, other.rewrite)
&& Objects.equals(rewrite_override, other.rewrite_override)
&& Objects.equals(caseInsensitive, other.caseInsensitive);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
import org.apache.lucene.search.MatchNoDocsQuery;
import org.apache.lucene.search.Query;
import org.apache.lucene.util.BytesRef;
import org.opensearch.Version;
import org.opensearch.common.geo.ShapeRelation;
import org.opensearch.common.time.DateFormatter;
import org.opensearch.common.time.DateMathParser;
Expand Down Expand Up @@ -75,6 +76,8 @@ public class RangeQueryBuilder extends AbstractQueryBuilder<RangeQueryBuilder> i
private static final ParseField TIME_ZONE_FIELD = new ParseField("time_zone");
private static final ParseField FORMAT_FIELD = new ParseField("format");
private static final ParseField RELATION_FIELD = new ParseField("relation");
private static final ParseField REWRITE_OVERRIDE = new ParseField("rewrite_override");


private final String fieldName;
private Object from;
Expand All @@ -85,6 +88,8 @@ public class RangeQueryBuilder extends AbstractQueryBuilder<RangeQueryBuilder> i
private String format;
private ShapeRelation relation;

private String rewrite_override;

/**
* A Query that matches documents within an range of terms.
*
Expand Down Expand Up @@ -116,6 +121,9 @@ public RangeQueryBuilder(StreamInput in) throws IOException {
throw new IllegalArgumentException("[range] query does not support relation [" + relationString + "]");
}
}
if (in.getVersion().after(Version.V_2_16_0)) {
rewrite_override = in.readOptionalString();
}
}

private boolean isRelationAllowed(ShapeRelation relation) {
Expand All @@ -136,6 +144,9 @@ protected void doWriteTo(StreamOutput out) throws IOException {
relationString = this.relation.getRelationName();
}
out.writeOptionalString(relationString);
if (out.getVersion().after(Version.V_2_16_0)) {
out.writeOptionalString(rewrite_override);
}
}

/**
Expand Down Expand Up @@ -271,6 +282,11 @@ public RangeQueryBuilder timeZone(String timeZone) {
return this;
}

public RangeQueryBuilder rewrite_override(String rewrite_override) {
this.rewrite_override = rewrite_override;
return this;
}

/**
* In case of date field, gets the from/to fields timezone adjustment
*/
Expand Down Expand Up @@ -344,6 +360,9 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep
if (relation != null) {
builder.field(RELATION_FIELD.getPreferredName(), relation.getRelationName());
}
if (rewrite_override != null) {
builder.field(REWRITE_OVERRIDE.getPreferredName(), rewrite_override);
}
printBoostAndQueryName(builder);
builder.endObject();
builder.endObject();
Expand All @@ -360,6 +379,8 @@ public static RangeQueryBuilder fromXContent(XContentParser parser) throws IOExc
String queryName = null;
String format = null;
String relation = null;
String rewrite_override = null;


String currentFieldName = null;
XContentParser.Token token;
Expand Down Expand Up @@ -403,7 +424,9 @@ public static RangeQueryBuilder fromXContent(XContentParser parser) throws IOExc
relation = parser.text();
} else if (AbstractQueryBuilder.NAME_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
queryName = parser.text();
} else {
} else if (REWRITE_OVERRIDE.match(currentFieldName, parser.getDeprecationHandler())) {
rewrite_override = parser.textOrNull();
}else {
throw new ParsingException(
parser.getTokenLocation(),
"[range] query does not support [" + currentFieldName + "]"
Expand Down Expand Up @@ -432,6 +455,7 @@ public static RangeQueryBuilder fromXContent(XContentParser parser) throws IOExc
if (relation != null) {
rangeQuery.relation(relation);
}
rangeQuery.rewrite_override(rewrite_override);
return rangeQuery;
}

Expand Down Expand Up @@ -529,7 +553,7 @@ protected Query doToQuery(QueryShardContext context) throws IOException {

@Override
protected int doHashCode() {
return Objects.hash(fieldName, from, to, timeZone, includeLower, includeUpper, format);
return Objects.hash(fieldName, from, to, timeZone, includeLower, includeUpper, format, rewrite_override);
}

@Override
Expand All @@ -540,6 +564,7 @@ protected boolean doEquals(RangeQueryBuilder other) {
&& Objects.equals(timeZone, other.timeZone)
&& Objects.equals(includeLower, other.includeLower)
&& Objects.equals(includeUpper, other.includeUpper)
&& Objects.equals(format, other.format);
&& Objects.equals(format, other.format)
&& Objects.equals(rewrite_override, other.rewrite_override);
}
}
Loading

0 comments on commit 67425be

Please sign in to comment.