Skip to content

Commit

Permalink
Adds various Query overrides to Keyword Field
Browse files Browse the repository at this point in the history
The keywordfield mapper provides access to various query types, e.g. the termsQuery, fuzzyQuery. These are inherited as is from the StringType. But we do not take into account the fact that keyword fields can have doc_values enabled. This PR adds the ability for various queries to first check if doc_values are enabled and if so out-source the work to lucene to decide if it's better to use index values or doc_values when running queries.

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>
  • Loading branch information
harshavamsi committed Oct 13, 2023
1 parent 1447d75 commit 292429b
Show file tree
Hide file tree
Showing 3 changed files with 382 additions and 17 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -38,11 +38,23 @@
import org.apache.lucene.document.FieldType;
import org.apache.lucene.document.SortedSetDocValuesField;
import org.apache.lucene.index.IndexOptions;
import org.apache.lucene.index.Term;
import org.apache.lucene.search.FuzzyQuery;
import org.apache.lucene.search.IndexOrDocValuesQuery;
import org.apache.lucene.search.MultiTermQuery;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.RegexpQuery;
import org.apache.lucene.search.TermInSetQuery;
import org.apache.lucene.search.TermRangeQuery;
import org.apache.lucene.search.WildcardQuery;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.automaton.Operations;
import org.opensearch.OpenSearchException;
import org.opensearch.common.Nullable;
import org.opensearch.common.lucene.BytesRefs;
import org.opensearch.common.lucene.Lucene;
import org.opensearch.common.lucene.search.AutomatonQueries;
import org.opensearch.common.unit.Fuzziness;
import org.opensearch.core.xcontent.XContentParser;
import org.opensearch.index.analysis.IndexAnalyzers;
import org.opensearch.index.analysis.NamedAnalyzer;
Expand All @@ -62,6 +74,8 @@
import java.util.Objects;
import java.util.function.Supplier;

import static org.opensearch.search.SearchService.ALLOW_EXPENSIVE_QUERIES;

