From 39e59d6edffa8202c8e2956a28e12425d398ad03 Mon Sep 17 00:00:00 2001 From: Luca Cavanna Date: Thu, 10 Sep 2020 17:50:26 +0200 Subject: [PATCH] Share more query execution code for runtime fields (#62229) For runtime fields we have written quite some lucene queries that work against runtime values that are the result of the execution of the different script contexts that runtime fields support. The all (but one) share the same main logic: use a two phase iterator, iterate over all documents, and decide whether the current doc matches or not based on what the script returns. I went ahead and shared this bit of code in the base class for all queries on top of runtime fields. --- .../BooleanScriptFieldScript.java | 4 +- .../runtimefields/DateScriptFieldScript.java | 4 +- .../DoubleScriptFieldScript.java | 4 +- .../runtimefields/IpScriptFieldScript.java | 4 +- .../runtimefields/LongScriptFieldScript.java | 4 +- .../StringScriptFieldScript.java | 4 +- .../AbstractBooleanScriptFieldQuery.java | 53 +++--------------- .../query/AbstractDoubleScriptFieldQuery.java | 53 +++--------------- .../query/AbstractIpScriptFieldQuery.java | 52 +++--------------- .../query/AbstractLongScriptFieldQuery.java | 55 ++++--------------- .../query/AbstractScriptFieldQuery.java | 54 +++++++++++++++++- .../query/AbstractStringScriptFieldQuery.java | 50 +++-------------- .../LongScriptFieldDistanceFeatureQuery.java | 33 ++++++----- .../query/LongScriptFieldExistsQuery.java | 5 +- .../query/LongScriptFieldRangeQuery.java | 5 +- .../query/LongScriptFieldTermQuery.java | 5 +- .../query/LongScriptFieldTermsQuery.java | 6 +- .../AbstractLongScriptFieldQueryTestCase.java | 5 +- ...gScriptFieldDistanceFeatureQueryTests.java | 25 +++++---- 19 files changed, 147 insertions(+), 278 deletions(-) diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/BooleanScriptFieldScript.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/BooleanScriptFieldScript.java index 7be6e8ba31a..131d1182df0 100644 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/BooleanScriptFieldScript.java +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/BooleanScriptFieldScript.java @@ -14,7 +14,6 @@ import org.elasticsearch.script.ScriptContext; import org.elasticsearch.script.ScriptFactory; import org.elasticsearch.search.lookup.SearchLookup; -import java.io.IOException; import java.util.Collections; import java.util.List; import java.util.Map; @@ -28,6 +27,7 @@ public abstract class BooleanScriptFieldScript extends AbstractScriptFieldScript ); } + @SuppressWarnings("unused") public static final String[] PARAMETERS = {}; public interface Factory extends ScriptFactory { @@ -35,7 +35,7 @@ public abstract class BooleanScriptFieldScript extends AbstractScriptFieldScript } public interface LeafFactory { - BooleanScriptFieldScript newInstance(LeafReaderContext ctx) throws IOException; + BooleanScriptFieldScript newInstance(LeafReaderContext ctx); } private int trues; diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/DateScriptFieldScript.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/DateScriptFieldScript.java index a08eb63896c..ac5485fafac 100644 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/DateScriptFieldScript.java +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/DateScriptFieldScript.java @@ -14,7 +14,6 @@ import org.elasticsearch.script.ScriptContext; import org.elasticsearch.script.ScriptFactory; import org.elasticsearch.search.lookup.SearchLookup; -import java.io.IOException; import java.time.temporal.ChronoField; import java.time.temporal.TemporalAccessor; import java.util.Collections; @@ -28,6 +27,7 @@ public abstract class DateScriptFieldScript extends AbstractLongScriptFieldScrip return Collections.singletonList(WhitelistLoader.loadFromResourceFiles(RuntimeFieldsPainlessExtension.class, "date_whitelist.txt")); } + @SuppressWarnings("unused") public static final String[] PARAMETERS = {}; public interface Factory extends ScriptFactory { @@ -35,7 +35,7 @@ public abstract class DateScriptFieldScript extends AbstractLongScriptFieldScrip } public interface LeafFactory { - DateScriptFieldScript newInstance(LeafReaderContext ctx) throws IOException; + DateScriptFieldScript newInstance(LeafReaderContext ctx); } private final DateFormatter formatter; diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/DoubleScriptFieldScript.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/DoubleScriptFieldScript.java index 0ade80bdb65..38dc4cf09d8 100644 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/DoubleScriptFieldScript.java +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/DoubleScriptFieldScript.java @@ -14,7 +14,6 @@ import org.elasticsearch.script.ScriptContext; import org.elasticsearch.script.ScriptFactory; import org.elasticsearch.search.lookup.SearchLookup; -import java.io.IOException; import java.util.Collections; import java.util.List; import java.util.Map; @@ -28,6 +27,7 @@ public abstract class DoubleScriptFieldScript extends AbstractScriptFieldScript ); } + @SuppressWarnings("unused") public static final String[] PARAMETERS = {}; public interface Factory extends ScriptFactory { @@ -35,7 +35,7 @@ public abstract class DoubleScriptFieldScript extends AbstractScriptFieldScript } public interface LeafFactory { - DoubleScriptFieldScript newInstance(LeafReaderContext ctx) throws IOException; + DoubleScriptFieldScript newInstance(LeafReaderContext ctx); } private double[] values = new double[1]; diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/IpScriptFieldScript.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/IpScriptFieldScript.java index 49f260f2073..529b28d9200 100644 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/IpScriptFieldScript.java +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/IpScriptFieldScript.java @@ -18,7 +18,6 @@ import org.elasticsearch.script.ScriptContext; import org.elasticsearch.script.ScriptFactory; import org.elasticsearch.search.lookup.SearchLookup; -import java.io.IOException; import java.net.Inet4Address; import java.net.Inet6Address; import java.net.InetAddress; @@ -47,6 +46,7 @@ public abstract class IpScriptFieldScript extends AbstractScriptFieldScript { return Collections.singletonList(WhitelistLoader.loadFromResourceFiles(RuntimeFieldsPainlessExtension.class, "ip_whitelist.txt")); } + @SuppressWarnings("unused") public static final String[] PARAMETERS = {}; public interface Factory extends ScriptFactory { @@ -54,7 +54,7 @@ public abstract class IpScriptFieldScript extends AbstractScriptFieldScript { } public interface LeafFactory { - IpScriptFieldScript newInstance(LeafReaderContext ctx) throws IOException; + IpScriptFieldScript newInstance(LeafReaderContext ctx); } private BytesRef[] values = new BytesRef[1]; diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/LongScriptFieldScript.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/LongScriptFieldScript.java index 3835127eccc..d546624c3a8 100644 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/LongScriptFieldScript.java +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/LongScriptFieldScript.java @@ -13,7 +13,6 @@ import org.elasticsearch.script.ScriptContext; import org.elasticsearch.script.ScriptFactory; import org.elasticsearch.search.lookup.SearchLookup; -import java.io.IOException; import java.util.Collections; import java.util.List; import java.util.Map; @@ -25,6 +24,7 @@ public abstract class LongScriptFieldScript extends AbstractLongScriptFieldScrip return Collections.singletonList(WhitelistLoader.loadFromResourceFiles(RuntimeFieldsPainlessExtension.class, "long_whitelist.txt")); } + @SuppressWarnings("unused") public static final String[] PARAMETERS = {}; public interface Factory extends ScriptFactory { @@ -32,7 +32,7 @@ public abstract class LongScriptFieldScript extends AbstractLongScriptFieldScrip } public interface LeafFactory { - LongScriptFieldScript newInstance(LeafReaderContext ctx) throws IOException; + LongScriptFieldScript newInstance(LeafReaderContext ctx); } public LongScriptFieldScript(String fieldName, Map params, SearchLookup searchLookup, LeafReaderContext ctx) { diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/StringScriptFieldScript.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/StringScriptFieldScript.java index e600e6e640f..861b0daec6b 100644 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/StringScriptFieldScript.java +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/StringScriptFieldScript.java @@ -13,7 +13,6 @@ import org.elasticsearch.script.ScriptContext; import org.elasticsearch.script.ScriptFactory; import org.elasticsearch.search.lookup.SearchLookup; -import java.io.IOException; import java.util.ArrayList; import java.util.Collections; import java.util.List; @@ -34,6 +33,7 @@ public abstract class StringScriptFieldScript extends AbstractScriptFieldScript ); } + @SuppressWarnings("unused") public static final String[] PARAMETERS = {}; public interface Factory extends ScriptFactory { @@ -41,7 +41,7 @@ public abstract class StringScriptFieldScript extends AbstractScriptFieldScript } public interface LeafFactory { - StringScriptFieldScript newInstance(LeafReaderContext ctx) throws IOException; + StringScriptFieldScript newInstance(LeafReaderContext ctx); } private final List results = new ArrayList<>(); diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractBooleanScriptFieldQuery.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractBooleanScriptFieldQuery.java index b4af5ee2f6f..060e733903c 100644 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractBooleanScriptFieldQuery.java +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractBooleanScriptFieldQuery.java @@ -6,32 +6,24 @@ package org.elasticsearch.xpack.runtimefields.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.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.elasticsearch.script.Script; import org.elasticsearch.xpack.runtimefields.BooleanScriptFieldScript; import org.elasticsearch.xpack.runtimefields.DoubleScriptFieldScript; -import java.io.IOException; -import java.util.Objects; - /** * Abstract base class for building queries based on {@link DoubleScriptFieldScript}. */ -abstract class AbstractBooleanScriptFieldQuery extends AbstractScriptFieldQuery { - private final BooleanScriptFieldScript.LeafFactory leafFactory; +abstract class AbstractBooleanScriptFieldQuery extends AbstractScriptFieldQuery { AbstractBooleanScriptFieldQuery(Script script, BooleanScriptFieldScript.LeafFactory leafFactory, String fieldName) { - super(script, fieldName); - this.leafFactory = Objects.requireNonNull(leafFactory); + super(script, fieldName, leafFactory::newInstance); + } + + @Override + protected boolean matches(BooleanScriptFieldScript scriptContext, int docId) { + scriptContext.runForDoc(docId); + return matches(scriptContext.trues(), scriptContext.falses()); } /** @@ -41,35 +33,6 @@ abstract class AbstractBooleanScriptFieldQuery extends AbstractScriptFieldQuery */ protected abstract boolean matches(int trues, int falses); - @Override - public final Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { - return new ConstantScoreWeight(this, boost) { - @Override - public boolean isCacheable(LeafReaderContext ctx) { - return false; // scripts aren't really cacheable at this point - } - - @Override - public Scorer scorer(LeafReaderContext ctx) throws IOException { - BooleanScriptFieldScript script = leafFactory.newInstance(ctx); - DocIdSetIterator approximation = DocIdSetIterator.all(ctx.reader().maxDoc()); - TwoPhaseIterator twoPhase = new TwoPhaseIterator(approximation) { - @Override - public boolean matches() throws IOException { - script.runForDoc(approximation().docID()); - return AbstractBooleanScriptFieldQuery.this.matches(script.trues(), script.falses()); - } - - @Override - public float matchCost() { - return MATCH_COST; - } - }; - return new ConstantScoreScorer(this, score(), scoreMode, twoPhase); - } - }; - } - @Override public final void visit(QueryVisitor visitor) { // No subclasses contain any Terms because those have to be strings. diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractDoubleScriptFieldQuery.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractDoubleScriptFieldQuery.java index 5c4841085e2..dacc337a8fa 100644 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractDoubleScriptFieldQuery.java +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractDoubleScriptFieldQuery.java @@ -6,31 +6,23 @@ package org.elasticsearch.xpack.runtimefields.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.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.elasticsearch.script.Script; import org.elasticsearch.xpack.runtimefields.DoubleScriptFieldScript; -import java.io.IOException; -import java.util.Objects; - /** * Abstract base class for building queries based on {@link DoubleScriptFieldScript}. */ -abstract class AbstractDoubleScriptFieldQuery extends AbstractScriptFieldQuery { - private final DoubleScriptFieldScript.LeafFactory leafFactory; +abstract class AbstractDoubleScriptFieldQuery extends AbstractScriptFieldQuery { AbstractDoubleScriptFieldQuery(Script script, DoubleScriptFieldScript.LeafFactory leafFactory, String fieldName) { - super(script, fieldName); - this.leafFactory = Objects.requireNonNull(leafFactory); + super(script, fieldName, leafFactory::newInstance); + } + + @Override + protected boolean matches(DoubleScriptFieldScript scriptContext, int docId) { + scriptContext.runForDoc(docId); + return matches(scriptContext.values(), scriptContext.count()); } /** @@ -38,35 +30,6 @@ abstract class AbstractDoubleScriptFieldQuery extends AbstractScriptFieldQuery { */ protected abstract boolean matches(double[] values, int count); - @Override - public final Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { - return new ConstantScoreWeight(this, boost) { - @Override - public boolean isCacheable(LeafReaderContext ctx) { - return false; // scripts aren't really cacheable at this point - } - - @Override - public Scorer scorer(LeafReaderContext ctx) throws IOException { - DoubleScriptFieldScript script = leafFactory.newInstance(ctx); - DocIdSetIterator approximation = DocIdSetIterator.all(ctx.reader().maxDoc()); - TwoPhaseIterator twoPhase = new TwoPhaseIterator(approximation) { - @Override - public boolean matches() throws IOException { - script.runForDoc(approximation().docID()); - return AbstractDoubleScriptFieldQuery.this.matches(script.values(), script.count()); - } - - @Override - public float matchCost() { - return MATCH_COST; - } - }; - return new ConstantScoreScorer(this, score(), scoreMode, twoPhase); - } - }; - } - @Override public final void visit(QueryVisitor visitor) { // No subclasses contain any Terms because those have to be strings. diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractIpScriptFieldQuery.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractIpScriptFieldQuery.java index c68982bd944..15b18cf7e9c 100644 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractIpScriptFieldQuery.java +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractIpScriptFieldQuery.java @@ -7,15 +7,6 @@ package org.elasticsearch.xpack.runtimefields.query; import org.apache.lucene.document.InetAddressPoint; -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.ScoreMode; -import org.apache.lucene.search.Scorer; -import org.apache.lucene.search.TwoPhaseIterator; -import org.apache.lucene.search.Weight; import org.apache.lucene.util.BytesRef; import org.elasticsearch.common.bytes.BytesArray; import org.elasticsearch.common.bytes.BytesReference; @@ -23,19 +14,21 @@ import org.elasticsearch.script.Script; import org.elasticsearch.xpack.runtimefields.IpScriptFieldScript; import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript; -import java.io.IOException; import java.net.InetAddress; -import java.util.Objects; /** * Abstract base class for building queries based on {@link StringScriptFieldScript}. */ -abstract class AbstractIpScriptFieldQuery extends AbstractScriptFieldQuery { - private final IpScriptFieldScript.LeafFactory leafFactory; +abstract class AbstractIpScriptFieldQuery extends AbstractScriptFieldQuery { AbstractIpScriptFieldQuery(Script script, IpScriptFieldScript.LeafFactory leafFactory, String fieldName) { - super(script, fieldName); - this.leafFactory = Objects.requireNonNull(leafFactory); + super(script, fieldName, leafFactory::newInstance); + } + + @Override + protected boolean matches(IpScriptFieldScript scriptContext, int docId) { + scriptContext.runForDoc(docId); + return matches(scriptContext.values(), scriptContext.count()); } /** @@ -43,35 +36,6 @@ abstract class AbstractIpScriptFieldQuery extends AbstractScriptFieldQuery { */ protected abstract boolean matches(BytesRef[] values, int conut); - @Override - public final Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { - return new ConstantScoreWeight(this, boost) { - @Override - public boolean isCacheable(LeafReaderContext ctx) { - return false; // scripts aren't really cacheable at this point - } - - @Override - public Scorer scorer(LeafReaderContext ctx) throws IOException { - IpScriptFieldScript script = leafFactory.newInstance(ctx); - DocIdSetIterator approximation = DocIdSetIterator.all(ctx.reader().maxDoc()); - TwoPhaseIterator twoPhase = new TwoPhaseIterator(approximation) { - @Override - public boolean matches() throws IOException { - script.runForDoc(approximation().docID()); - return AbstractIpScriptFieldQuery.this.matches(script.values(), script.count()); - } - - @Override - public float matchCost() { - return MATCH_COST; - } - }; - return new ConstantScoreScorer(this, score(), scoreMode, twoPhase); - } - }; - } - protected static InetAddress decode(BytesRef ref) { return InetAddressPoint.decode(BytesReference.toBytes(new BytesArray(ref))); } diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractLongScriptFieldQuery.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractLongScriptFieldQuery.java index 8fb3f99ab05..9b5c9c06be8 100644 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractLongScriptFieldQuery.java +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractLongScriptFieldQuery.java @@ -7,35 +7,29 @@ package org.elasticsearch.xpack.runtimefields.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.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.elasticsearch.common.CheckedFunction; import org.elasticsearch.script.Script; import org.elasticsearch.xpack.runtimefields.AbstractLongScriptFieldScript; -import java.io.IOException; -import java.util.Objects; +import java.util.function.Function; /** * Abstract base class for building queries based on {@link AbstractLongScriptFieldScript}. */ -abstract class AbstractLongScriptFieldQuery extends AbstractScriptFieldQuery { - private final CheckedFunction leafFactory; +abstract class AbstractLongScriptFieldQuery extends AbstractScriptFieldQuery { AbstractLongScriptFieldQuery( Script script, - CheckedFunction leafFactory, + Function scriptContextFunction, String fieldName ) { - super(script, fieldName); - this.leafFactory = Objects.requireNonNull(leafFactory); + super(script, fieldName, scriptContextFunction); + } + + @Override + protected boolean matches(AbstractLongScriptFieldScript scriptContext, int docId) { + scriptContext.runForDoc(docId); + return AbstractLongScriptFieldQuery.this.matches(scriptContext.values(), scriptContext.count()); } /** @@ -43,35 +37,6 @@ abstract class AbstractLongScriptFieldQuery extends AbstractScriptFieldQuery { */ protected abstract boolean matches(long[] values, int count); - @Override - public final Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { - return new ConstantScoreWeight(this, boost) { - @Override - public boolean isCacheable(LeafReaderContext ctx) { - return false; // scripts aren't really cacheable at this point - } - - @Override - public Scorer scorer(LeafReaderContext ctx) throws IOException { - AbstractLongScriptFieldScript script = leafFactory.apply(ctx); - DocIdSetIterator approximation = DocIdSetIterator.all(ctx.reader().maxDoc()); - TwoPhaseIterator twoPhase = new TwoPhaseIterator(approximation) { - @Override - public boolean matches() throws IOException { - script.runForDoc(approximation().docID()); - return AbstractLongScriptFieldQuery.this.matches(script.values(), script.count()); - } - - @Override - public float matchCost() { - return MATCH_COST; - } - }; - return new ConstantScoreScorer(this, score(), scoreMode, twoPhase); - } - }; - } - @Override public final void visit(QueryVisitor visitor) { // No subclasses contain any Terms because those have to be strings. diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractScriptFieldQuery.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractScriptFieldQuery.java index a69dab5512c..3a994da235d 100644 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractScriptFieldQuery.java +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractScriptFieldQuery.java @@ -6,15 +6,27 @@ package org.elasticsearch.xpack.runtimefields.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.ScoreMode; +import org.apache.lucene.search.Scorer; +import org.apache.lucene.search.TwoPhaseIterator; +import org.apache.lucene.search.Weight; import org.elasticsearch.script.Script; +import org.elasticsearch.xpack.runtimefields.AbstractScriptFieldScript; +import java.io.IOException; import java.util.Objects; +import java.util.function.Function; /** * Abstract base class for building queries based on script fields. */ -abstract class AbstractScriptFieldQuery extends Query { +abstract class AbstractScriptFieldQuery extends Query { /** * We don't have the infrastructure to estimate the match cost of a script * so we just use a big number. @@ -23,10 +35,16 @@ abstract class AbstractScriptFieldQuery extends Query { private final Script script; private final String fieldName; + private final Function scriptContextFunction; - AbstractScriptFieldQuery(Script script, String fieldName) { + AbstractScriptFieldQuery(Script script, String fieldName, Function scriptContextFunction) { this.script = Objects.requireNonNull(script); this.fieldName = Objects.requireNonNull(fieldName); + this.scriptContextFunction = scriptContextFunction; + } + + final Function scriptContextFunction() { + return scriptContextFunction; } final Script script() { @@ -37,6 +55,36 @@ abstract class AbstractScriptFieldQuery extends Query { return fieldName; } + @Override + public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { + return new ConstantScoreWeight(this, boost) { + @Override + public boolean isCacheable(LeafReaderContext ctx) { + return false; // scripts aren't really cacheable at this point + } + + @Override + public Scorer scorer(LeafReaderContext ctx) { + S scriptContext = scriptContextFunction.apply(ctx); + DocIdSetIterator approximation = DocIdSetIterator.all(ctx.reader().maxDoc()); + TwoPhaseIterator twoPhase = new TwoPhaseIterator(approximation) { + @Override + public boolean matches() { + return AbstractScriptFieldQuery.this.matches(scriptContext, approximation.docID()); + } + + @Override + public float matchCost() { + return MATCH_COST; + } + }; + return new ConstantScoreScorer(this, score(), scoreMode, twoPhase); + } + }; + } + + protected abstract boolean matches(S scriptContext, int docId); + @Override public int hashCode() { return Objects.hash(getClass(), script, fieldName); @@ -47,7 +95,7 @@ abstract class AbstractScriptFieldQuery extends Query { if (obj == null || getClass() != obj.getClass()) { return false; } - AbstractScriptFieldQuery other = (AbstractScriptFieldQuery) obj; + AbstractScriptFieldQuery other = (AbstractScriptFieldQuery) obj; return script.equals(other.script) && fieldName.equals(other.fieldName); } } diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractStringScriptFieldQuery.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractStringScriptFieldQuery.java index b17bfc5ac13..6a14019c70a 100644 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractStringScriptFieldQuery.java +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractStringScriptFieldQuery.java @@ -6,63 +6,27 @@ package org.elasticsearch.xpack.runtimefields.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.ScoreMode; -import org.apache.lucene.search.Scorer; -import org.apache.lucene.search.TwoPhaseIterator; -import org.apache.lucene.search.Weight; import org.elasticsearch.script.Script; import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript; -import java.io.IOException; import java.util.List; -import java.util.Objects; /** * Abstract base class for building queries based on {@link StringScriptFieldScript}. */ -abstract class AbstractStringScriptFieldQuery extends AbstractScriptFieldQuery { - private final StringScriptFieldScript.LeafFactory leafFactory; +abstract class AbstractStringScriptFieldQuery extends AbstractScriptFieldQuery { AbstractStringScriptFieldQuery(Script script, StringScriptFieldScript.LeafFactory leafFactory, String fieldName) { - super(script, fieldName); - this.leafFactory = Objects.requireNonNull(leafFactory); + super(script, fieldName, leafFactory::newInstance); + } + + @Override + protected final boolean matches(StringScriptFieldScript scriptContext, int docId) { + return matches(scriptContext.resultsForDoc(docId)); } /** * Does the value match this query? */ protected abstract boolean matches(List values); - - @Override - public final Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { - return new ConstantScoreWeight(this, boost) { - @Override - public boolean isCacheable(LeafReaderContext ctx) { - return false; // scripts aren't really cacheable at this point - } - - @Override - public Scorer scorer(LeafReaderContext ctx) throws IOException { - StringScriptFieldScript script = leafFactory.newInstance(ctx); - DocIdSetIterator approximation = DocIdSetIterator.all(ctx.reader().maxDoc()); - TwoPhaseIterator twoPhase = new TwoPhaseIterator(approximation) { - @Override - public boolean matches() throws IOException { - return AbstractStringScriptFieldQuery.this.matches(script.resultsForDoc(approximation().docID())); - } - - @Override - public float matchCost() { - return MATCH_COST; - } - }; - return new ConstantScoreScorer(this, score(), scoreMode, twoPhase); - } - }; - } } diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldDistanceFeatureQuery.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldDistanceFeatureQuery.java index 29bc3a4f9a3..e5b947eff47 100644 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldDistanceFeatureQuery.java +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldDistanceFeatureQuery.java @@ -16,35 +16,39 @@ 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.elasticsearch.common.CheckedFunction; import org.elasticsearch.script.Script; import org.elasticsearch.xpack.runtimefields.AbstractLongScriptFieldScript; import java.io.IOException; import java.util.Objects; import java.util.Set; +import java.util.function.Function; -public final class LongScriptFieldDistanceFeatureQuery extends AbstractScriptFieldQuery { - private final CheckedFunction leafFactory; +public final class LongScriptFieldDistanceFeatureQuery extends AbstractScriptFieldQuery { private final long origin; private final long pivot; private final float boost; public LongScriptFieldDistanceFeatureQuery( Script script, - CheckedFunction leafFactory, + Function leafFactory, String fieldName, long origin, long pivot, float boost ) { - super(script, fieldName); - this.leafFactory = leafFactory; + super(script, fieldName, leafFactory); this.origin = origin; this.pivot = pivot; this.boost = boost; } + @Override + protected boolean matches(AbstractLongScriptFieldScript scriptContext, int docId) { + scriptContext.runForDoc(docId); + return scriptContext.count() > 0; + } + @Override public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { return new Weight(this) { @@ -57,13 +61,13 @@ public final class LongScriptFieldDistanceFeatureQuery extends AbstractScriptFie public void extractTerms(Set terms) {} @Override - public Scorer scorer(LeafReaderContext context) throws IOException { - return new DistanceScorer(this, leafFactory.apply(context), context.reader().maxDoc(), boost); + public Scorer scorer(LeafReaderContext context) { + return new DistanceScorer(this, scriptContextFunction().apply(context), context.reader().maxDoc(), boost); } @Override - public Explanation explain(LeafReaderContext context, int doc) throws IOException { - AbstractLongScriptFieldScript script = leafFactory.apply(context); + public Explanation explain(LeafReaderContext context, int doc) { + AbstractLongScriptFieldScript script = scriptContextFunction().apply(context); script.runForDoc(doc); long value = valueWithMinAbsoluteDistance(script); float weight = LongScriptFieldDistanceFeatureQuery.this.boost * boost; @@ -91,9 +95,8 @@ public final class LongScriptFieldDistanceFeatureQuery extends AbstractScriptFie this.script = script; twoPhase = new TwoPhaseIterator(DocIdSetIterator.all(maxDoc)) { @Override - public boolean matches() throws IOException { - script.runForDoc(approximation().docID()); - return script.count() > 0; + public boolean matches() { + return LongScriptFieldDistanceFeatureQuery.this.matches(script, approximation.docID()); } @Override @@ -121,12 +124,12 @@ public final class LongScriptFieldDistanceFeatureQuery extends AbstractScriptFie } @Override - public float getMaxScore(int upTo) throws IOException { + public float getMaxScore(int upTo) { return weight; } @Override - public float score() throws IOException { + public float score() { if (script.count() == 0) { return 0; } diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldExistsQuery.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldExistsQuery.java index 3a6e28bdd79..7289fc68043 100644 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldExistsQuery.java +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldExistsQuery.java @@ -7,16 +7,15 @@ package org.elasticsearch.xpack.runtimefields.query; import org.apache.lucene.index.LeafReaderContext; -import org.elasticsearch.common.CheckedFunction; import org.elasticsearch.script.Script; import org.elasticsearch.xpack.runtimefields.AbstractLongScriptFieldScript; -import java.io.IOException; +import java.util.function.Function; public class LongScriptFieldExistsQuery extends AbstractLongScriptFieldQuery { public LongScriptFieldExistsQuery( Script script, - CheckedFunction leafFactory, + Function leafFactory, String fieldName ) { super(script, leafFactory, fieldName); diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldRangeQuery.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldRangeQuery.java index 9d20c42c84e..ce72cc22e60 100644 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldRangeQuery.java +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldRangeQuery.java @@ -7,12 +7,11 @@ package org.elasticsearch.xpack.runtimefields.query; import org.apache.lucene.index.LeafReaderContext; -import org.elasticsearch.common.CheckedFunction; import org.elasticsearch.script.Script; import org.elasticsearch.xpack.runtimefields.AbstractLongScriptFieldScript; -import java.io.IOException; import java.util.Objects; +import java.util.function.Function; public class LongScriptFieldRangeQuery extends AbstractLongScriptFieldQuery { private final long lowerValue; @@ -20,7 +19,7 @@ public class LongScriptFieldRangeQuery extends AbstractLongScriptFieldQuery { public LongScriptFieldRangeQuery( Script script, - CheckedFunction leafFactory, + Function leafFactory, String fieldName, long lowerValue, long upperValue diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldTermQuery.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldTermQuery.java index e3cc24abd76..66023e631c6 100644 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldTermQuery.java +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldTermQuery.java @@ -7,19 +7,18 @@ package org.elasticsearch.xpack.runtimefields.query; import org.apache.lucene.index.LeafReaderContext; -import org.elasticsearch.common.CheckedFunction; import org.elasticsearch.script.Script; import org.elasticsearch.xpack.runtimefields.AbstractLongScriptFieldScript; -import java.io.IOException; import java.util.Objects; +import java.util.function.Function; public class LongScriptFieldTermQuery extends AbstractLongScriptFieldQuery { private final long term; public LongScriptFieldTermQuery( Script script, - CheckedFunction leafFactory, + Function leafFactory, String fieldName, long term ) { diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldTermsQuery.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldTermsQuery.java index 4f9013882b3..b8039327df1 100644 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldTermsQuery.java +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldTermsQuery.java @@ -7,21 +7,19 @@ package org.elasticsearch.xpack.runtimefields.query; import com.carrotsearch.hppc.LongSet; - import org.apache.lucene.index.LeafReaderContext; -import org.elasticsearch.common.CheckedFunction; import org.elasticsearch.script.Script; import org.elasticsearch.xpack.runtimefields.AbstractLongScriptFieldScript; -import java.io.IOException; import java.util.Objects; +import java.util.function.Function; public class LongScriptFieldTermsQuery extends AbstractLongScriptFieldQuery { private final LongSet terms; public LongScriptFieldTermsQuery( Script script, - CheckedFunction leafFactory, + Function leafFactory, String fieldName, LongSet terms ) { diff --git a/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/query/AbstractLongScriptFieldQueryTestCase.java b/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/query/AbstractLongScriptFieldQueryTestCase.java index 7958d906a85..68721da78b3 100644 --- a/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/query/AbstractLongScriptFieldQueryTestCase.java +++ b/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/query/AbstractLongScriptFieldQueryTestCase.java @@ -7,14 +7,13 @@ package org.elasticsearch.xpack.runtimefields.query; import org.apache.lucene.index.LeafReaderContext; -import org.elasticsearch.common.CheckedFunction; import org.elasticsearch.xpack.runtimefields.AbstractLongScriptFieldScript; -import java.io.IOException; +import java.util.function.Function; public abstract class AbstractLongScriptFieldQueryTestCase extends AbstractScriptFieldQueryTestCase< T> { - protected final CheckedFunction leafFactory = ctx -> null; + protected final Function leafFactory = ctx -> null; @Override public final void testVisit() { diff --git a/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldDistanceFeatureQueryTests.java b/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldDistanceFeatureQueryTests.java index bb9d0a804ad..cd389f7f00b 100644 --- a/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldDistanceFeatureQueryTests.java +++ b/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldDistanceFeatureQueryTests.java @@ -15,7 +15,6 @@ import org.apache.lucene.search.ScoreMode; import org.apache.lucene.search.TopDocs; import org.apache.lucene.store.Directory; import org.apache.lucene.util.BytesRef; -import org.elasticsearch.common.CheckedFunction; import org.elasticsearch.script.Script; import org.elasticsearch.search.lookup.SearchLookup; import org.elasticsearch.test.ESTestCase; @@ -25,11 +24,12 @@ import org.elasticsearch.xpack.runtimefields.DateScriptFieldScript; import java.io.IOException; import java.util.Collections; import java.util.List; +import java.util.function.Function; import static org.hamcrest.Matchers.equalTo; public class LongScriptFieldDistanceFeatureQueryTests extends AbstractScriptFieldQueryTestCase { - private final CheckedFunction leafFactory = ctx -> null; + private final Function leafFactory = ctx -> null; @Override protected LongScriptFieldDistanceFeatureQuery createTestInstance() { @@ -90,15 +90,20 @@ public class LongScriptFieldDistanceFeatureQueryTests extends AbstractScriptFiel ); try (DirectoryReader reader = iw.getReader()) { IndexSearcher searcher = newSearcher(reader); - CheckedFunction leafFactory = - ctx -> new DateScriptFieldScript("test", Collections.emptyMap(), new SearchLookup(null, null, null), null, ctx) { - @Override - public void execute() { - for (Object timestamp : (List) getSource().get("timestamp")) { - emit(((Number) timestamp).longValue()); - } + Function leafFactory = ctx -> new DateScriptFieldScript( + "test", + Collections.emptyMap(), + new SearchLookup(null, null, null), + null, + ctx + ) { + @Override + public void execute() { + for (Object timestamp : (List) getSource().get("timestamp")) { + emit(((Number) timestamp).longValue()); } - }; + } + }; LongScriptFieldDistanceFeatureQuery query = new LongScriptFieldDistanceFeatureQuery( randomScript(), leafFactory,