org.opensearch.index.query.ScriptQueryBuilder Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of opensearch Show documentation
Show all versions of opensearch Show documentation
OpenSearch subproject :server
/*
* 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.
*/
/*
* 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.
*/
/*
* Modifications Copyright OpenSearch Contributors. See
* GitHub history for details.
*/
package org.opensearch.index.query;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.ConstantScoreScorer;
import org.apache.lucene.search.ConstantScoreWeight;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.QueryVisitor;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.search.TwoPhaseIterator;
import org.apache.lucene.search.Weight;
import org.opensearch.OpenSearchException;
import org.opensearch.common.Nullable;
import org.opensearch.common.lucene.search.function.Functions;
import org.opensearch.core.common.ParsingException;
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.core.common.io.stream.StreamOutput;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.core.xcontent.XContentParser;
import org.opensearch.script.FilterScript;
import org.opensearch.script.Script;
import java.io.IOException;
import java.util.Objects;
import static org.opensearch.search.SearchService.ALLOW_EXPENSIVE_QUERIES;
/**
* Query builder for script queries
*
* @opensearch.internal
*/
public class ScriptQueryBuilder extends AbstractQueryBuilder {
public static final String NAME = "script";
private final Script script;
public ScriptQueryBuilder(Script script) {
if (script == null) {
throw new IllegalArgumentException("script cannot be null");
}
this.script = script;
}
/**
* Read from a stream.
*/
public ScriptQueryBuilder(StreamInput in) throws IOException {
super(in);
script = new Script(in);
}
@Override
protected void doWriteTo(StreamOutput out) throws IOException {
script.writeTo(out);
}
public Script script() {
return this.script;
}
@Override
public String getWriteableName() {
return NAME;
}
@Override
protected void doXContent(XContentBuilder builder, Params builderParams) throws IOException {
builder.startObject(NAME);
builder.field(Script.SCRIPT_PARSE_FIELD.getPreferredName(), script);
printBoostAndQueryName(builder);
builder.endObject();
}
public static ScriptQueryBuilder fromXContent(XContentParser parser) throws IOException {
// also, when caching, since its isCacheable is false, will result in loading all bit set...
Script script = null;
float boost = AbstractQueryBuilder.DEFAULT_BOOST;
String queryName = null;
XContentParser.Token token;
String currentFieldName = null;
while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
if (token == XContentParser.Token.FIELD_NAME) {
currentFieldName = parser.currentName();
} else if (token == XContentParser.Token.START_OBJECT) {
if (Script.SCRIPT_PARSE_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
script = Script.parse(parser);
} else {
throw new ParsingException(parser.getTokenLocation(), "[script] query does not support [" + currentFieldName + "]");
}
} else if (token.isValue()) {
if (AbstractQueryBuilder.NAME_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
queryName = parser.text();
} else if (AbstractQueryBuilder.BOOST_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
boost = parser.floatValue();
} else if (Script.SCRIPT_PARSE_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
script = Script.parse(parser);
} else {
throw new ParsingException(parser.getTokenLocation(), "[script] query does not support [" + currentFieldName + "]");
}
} else {
if (token != XContentParser.Token.START_ARRAY) {
throw new AssertionError("Impossible token received: " + token.name());
}
throw new ParsingException(
parser.getTokenLocation(),
"[script] query does not support an array of scripts. Use a bool query with a clause per script instead."
);
}
}
if (script == null) {
throw new ParsingException(parser.getTokenLocation(), "script must be provided with a [script] filter");
}
return new ScriptQueryBuilder(script).boost(boost).queryName(queryName);
}
@Override
protected Query doToQuery(QueryShardContext context) throws IOException {
if (context.allowExpensiveQueries() == false) {
throw new OpenSearchException(
"[script] queries cannot be executed when '" + ALLOW_EXPENSIVE_QUERIES.getKey() + "' is set to false."
);
}
FilterScript.Factory factory = context.compile(script, FilterScript.CONTEXT);
FilterScript.LeafFactory filterScript = factory.newFactory(script.getParams(), context.lookup());
return new ScriptQuery(script, filterScript, queryName);
}
/**
* Internal script query
*
* @opensearch.internal
*/
static class ScriptQuery extends Query {
final Script script;
final FilterScript.LeafFactory filterScript;
final String queryName;
ScriptQuery(Script script, FilterScript.LeafFactory filterScript, @Nullable String queryName) {
this.script = script;
this.filterScript = filterScript;
this.queryName = queryName;
}
@Override
public String toString(String field) {
StringBuilder buffer = new StringBuilder();
buffer.append("ScriptQuery(");
buffer.append(script);
buffer.append(Functions.nameOrEmptyArg(queryName));
buffer.append(")");
return buffer.toString();
}
@Override
public void visit(QueryVisitor visitor) {
visitor.visitLeaf(this);
}
@Override
public boolean equals(Object obj) {
if (sameClassAs(obj) == false) return false;
ScriptQuery other = (ScriptQuery) obj;
return Objects.equals(script, other.script);
}
@Override
public int hashCode() {
int h = classHash();
h = 31 * h + script.hashCode();
return h;
}
@Override
public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
return new ConstantScoreWeight(this, boost) {
@Override
public Scorer scorer(LeafReaderContext context) throws IOException {
DocIdSetIterator approximation = DocIdSetIterator.all(context.reader().maxDoc());
final FilterScript leafScript = filterScript.newInstance(context);
TwoPhaseIterator twoPhase = new TwoPhaseIterator(approximation) {
@Override
public boolean matches() throws IOException {
leafScript.setDocument(approximation.docID());
return leafScript.execute();
}
@Override
public float matchCost() {
// TODO: how can we compute this?
return 1000f;
}
};
return new ConstantScoreScorer(this, score(), scoreMode, twoPhase);
}
@Override
public boolean isCacheable(LeafReaderContext ctx) {
// TODO: Change this to true when we can assume that scripts are pure functions
// ie. the return value is always the same given the same conditions and may not
// depend on the current timestamp, other documents, etc.
return false;
}
};
}
}
@Override
protected int doHashCode() {
return Objects.hash(script);
}
@Override
protected boolean doEquals(ScriptQueryBuilder other) {
return Objects.equals(script, other.script);
}
}