/**
* A field mapper for keywords. This mapper accepts strings and indexes them as-is.
*
Expand Down Expand Up @@ -317,7 +331,7 @@ public IndexFieldData.Builder fielddataBuilder(String fullyQualifiedIndexName, S
@Override
public ValueFetcher valueFetcher(QueryShardContext context, SearchLookup searchLookup, String format) {
if (format != null) {
throw new IllegalArgumentException("Field [" + name() + "] of type [" + typeName() + "] doesn't support formats.");
throw new IllegalArgumentException("Field [" + name() + "] of type [" + typeName() + "] doesn't " + "support formats.");
}

return new SourceValueFetcher(name(), context, nullValue) {
Expand Down Expand Up @@ -372,17 +386,220 @@ protected BytesRef indexedValueForSearch(Object value) {
return getTextSearchInfo().getSearchAnalyzer().normalize(name(), value.toString());
}

@Override
public Query termsQuery(List<?> values, QueryShardContext context) {
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 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);
}
// has index enabled, we're going to return the query as is
return super.termsQuery(values, context);
}

@Override
public Query prefixQuery(
String value,
@Nullable MultiTermQuery.RewriteMethod method,
boolean caseInsensitive,
QueryShardContext context
) {
if (context.allowExpensiveQueries() == false) {
throw new OpenSearchException(
"[prefix] queries cannot be executed when '"
+ ALLOW_EXPENSIVE_QUERIES.getKey()
+ "' is set to false. For optimised prefix queries on text "
+ "fields please enable [index_prefixes]."
);
}
failIfNotIndexedAndNoDocValues();
if (isSearchable() && hasDocValues()) {
Query indexQuery = super.prefixQuery(value, method, caseInsensitive, context);
Query dvQuery = super.prefixQuery(value, MultiTermQuery.DOC_VALUES_REWRITE, caseInsensitive, context);
return new IndexOrDocValuesQuery(indexQuery, dvQuery);
}
if (hasDocValues()) {
return super.prefixQuery(value, MultiTermQuery.DOC_VALUES_REWRITE, caseInsensitive, context);
}
return super.prefixQuery(value, method, caseInsensitive, context);
}

@Override
public Query regexpQuery(
String value,
int syntaxFlags,
int matchFlags,
int maxDeterminizedStates,
@Nullable MultiTermQuery.RewriteMethod method,
QueryShardContext context
) {
if (context.allowExpensiveQueries() == false) {
throw new OpenSearchException(
"[regexp] queries cannot be executed when '" + ALLOW_EXPENSIVE_QUERIES.getKey() + "' is set to " + "false."
);
}
failIfNotIndexedAndNoDocValues();
if (isSearchable() && hasDocValues()) {
Query indexQuery = super.regexpQuery(value, syntaxFlags, matchFlags, maxDeterminizedStates, method, context);
Query dvQuery = super.regexpQuery(
value,
syntaxFlags,
matchFlags,
maxDeterminizedStates,
MultiTermQuery.DOC_VALUES_REWRITE,
context
);
return new IndexOrDocValuesQuery(indexQuery, dvQuery);
}
if (hasDocValues()) {
return new RegexpQuery(
new Term(name(), indexedValueForSearch(value)),
syntaxFlags,
matchFlags,
RegexpQuery.DEFAULT_PROVIDER,
maxDeterminizedStates,
MultiTermQuery.DOC_VALUES_REWRITE
);
}
return super.regexpQuery(value, syntaxFlags, matchFlags, maxDeterminizedStates, method, context);
}

@Override
public Query rangeQuery(Object lowerTerm, Object upperTerm, boolean includeLower, boolean includeUpper, QueryShardContext context) {
if (context.allowExpensiveQueries() == false) {
throw new OpenSearchException(
"[range] queries on [text] or [keyword] fields cannot be executed when '"
+ ALLOW_EXPENSIVE_QUERIES.getKey()
+ "' is set to false."
);
}
failIfNotIndexedAndNoDocValues();
if (isSearchable() && hasDocValues()) {
Query indexQuery = new TermRangeQuery(
name(),
lowerTerm == null ? null : indexedValueForSearch(lowerTerm),
upperTerm == null ? null : indexedValueForSearch(upperTerm),
includeLower,
includeUpper
);
Query dvQuery = new TermRangeQuery(
name(),
lowerTerm == null ? null : indexedValueForSearch(lowerTerm),
upperTerm == null ? null : indexedValueForSearch(upperTerm),
includeLower,
includeUpper,
MultiTermQuery.DOC_VALUES_REWRITE
);
return new IndexOrDocValuesQuery(indexQuery, dvQuery);
}
if (hasDocValues()) {
return new TermRangeQuery(
name(),
lowerTerm == null ? null : indexedValueForSearch(lowerTerm),
upperTerm == null ? null : indexedValueForSearch(upperTerm),
includeLower,
includeUpper,
MultiTermQuery.DOC_VALUES_REWRITE
);
}
return new TermRangeQuery(
name(),
lowerTerm == null ? null : indexedValueForSearch(lowerTerm),
upperTerm == null ? null : indexedValueForSearch(upperTerm),
includeLower,
includeUpper
);
}

@Override
public Query fuzzyQuery(
Object value,
Fuzziness fuzziness,
int prefixLength,
int maxExpansions,
boolean transpositions,
@Nullable MultiTermQuery.RewriteMethod method,
QueryShardContext context
) {
failIfNotIndexedAndNoDocValues();
if (context.allowExpensiveQueries() == false) {
throw new OpenSearchException(
"[fuzzy] queries cannot be executed when '" + ALLOW_EXPENSIVE_QUERIES.getKey() + "' is set to " + "false."
);
}
if (isSearchable() && hasDocValues()) {
Query indexQuery = super.fuzzyQuery(value, fuzziness, prefixLength, maxExpansions, transpositions, context);
Query dvQuery = super.fuzzyQuery(
value,
fuzziness,
prefixLength,
maxExpansions,
transpositions,
MultiTermQuery.DOC_VALUES_REWRITE,
context
);
return new IndexOrDocValuesQuery(indexQuery, dvQuery);
}
if (hasDocValues()) {
return new FuzzyQuery(
new Term(name(), indexedValueForSearch(value)),
fuzziness.asDistance(BytesRefs.toString(value)),
prefixLength,
maxExpansions,
transpositions,
MultiTermQuery.DOC_VALUES_REWRITE
);
}
return super.fuzzyQuery(value, fuzziness, prefixLength, maxExpansions, transpositions, context);
}

@Override
public Query wildcardQuery(
String value,
@Nullable MultiTermQuery.RewriteMethod method,
boolean caseInsensitve,
boolean caseInsensitive,
QueryShardContext context
) {
// keyword field types are always normalized, so ignore case sensitivity and force normalize the wildcard
if (context.allowExpensiveQueries() == false) {
throw new OpenSearchException(
"[wildcard] queries cannot be executed when '" + ALLOW_EXPENSIVE_QUERIES.getKey() + "' is set to " + "false."
);
}
failIfNotIndexedAndNoDocValues();
// keyword field types are always normalized, so ignore case sensitivity and force normalize the
// wildcard
// query text
return super.wildcardQuery(value, method, caseInsensitve, true, context);
if (isSearchable() && hasDocValues()) {
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);
}
if (hasDocValues()) {
Term term;
value = normalizeWildcardPattern(name(), value, getTextSearchInfo().getSearchAnalyzer());
term = new Term(name(), value);
if (caseInsensitive) {
return AutomatonQueries.caseInsensitiveWildcardQuery(term, method);
}
return new WildcardQuery(term, Operations.DEFAULT_DETERMINIZE_WORK_LIMIT, MultiTermQuery.DOC_VALUES_REWRITE);
}
return super.wildcardQuery(value, method, caseInsensitive, true, context);
}

}

private final boolean indexed;
Expand Down Expand Up @@ -422,8 +639,10 @@ protected KeywordFieldMapper(
this.indexAnalyzers = builder.indexAnalyzers;
}

/** Values that have more chars than the return value of this method will
* be skipped at parsing time. */
/**
* Values that have more chars than the return value of this method will
* be skipped at parsing time.
*/
public int ignoreAbove() {
return ignoreAbove;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -269,6 +269,21 @@ public Query fuzzyQuery(
);
}

// Fuzzy Query with re-write method
public Query fuzzyQuery(
Object value,
Fuzziness fuzziness,
int prefixLength,
int maxExpansions,
boolean transpositions,
@Nullable MultiTermQuery.RewriteMethod method,
QueryShardContext context
) {
throw new IllegalArgumentException(
"Can only use fuzzy 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 @@ -433,6 +448,15 @@ protected final void failIfNotIndexed() {
}
}

protected final void failIfNotIndexedAndNoDocValues() {
// we fail if a field is both not indexed and does not have doc_values enabled
if (isIndexed == false && hasDocValues() == false) {
throw new IllegalArgumentException(
"Cannot search on field [" + name() + "] since it is both not indexed," + " and does not have doc_values enabled."
);
}
}

public boolean eagerGlobalOrdinals() {
return eagerGlobalOrdinals;
}
Expand Down
Loading

0 comments on commit 292429b

Please sign in to comment.