All Downloads are FREE. Search and download functionalities are using the official Maven repository.

org.elasticsearch.search.sort.ScriptSortBuilder Maven / Gradle / Ivy

There is a newer version: 8.14.1
Show newest version
/*
 * 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); } } }




© 2015 - 2024 Weber Informatics LLC | Privacy Policy