org.elasticsearch.search.sort.ScriptSortBuilder Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of elasticsearch Show documentation
Show all versions of elasticsearch Show documentation
Elasticsearch subproject :core
/*
* Licensed to Elasticsearch under one or more contributor
* license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright
* ownership. Elasticsearch licenses this file to you under
* the Apache License, Version 2.0 (the "License"); you may
* not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
package org.elasticsearch.search.sort;
import org.apache.lucene.index.BinaryDocValues;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.search.SortField;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.BytesRefBuilder;
import org.elasticsearch.common.ParseField;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.io.stream.Writeable;
import org.elasticsearch.common.xcontent.ConstructingObjectParser;
import org.elasticsearch.common.xcontent.ObjectParser.ValueType;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.index.fielddata.FieldData;
import org.elasticsearch.index.fielddata.IndexFieldData;
import org.elasticsearch.index.fielddata.IndexFieldData.XFieldComparatorSource.Nested;
import org.elasticsearch.index.fielddata.NumericDoubleValues;
import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
import org.elasticsearch.index.fielddata.fieldcomparator.BytesRefFieldComparatorSource;
import org.elasticsearch.index.fielddata.fieldcomparator.DoubleValuesComparatorSource;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryParseContext;
import org.elasticsearch.index.query.QueryShardContext;
import org.elasticsearch.index.query.QueryShardException;
import org.elasticsearch.script.LeafSearchScript;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptContext;
import org.elasticsearch.script.SearchScript;
import org.elasticsearch.search.DocValueFormat;
import org.elasticsearch.search.MultiValueMode;
import java.io.IOException;
import java.util.Locale;
import java.util.Objects;
import static org.elasticsearch.common.xcontent.ConstructingObjectParser.constructorArg;
/**
* Script sort builder allows to sort based on a custom script expression.
*/
public class ScriptSortBuilder extends SortBuilder {
public static final String NAME = "_script";
public static final ParseField TYPE_FIELD = new ParseField("type");
public static final ParseField SCRIPT_FIELD = new ParseField("script");
public static final ParseField SORTMODE_FIELD = new ParseField("mode");
private final Script script;
private final ScriptSortType type;
private SortMode sortMode;
private QueryBuilder nestedFilter;
private String nestedPath;
/**
* Constructs a script sort builder with the given script.
*
* @param script
* The script to use.
* @param type
* The type of the script, can be either {@link ScriptSortType#STRING} or
* {@link ScriptSortType#NUMBER}
*/
public ScriptSortBuilder(Script script, ScriptSortType type) {
Objects.requireNonNull(script, "script cannot be null");
Objects.requireNonNull(type, "type cannot be null");
this.script = script;
this.type = type;
}
ScriptSortBuilder(ScriptSortBuilder original) {
this.script = original.script;
this.type = original.type;
this.order = original.order;
this.sortMode = original.sortMode;
this.nestedFilter = original.nestedFilter;
this.nestedPath = original.nestedPath;
}
/**
* Read from a stream.
*/
public ScriptSortBuilder(StreamInput in) throws IOException {
script = new Script(in);
type = ScriptSortType.readFromStream(in);
order = SortOrder.readFromStream(in);
sortMode = in.readOptionalWriteable(SortMode::readFromStream);
nestedPath = in.readOptionalString();
nestedFilter = in.readOptionalNamedWriteable(QueryBuilder.class);
}
@Override
public void writeTo(StreamOutput out) throws IOException {
script.writeTo(out);
type.writeTo(out);
order.writeTo(out);
out.writeOptionalWriteable(sortMode);
out.writeOptionalString(nestedPath);
out.writeOptionalNamedWriteable(nestedFilter);
}
/**
* Get the script used in this sort.
*/
public Script script() {
return this.script;
}
/**
* Get the type used in this sort.
*/
public ScriptSortType type() {
return this.type;
}
/**
* Defines which distance to use for sorting in the case a document contains multiple values.
* For {@link ScriptSortType#STRING}, the set of possible values is restricted to {@link SortMode#MIN} and {@link SortMode#MAX}
*/
public ScriptSortBuilder sortMode(SortMode sortMode) {
Objects.requireNonNull(sortMode, "sort mode cannot be null.");
if (ScriptSortType.STRING.equals(type) && (sortMode == SortMode.SUM || sortMode == SortMode.AVG ||
sortMode == SortMode.MEDIAN)) {
throw new IllegalArgumentException("script sort of type [string] doesn't support mode [" + sortMode + "]");
}
this.sortMode = sortMode;
return this;
}
/**
* Get the sort mode.
*/
public SortMode sortMode() {
return this.sortMode;
}
/**
* Sets the nested filter that the nested objects should match with in order to be taken into account
* for sorting.
*/
public ScriptSortBuilder setNestedFilter(QueryBuilder nestedFilter) {
this.nestedFilter = nestedFilter;
return this;
}
/**
* Gets the nested filter.
*/
public QueryBuilder getNestedFilter() {
return this.nestedFilter;
}
/**
* Sets the nested path if sorting occurs on a field that is inside a nested object. For sorting by script this
* needs to be specified.
*/
public ScriptSortBuilder setNestedPath(String nestedPath) {
this.nestedPath = nestedPath;
return this;
}
/**
* Gets the nested path.
*/
public String getNestedPath() {
return this.nestedPath;
}
@Override
public XContentBuilder toXContent(XContentBuilder builder, Params builderParams) throws IOException {
builder.startObject();
builder.startObject(NAME);
builder.field(SCRIPT_FIELD.getPreferredName(), script);
builder.field(TYPE_FIELD.getPreferredName(), type);
builder.field(ORDER_FIELD.getPreferredName(), order);
if (sortMode != null) {
builder.field(SORTMODE_FIELD.getPreferredName(), sortMode);
}
if (nestedPath != null) {
builder.field(NESTED_PATH_FIELD.getPreferredName(), nestedPath);
}
if (nestedFilter != null) {
builder.field(NESTED_FILTER_FIELD.getPreferredName(), nestedFilter, builderParams);
}
builder.endObject();
builder.endObject();
return builder;
}
private static ConstructingObjectParser PARSER = new ConstructingObjectParser<>(NAME,
a -> new ScriptSortBuilder((Script) a[0], (ScriptSortType) a[1]));
static {
PARSER.declareField(constructorArg(), Script::parse, Script.SCRIPT_PARSE_FIELD, ValueType.OBJECT_OR_STRING);
PARSER.declareField(constructorArg(), p -> ScriptSortType.fromString(p.text()), TYPE_FIELD, ValueType.STRING);
PARSER.declareString((b, v) -> b.order(SortOrder.fromString(v)), ORDER_FIELD);
PARSER.declareString((b, v) -> b.sortMode(SortMode.fromString(v)), SORTMODE_FIELD);
PARSER.declareString(ScriptSortBuilder::setNestedPath , NESTED_PATH_FIELD);
PARSER.declareObject(ScriptSortBuilder::setNestedFilter, SortBuilder::parseNestedFilter, NESTED_FILTER_FIELD);
}
/**
* Creates a new {@link ScriptSortBuilder} from the query held by the {@link QueryParseContext} in
* {@link org.elasticsearch.common.xcontent.XContent} format.
*
* @param context the input parse context. The state on the parser contained in this context will be changed as a side effect of this
* method call
* @param elementName in some sort syntax variations the field name precedes the xContent object that specifies further parameters, e.g.
* in '{ "foo": { "order" : "asc"} }'. When parsing the inner object, the field name can be passed in via this argument
*/
public static ScriptSortBuilder fromXContent(QueryParseContext context, String elementName) throws IOException {
return PARSER.apply(context.parser(), context);
}
@Override
public SortFieldAndFormat build(QueryShardContext context) throws IOException {
final SearchScript searchScript = context.getSearchScript(script, ScriptContext.Standard.SEARCH);
MultiValueMode valueMode = null;
if (sortMode != null) {
valueMode = MultiValueMode.fromString(sortMode.toString());
}
boolean reverse = (order == SortOrder.DESC);
if (valueMode == null) {
valueMode = reverse ? MultiValueMode.MAX : MultiValueMode.MIN;
}
final Nested nested = resolveNested(context, nestedPath, nestedFilter);
final IndexFieldData.XFieldComparatorSource fieldComparatorSource;
switch (type) {
case STRING:
fieldComparatorSource = new BytesRefFieldComparatorSource(null, null, valueMode, nested) {
LeafSearchScript leafScript;
@Override
protected SortedBinaryDocValues getValues(LeafReaderContext context) throws IOException {
leafScript = searchScript.getLeafSearchScript(context);
final BinaryDocValues values = new BinaryDocValues() {
final BytesRefBuilder spare = new BytesRefBuilder();
@Override
public BytesRef get(int docID) {
leafScript.setDocument(docID);
spare.copyChars(leafScript.run().toString());
return spare.get();
}
};
return FieldData.singleton(values, null);
}
@Override
protected void setScorer(Scorer scorer) {
leafScript.setScorer(scorer);
}
};
break;
case NUMBER:
fieldComparatorSource = new DoubleValuesComparatorSource(null, Double.MAX_VALUE, valueMode, nested) {
LeafSearchScript leafScript;
@Override
protected SortedNumericDoubleValues getValues(LeafReaderContext context) throws IOException {
leafScript = searchScript.getLeafSearchScript(context);
final NumericDoubleValues values = new NumericDoubleValues() {
@Override
public double get(int docID) {
leafScript.setDocument(docID);
return leafScript.runAsDouble();
}
};
return FieldData.singleton(values, null);
}
@Override
protected void setScorer(Scorer scorer) {
leafScript.setScorer(scorer);
}
};
break;
default:
throw new QueryShardException(context, "custom script sort type [" + type + "] not supported");
}
return new SortFieldAndFormat(new SortField("_script", fieldComparatorSource, reverse), DocValueFormat.RAW);
}
@Override
public boolean equals(Object object) {
if (this == object) {
return true;
}
if (object == null || getClass() != object.getClass()) {
return false;
}
ScriptSortBuilder other = (ScriptSortBuilder) object;
return Objects.equals(script, other.script) &&
Objects.equals(type, other.type) &&
Objects.equals(order, other.order) &&
Objects.equals(sortMode, other.sortMode) &&
Objects.equals(nestedFilter, other.nestedFilter) &&
Objects.equals(nestedPath, other.nestedPath);
}
@Override
public int hashCode() {
return Objects.hash(script, type, order, sortMode, nestedFilter, nestedPath);
}
@Override
public String getWriteableName() {
return NAME;
}
public enum ScriptSortType implements Writeable {
/** script sort for a string value **/
STRING,
/** script sort for a numeric value **/
NUMBER;
@Override
public void writeTo(final StreamOutput out) throws IOException {
out.writeVInt(ordinal());
}
/**
* Read from a stream.
*/
static ScriptSortType readFromStream(final StreamInput in) throws IOException {
int ordinal = in.readVInt();
if (ordinal < 0 || ordinal >= values().length) {
throw new IOException("Unknown ScriptSortType ordinal [" + ordinal + "]");
}
return values()[ordinal];
}
public static ScriptSortType fromString(final String str) {
Objects.requireNonNull(str, "input string is null");
switch (str.toLowerCase(Locale.ROOT)) {
case ("string"):
return ScriptSortType.STRING;
case ("number"):
return ScriptSortType.NUMBER;
default:
throw new IllegalArgumentException("Unknown ScriptSortType [" + str + "]");
}
}
@Override
public String toString() {
return name().toLowerCase(Locale.ROOT);
}
}
}