diff --git a/server/src/main/java/org/elasticsearch/index/query/QueryShardContext.java b/server/src/main/java/org/elasticsearch/index/query/QueryShardContext.java index b550c6b5552b2..be449f625c399 100644 --- a/server/src/main/java/org/elasticsearch/index/query/QueryShardContext.java +++ b/server/src/main/java/org/elasticsearch/index/query/QueryShardContext.java @@ -297,8 +297,7 @@ MappedFieldType failIfFieldMappingNotFound(String name, MappedFieldType fieldMap public SearchLookup lookup() { if (lookup == null) { - lookup = new SearchLookup(getMapperService(), - mappedFieldType -> indexFieldDataService.apply(mappedFieldType, fullyQualifiedIndex.getName())); + lookup = new SearchLookup(getMapperService(), this::getForField); } return lookup; } diff --git a/server/src/main/java/org/elasticsearch/search/lookup/LeafDocLookup.java b/server/src/main/java/org/elasticsearch/search/lookup/LeafDocLookup.java index e06f1981f099d..1098f979cc285 100644 --- a/server/src/main/java/org/elasticsearch/search/lookup/LeafDocLookup.java +++ b/server/src/main/java/org/elasticsearch/search/lookup/LeafDocLookup.java @@ -45,8 +45,7 @@ public class LeafDocLookup implements Map> { private int docId = -1; - LeafDocLookup(MapperService mapperService, Function> fieldDataLookup, - LeafReaderContext reader) { + LeafDocLookup(MapperService mapperService, Function> fieldDataLookup, LeafReaderContext reader) { this.mapperService = mapperService; this.fieldDataLookup = fieldDataLookup; this.reader = reader; diff --git a/server/src/test/java/org/elasticsearch/search/lookup/LeafDocLookupTests.java b/server/src/test/java/org/elasticsearch/search/lookup/LeafDocLookupTests.java index 6aaa4a1ed6a58..0df1553d319dc 100644 --- a/server/src/test/java/org/elasticsearch/search/lookup/LeafDocLookupTests.java +++ b/server/src/test/java/org/elasticsearch/search/lookup/LeafDocLookupTests.java @@ -51,9 +51,7 @@ public void setUp() throws Exception { docValues = mock(ScriptDocValues.class); IndexFieldData fieldData = createFieldData(docValues); - docLookup = new LeafDocLookup(mapperService, - ignored -> fieldData, - null); + docLookup = new LeafDocLookup(mapperService, ignored -> fieldData, null); } public void testBasicLookup() { diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/AbstractRuntimeValues.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/AbstractRuntimeValues.java new file mode 100644 index 0000000000000..b9f81572e2177 --- /dev/null +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/AbstractRuntimeValues.java @@ -0,0 +1,134 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License; + * you may not use this file except in compliance with the Elastic License. + */ + +package org.elasticsearch.xpack.runtimefields; + +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 java.io.IOException; +import java.util.Objects; +import java.util.function.IntConsumer; + +/** + * Abstract base for implementing doc values and queries against values + * calculated at runtime. + */ +public abstract class AbstractRuntimeValues { + protected int count; + private boolean sort; + + private int docId = -1; + private int maxDoc; + + protected final IntConsumer leafCursor(LeafReaderContext ctx) throws IOException { + IntConsumer leafLoader = newLeafLoader(ctx); + docId = -1; + maxDoc = ctx.reader().maxDoc(); + return new IntConsumer() { + @Override + public void accept(int targetDocId) { + if (docId == targetDocId) { + return; + } + docId = targetDocId; + count = 0; + leafLoader.accept(targetDocId); + if (sort) { + sort(); + } + } + }; + } + + protected final void alwaysSortResults() { + sort = true; + } + + protected final int docId() { + return docId; + } + + protected final int maxDoc() { + return maxDoc; + } + + protected abstract IntConsumer newLeafLoader(LeafReaderContext ctx) throws IOException; + + protected abstract void sort(); + + protected abstract class AbstractRuntimeQuery extends Query { + protected final String fieldName; + + protected AbstractRuntimeQuery(String fieldName) { + this.fieldName = fieldName; + } + + @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 { + IntConsumer leafCursor = leafCursor(ctx); + DocIdSetIterator approximation = DocIdSetIterator.all(ctx.reader().maxDoc()); + TwoPhaseIterator twoPhase = new TwoPhaseIterator(approximation) { + @Override + public boolean matches() throws IOException { + leafCursor.accept(approximation.docID()); + return AbstractRuntimeQuery.this.matches(); + } + + @Override + public float matchCost() { + // TODO we don't have a good way of estimating the complexity of the script so we just go with 9000 + return approximation().cost() * 9000f; + } + }; + return new ConstantScoreScorer(this, score(), scoreMode, twoPhase); + } + }; + } + + protected abstract boolean matches(); + + @Override + public final String toString(String field) { + if (fieldName.contentEquals(field)) { + return bareToString(); + } + return fieldName + ":" + bareToString(); + } + + protected abstract String bareToString(); + + @Override + public int hashCode() { + return Objects.hash(fieldName); + } + + @Override + public boolean equals(Object obj) { + if (obj == null || getClass() != obj.getClass()) { + return false; + } + AbstractRuntimeQuery other = (AbstractRuntimeQuery) obj; + return fieldName.equals(other.fieldName); + } + } +} diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/AbstractScriptFieldScript.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/AbstractScriptFieldScript.java index ec49d38ef0a30..fc4058f9c02d8 100644 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/AbstractScriptFieldScript.java +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/AbstractScriptFieldScript.java @@ -23,7 +23,7 @@ public abstract class AbstractScriptFieldScript { private final LeafSearchLookup leafSearchLookup; public AbstractScriptFieldScript(Map params, SearchLookup searchLookup, LeafReaderContext ctx) { - this.leafSearchLookup = searchLookup.getLeafSearchLookup(ctx); + leafSearchLookup = searchLookup.getLeafSearchLookup(ctx); // TODO how do other scripts get stored fields exposed? Through asMap? I don't see any getters for them. this.params = params; } @@ -31,8 +31,8 @@ public AbstractScriptFieldScript(Map params, SearchLookup search /** * Set the document to run the script against. */ - public final void setDocument(int docId) { - this.leafSearchLookup.setDocument(docId); + public final void setDocId(int docId) { + leafSearchLookup.setDocument(docId); } /** diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/DoubleRuntimeFieldHelper.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/DoubleRuntimeFieldHelper.java new file mode 100644 index 0000000000000..4e825db4b3c57 --- /dev/null +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/DoubleRuntimeFieldHelper.java @@ -0,0 +1,256 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License; + * you may not use this file except in compliance with the Elastic License. + */ + +package org.elasticsearch.xpack.runtimefields; + +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.Term; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.QueryVisitor; +import org.apache.lucene.util.ArrayUtil; +import org.elasticsearch.common.CheckedFunction; +import org.elasticsearch.index.fielddata.SortedNumericDoubleValues; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Objects; +import java.util.function.DoubleConsumer; +import java.util.function.IntConsumer; + +/** + * Manages the creation of doc values and queries for {@code double} fields. + */ +public final class DoubleRuntimeFieldHelper { + @FunctionalInterface + public interface NewLeafLoader { + IntConsumer leafLoader(LeafReaderContext ctx, DoubleConsumer sync) throws IOException; + } + + private final NewLeafLoader newLeafLoader; + + public DoubleRuntimeFieldHelper(NewLeafLoader newLeafLoader) { + this.newLeafLoader = newLeafLoader; + } + + public CheckedFunction docValues() { + return new Values().docValues(); + } + + public Query existsQuery(String fieldName) { + return new Values().new ExistsQuery(fieldName); + } + + public Query rangeQuery(String fieldName, double lowerValue, double upperValue) { + return new Values().new RangeQuery(fieldName, lowerValue, upperValue); + } + + public Query termQuery(String fieldName, double value) { + return new Values().new TermQuery(fieldName, value); + } + + public Query termsQuery(String fieldName, double... value) { + return new Values().new TermsQuery(fieldName, value); + } + + private class Values extends AbstractRuntimeValues { + private double[] values = new double[1]; + + @Override + protected IntConsumer newLeafLoader(LeafReaderContext ctx) throws IOException { + return newLeafLoader.leafLoader(ctx, this::add); + } + + private void add(double value) { + int newCount = count + 1; + if (values.length < newCount) { + values = Arrays.copyOf(values, ArrayUtil.oversize(newCount, 8)); + } + values[count] = value; + count = newCount; + } + + @Override + protected void sort() { + Arrays.sort(values, 0, count); + } + + private CheckedFunction docValues() { + alwaysSortResults(); + return DocValues::new; + } + + private class DocValues extends SortedNumericDoubleValues { + private final IntConsumer leafCursor; + private int next; + + DocValues(LeafReaderContext ctx) throws IOException { + leafCursor = leafCursor(ctx); + } + + @Override + public double nextValue() throws IOException { + return values[next++]; + } + + @Override + public int docValueCount() { + return count; + } + + @Override + public boolean advanceExact(int target) throws IOException { + leafCursor.accept(target); + next = 0; + return count > 0; + } + } + + private class ExistsQuery extends AbstractRuntimeQuery { + private ExistsQuery(String fieldName) { + super(fieldName); + } + + @Override + protected boolean matches() { + return count > 0; + } + + @Override + protected String bareToString() { + return "*"; + } + } + + private class RangeQuery extends AbstractRuntimeQuery { + private final double lowerValue; + private final double upperValue; + + private RangeQuery(String fieldName, double lowerValue, double upperValue) { + super(fieldName); + this.lowerValue = lowerValue; + this.upperValue = upperValue; + } + + @Override + protected boolean matches() { + for (int i = 0; i < count; i++) { + if (lowerValue <= values[i] && values[i] <= upperValue) { + return true; + } + } + return false; + } + + @Override + protected String bareToString() { + return "[" + lowerValue + "," + upperValue + "]"; + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), lowerValue, upperValue); + } + + @Override + public boolean equals(Object obj) { + if (false == super.equals(obj)) { + return false; + } + RangeQuery other = (RangeQuery) obj; + return lowerValue == other.lowerValue && upperValue == other.upperValue; + } + } + + private class TermQuery extends AbstractRuntimeQuery { + private final double term; + + private TermQuery(String fieldName, double term) { + super(fieldName); + this.term = term; + } + + @Override + protected boolean matches() { + for (int i = 0; i < count; i++) { + if (term == values[i]) { + return true; + } + } + return false; + } + + @Override + public void visit(QueryVisitor visitor) { + visitor.consumeTerms(this, new Term(fieldName, Double.toString(term))); + } + + @Override + protected String bareToString() { + return Double.toString(term); + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), term); + } + + @Override + public boolean equals(Object obj) { + if (false == super.equals(obj)) { + return false; + } + TermQuery other = (TermQuery) obj; + return term == other.term; + } + } + + private class TermsQuery extends AbstractRuntimeQuery { + private final double[] terms; + + private TermsQuery(String fieldName, double[] terms) { + super(fieldName); + this.terms = terms.clone(); + Arrays.sort(terms); + } + + @Override + protected boolean matches() { + for (int i = 0; i < count; i++) { + if (Arrays.binarySearch(terms, values[i]) >= 0) { + return true; + } + } + return false; + } + + @Override + public void visit(QueryVisitor visitor) { + for (double term : terms) { + visitor.consumeTerms(this, new Term(fieldName, Double.toString(term))); + } + } + + @Override + protected String bareToString() { + return "{" + Arrays.toString(terms) + "}"; + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), Arrays.hashCode(terms)); + } + + @Override + public boolean equals(Object obj) { + if (false == super.equals(obj)) { + return false; + } + TermsQuery other = (TermsQuery) obj; + return Arrays.equals(terms, other.terms); + } + } + } +} 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 8b66738b8bc74..61d5364a30e6b 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 @@ -25,14 +25,41 @@ static List whitelist() { return List.of(WhitelistLoader.loadFromResourceFiles(RuntimeFieldsPainlessExtension.class, "double_whitelist.txt")); } + /** + * Magic constant that painless needs to name the parameters. There aren't any so it is empty. + */ public static final String[] PARAMETERS = {}; + /** + * Factory for building instances of the script for a particular search context. + */ public interface Factory extends ScriptFactory { LeafFactory newFactory(Map params, SearchLookup searchLookup); } + /** + * Factory for building the script for a particular leaf or for building + * runtime values which manages the creation of doc values and queries. + */ public interface LeafFactory { + /** + * Build a new script. + */ DoubleScriptFieldScript newInstance(LeafReaderContext ctx, DoubleConsumer sync) throws IOException; + + /** + * Build an {@link DoubleRuntimeFieldHelper} to manage creation of doc + * values and queries using the script. + */ + default DoubleRuntimeFieldHelper runtimeFieldHelper() { + return new DoubleRuntimeFieldHelper((ctx, sync) -> { + DoubleScriptFieldScript script = newInstance(ctx, sync); + return docId -> { + script.setDocId(docId); + script.execute(); + }; + }); + } } private final DoubleConsumer sync; diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/ForceNoBulkScoringQuery.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/ForceNoBulkScoringQuery.java new file mode 100644 index 0000000000000..518b4fb9a06df --- /dev/null +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/ForceNoBulkScoringQuery.java @@ -0,0 +1,113 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License; + * you may not use this file except in compliance with the Elastic License. + */ + +package org.elasticsearch.xpack.runtimefields; + +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.Term; +import org.apache.lucene.search.BulkScorer; +import org.apache.lucene.search.Explanation; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.Matches; +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.Weight; + +import java.io.IOException; +import java.util.Objects; +import java.util.Set; + +/** + * Forces wrapped queries to score documents one at a time which is important + * for runtime queries so they can more in lock step with one another. + *

+ * Inspired by the ForceNoBulkScoringQuery in Lucene's monitor project. + */ +class ForceNoBulkScoringQuery extends Query { + private final Query delegate; + + ForceNoBulkScoringQuery(Query inner) { + this.delegate = inner; + } + + @Override + public Query rewrite(IndexReader reader) throws IOException { + Query rewritten = delegate.rewrite(reader); + if (rewritten != delegate) { + return new ForceNoBulkScoringQuery(rewritten); + } + return this; + } + + @Override + public void visit(QueryVisitor visitor) { + delegate.visit(visitor); + } + + @Override + public boolean equals(Object o) { + if (o == null || getClass() != o.getClass()) { + return false; + } + ForceNoBulkScoringQuery that = (ForceNoBulkScoringQuery) o; + return Objects.equals(delegate, that.delegate); + } + + @Override + public int hashCode() { + return Objects.hash(delegate); + } + + @Override + public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { + Weight innerWeight = delegate.createWeight(searcher, scoreMode, boost); + return new Weight(this) { + @Override + public BulkScorer bulkScorer(LeafReaderContext context) throws IOException { + /* + * Intentionally do not delegate this to the innerWeight so + * we don't use its BulkScorer. This is the magic that causes + * us to skip bulk scoring. + */ + return super.bulkScorer(context); + } + + @Override + public boolean isCacheable(LeafReaderContext ctx) { + return innerWeight.isCacheable(ctx); + } + + @Override + public Explanation explain(LeafReaderContext leafReaderContext, int i) throws IOException { + return innerWeight.explain(leafReaderContext, i); + } + + @Override + public Scorer scorer(LeafReaderContext leafReaderContext) throws IOException { + return innerWeight.scorer(leafReaderContext); + } + + @Override + public Matches matches(LeafReaderContext context, int doc) throws IOException { + return innerWeight.matches(context, doc); + } + + @Override + @Deprecated + public void extractTerms(Set terms) { + innerWeight.extractTerms(terms); + } + }; + } + + @Override + public String toString(String s) { + return "NoBulkScorer(" + delegate.toString(s) + ")"; + } +} diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/LongRuntimeFieldHelper.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/LongRuntimeFieldHelper.java new file mode 100644 index 0000000000000..d52a68d294dae --- /dev/null +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/LongRuntimeFieldHelper.java @@ -0,0 +1,284 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License; + * you may not use this file except in compliance with the Elastic License. + */ + +package org.elasticsearch.xpack.runtimefields; + +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.SortedNumericDocValues; +import org.apache.lucene.index.Term; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.QueryVisitor; +import org.apache.lucene.util.ArrayUtil; +import org.elasticsearch.common.CheckedFunction; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Objects; +import java.util.function.IntConsumer; +import java.util.function.LongConsumer; + +/** + * Manages the creation of doc values and queries for {@code long} fields. + */ +public final class LongRuntimeFieldHelper { + @FunctionalInterface + public interface NewLeafLoader { + IntConsumer leafLoader(LeafReaderContext ctx, LongConsumer sync) throws IOException; + } + + private final NewLeafLoader newLeafLoader; + + public LongRuntimeFieldHelper(NewLeafLoader newLeafLoader) { + this.newLeafLoader = newLeafLoader; + } + + public CheckedFunction docValues() { + return new Values().docValues(); + } + + public Query existsQuery(String fieldName) { + return new Values().new ExistsQuery(fieldName); + } + + public Query rangeQuery(String fieldName, long lowerValue, long upperValue) { + return new Values().new RangeQuery(fieldName, lowerValue, upperValue); + } + + public Query termQuery(String fieldName, long value) { + return new Values().new TermQuery(fieldName, value); + } + + public Query termsQuery(String fieldName, long... value) { + return new Values().new TermsQuery(fieldName, value); + } + + private class Values extends AbstractRuntimeValues { + private long[] values = new long[1]; + + @Override + protected IntConsumer newLeafLoader(LeafReaderContext ctx) throws IOException { + return newLeafLoader.leafLoader(ctx, this::add); + } + + private void add(long value) { + int newCount = count + 1; + if (values.length < newCount) { + values = Arrays.copyOf(values, ArrayUtil.oversize(newCount, 8)); + } + values[count] = value; + count = newCount; + } + + @Override + protected void sort() { + Arrays.sort(values, 0, count); + } + + private CheckedFunction docValues() { + alwaysSortResults(); + return DocValues::new; + } + + private class DocValues extends SortedNumericDocValues { + private final IntConsumer leafCursor; + private int next; + + DocValues(LeafReaderContext ctx) throws IOException { + leafCursor = leafCursor(ctx); + } + + @Override + public long nextValue() throws IOException { + return values[next++]; + } + + @Override + public int docValueCount() { + return count; + } + + @Override + public boolean advanceExact(int target) throws IOException { + leafCursor.accept(target); + next = 0; + return count > 0; + } + + @Override + public int docID() { + return docId(); + } + + @Override + public int nextDoc() throws IOException { + return advance(docId() + 1); + } + + @Override + public int advance(int target) throws IOException { + int current = target; + while (current < maxDoc()) { + if (advanceExact(current)) { + return current; + } + current++; + } + return NO_MORE_DOCS; + } + + @Override + public long cost() { + // TODO we have no idea what this should be and no real way to get one + return 1000; + } + } + + private class ExistsQuery extends AbstractRuntimeQuery { + private ExistsQuery(String fieldName) { + super(fieldName); + } + + @Override + protected boolean matches() { + return count > 0; + } + + @Override + protected String bareToString() { + return "*"; + } + } + + private class RangeQuery extends AbstractRuntimeQuery { + private final long lowerValue; + private final long upperValue; + + private RangeQuery(String fieldName, long lowerValue, long upperValue) { + super(fieldName); + this.lowerValue = lowerValue; + this.upperValue = upperValue; + } + + @Override + protected boolean matches() { + for (int i = 0; i < count; i++) { + if (lowerValue <= values[i] && values[i] <= upperValue) { + return true; + } + } + return false; + } + + @Override + protected String bareToString() { + return "[" + lowerValue + "," + upperValue + "]"; + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), lowerValue, upperValue); + } + + @Override + public boolean equals(Object obj) { + if (false == super.equals(obj)) { + return false; + } + RangeQuery other = (RangeQuery) obj; + return lowerValue == other.lowerValue && upperValue == other.upperValue; + } + } + + private class TermQuery extends AbstractRuntimeQuery { + private final long term; + + private TermQuery(String fieldName, long term) { + super(fieldName); + this.term = term; + } + + @Override + protected boolean matches() { + for (int i = 0; i < count; i++) { + if (term == values[i]) { + return true; + } + } + return false; + } + + @Override + public void visit(QueryVisitor visitor) { + visitor.consumeTerms(this, new Term(fieldName, Long.toString(term))); + } + + @Override + protected String bareToString() { + return Long.toString(term); + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), term); + } + + @Override + public boolean equals(Object obj) { + if (false == super.equals(obj)) { + return false; + } + TermQuery other = (TermQuery) obj; + return term == other.term; + } + } + + private class TermsQuery extends AbstractRuntimeQuery { + private final long[] terms; + + private TermsQuery(String fieldName, long[] terms) { + super(fieldName); + this.terms = terms.clone(); + Arrays.sort(terms); + } + + @Override + protected boolean matches() { + for (int i = 0; i < count; i++) { + if (Arrays.binarySearch(terms, values[i]) >= 0) { + return true; + } + } + return false; + } + + @Override + public void visit(QueryVisitor visitor) { + for (long term : terms) { + visitor.consumeTerms(this, new Term(fieldName, Long.toString(term))); + } + } + + @Override + protected String bareToString() { + return "{" + Arrays.toString(terms) + "}"; + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), Arrays.hashCode(terms)); + } + + @Override + public boolean equals(Object obj) { + if (false == super.equals(obj)) { + return false; + } + TermsQuery other = (TermsQuery) obj; + return Arrays.equals(terms, other.terms); + } + } + } +} 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 9850cab424642..6a316b71be2c4 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 @@ -25,14 +25,41 @@ static List whitelist() { return List.of(WhitelistLoader.loadFromResourceFiles(RuntimeFieldsPainlessExtension.class, "long_whitelist.txt")); } + /** + * Magic constant that painless needs to name the parameters. There aren't any so it is empty. + */ public static final String[] PARAMETERS = {}; + /** + * Factory for building instances of the script for a particular search context. + */ public interface Factory extends ScriptFactory { LeafFactory newFactory(Map params, SearchLookup searchLookup); } + /** + * Factory for building the script for a particular leaf or for building + * runtime values which manages the creation of doc values and queries. + */ public interface LeafFactory { + /** + * Build a new script. + */ LongScriptFieldScript newInstance(LeafReaderContext ctx, LongConsumer sync) throws IOException; + + /** + * Build an {@link LongRuntimeFieldHelper} to manage creation of doc + * values and queries using the script. + */ + default LongRuntimeFieldHelper runtimeFieldHelper() { + return new LongRuntimeFieldHelper((ctx, sync) -> { + LongScriptFieldScript script = newInstance(ctx, sync); + return docId -> { + script.setDocId(docId); + script.execute(); + }; + }); + } } private final LongConsumer sync; diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/StringRuntimeFieldHelper.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/StringRuntimeFieldHelper.java new file mode 100644 index 0000000000000..ba3375ad9de70 --- /dev/null +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/StringRuntimeFieldHelper.java @@ -0,0 +1,491 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License; + * you may not use this file except in compliance with the Elastic License. + */ + +package org.elasticsearch.xpack.runtimefields; + +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.Term; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.QueryVisitor; +import org.apache.lucene.util.ArrayUtil; +import org.apache.lucene.util.BytesRef; +import org.apache.lucene.util.BytesRefBuilder; +import org.apache.lucene.util.RamUsageEstimator; +import org.apache.lucene.util.automaton.Automaton; +import org.apache.lucene.util.automaton.ByteRunAutomaton; +import org.apache.lucene.util.automaton.CompiledAutomaton; +import org.apache.lucene.util.automaton.RegExp; +import org.elasticsearch.common.CheckedFunction; +import org.elasticsearch.index.fielddata.SortedBinaryDocValues; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Objects; +import java.util.function.Consumer; +import java.util.function.IntConsumer; + +/** + * Manages the creation of doc values and queries for {@link String} fields. + */ +public final class StringRuntimeFieldHelper { + @FunctionalInterface + public interface NewLeafLoader { + IntConsumer leafLoader(LeafReaderContext ctx, Consumer sync) throws IOException; + } + + private final NewLeafLoader newLeafLoader; + + public StringRuntimeFieldHelper(NewLeafLoader newLeafLoader) { + this.newLeafLoader = newLeafLoader; + } + + public CheckedFunction docValues() { + return new Values().docValues(); + } + + public Query existsQuery(String fieldName) { + return new Values().new ExistsQuery(fieldName); + } + + public Query fuzzyQuery(String fieldName, String value, int maxEdits, int prefixLength, int maxExpansions, boolean transpositions) { + return new Values().new FuzzyQuery(fieldName, value, maxEdits, prefixLength, maxExpansions, transpositions); + } + + public Query prefixQuery(String fieldName, String value) { + return new Values().new PrefixQuery(fieldName, value); + } + + public Query rangeQuery(String fieldName, String lowerValue, String upperValue, boolean includeLower, boolean includeUpper) { + return new Values().new RangeQuery(fieldName, lowerValue, upperValue, includeLower, includeUpper); + } + + public Query regexpQuery(String fieldName, String pattern, int flags, int maxDeterminizedStates) { + return new Values().new RegexpQuery(fieldName, pattern, flags, maxDeterminizedStates); + } + + public Query termQuery(String fieldName, String value) { + return new Values().new TermQuery(fieldName, value); + } + + public Query termsQuery(String fieldName, String... values) { + return new Values().new TermsQuery(fieldName, values); + } + + public Query wildcardQuery(String fieldName, String pattern) { + return new Values().new WildcardQuery(fieldName, pattern); + } + + private class Values extends AbstractRuntimeValues { + private String[] values = new String[1]; + + @Override + protected IntConsumer newLeafLoader(LeafReaderContext ctx) throws IOException { + return newLeafLoader.leafLoader(ctx, this::add); + } + + private void add(String value) { + if (value == null) { + return; + } + int newCount = count + 1; + if (values.length < newCount) { + values = Arrays.copyOf(values, ArrayUtil.oversize(newCount, RamUsageEstimator.NUM_BYTES_OBJECT_REF)); + } + values[count] = value; + count = newCount; + } + + @Override + protected void sort() { + Arrays.sort(values, 0, count); + } + + private CheckedFunction docValues() { + alwaysSortResults(); + return DocValues::new; + } + + private class DocValues extends SortedBinaryDocValues { + private final BytesRefBuilder ref = new BytesRefBuilder(); + private final IntConsumer leafCursor; + private int next; + + DocValues(LeafReaderContext ctx) throws IOException { + leafCursor = leafCursor(ctx); + } + + @Override + public boolean advanceExact(int docId) throws IOException { + leafCursor.accept(docId); + next = 0; + return count > 0; + } + + @Override + public int docValueCount() { + return count; + } + + @Override + public BytesRef nextValue() throws IOException { + ref.copyChars(values[next++]); + return ref.get(); + } + } + + private class ExistsQuery extends AbstractRuntimeQuery { + private ExistsQuery(String fieldName) { + super(fieldName); + } + + @Override + protected boolean matches() { + return count > 0; + } + + @Override + protected String bareToString() { + return "*"; + } + } + + private class FuzzyQuery extends AbstractRuntimeQuery { + private final BytesRefBuilder scratch = new BytesRefBuilder(); + private final String term; + private final org.apache.lucene.search.FuzzyQuery delegate; + private final CompiledAutomaton automaton; + + private FuzzyQuery(String fieldName, String term, int maxEdits, int prefixLength, int maxExpansions, boolean transpositions) { + super(fieldName); + this.term = term; + delegate = new org.apache.lucene.search.FuzzyQuery( + new Term(fieldName, term), + maxEdits, + prefixLength, + maxExpansions, + transpositions + ); + automaton = delegate.getAutomata(); + if (automaton.type != CompiledAutomaton.AUTOMATON_TYPE.NORMAL) { + // TODO I'll bet we have to actually implement all of these types + throw new IllegalArgumentException("Can't compile automaton for [" + delegate + "]"); + } + } + + @Override + protected boolean matches() { + for (int i = 0; i < count; i++) { + scratch.copyChars(values[i]); + if (automaton.runAutomaton.run(scratch.bytes(), 0, scratch.length())) { + return true; + } + } + return false; + } + + @Override + public void visit(QueryVisitor visitor) { + visitor.consumeTerms(this, new Term(fieldName, term)); + } + + @Override + protected String bareToString() { + return term + "~" + delegate.getMaxEdits(); + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), term); + } + + @Override + public boolean equals(Object obj) { + if (false == super.equals(obj)) { + return false; + } + FuzzyQuery other = (FuzzyQuery) obj; + return delegate.equals(other.delegate); + } + } + + private class PrefixQuery extends AbstractRuntimeQuery { + private final String prefix; + + private PrefixQuery(String fieldName, String prefix) { + super(fieldName); + this.prefix = prefix; + } + + @Override + protected boolean matches() { + for (int i = 0; i < count; i++) { + if (values[i].startsWith(prefix)) { + return true; + } + } + return false; + } + + @Override + protected String bareToString() { + return prefix + "*"; + } + + @Override + public void visit(QueryVisitor visitor) { + visitor.consumeTermsMatching( + this, + fieldName, + () -> new ByteRunAutomaton( + org.apache.lucene.search.PrefixQuery.toAutomaton(new BytesRef(prefix)), + true, + Integer.MAX_VALUE + ) + ); + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), prefix); + } + + @Override + public boolean equals(Object obj) { + if (false == super.equals(obj)) { + return false; + } + PrefixQuery other = (PrefixQuery) obj; + return prefix.equals(other.prefix); + } + } + + private class TermQuery extends AbstractRuntimeQuery { + private final String term; + + private TermQuery(String fieldName, String term) { + super(fieldName); + this.term = term; + } + + @Override + protected boolean matches() { + for (int i = 0; i < count; i++) { + if (term.equals(values[i])) { + return true; + } + } + return false; + } + + @Override + public void visit(QueryVisitor visitor) { + visitor.consumeTerms(this, new Term(fieldName, term)); + } + + @Override + protected String bareToString() { + return term; + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), term); + } + + @Override + public boolean equals(Object obj) { + if (false == super.equals(obj)) { + return false; + } + TermQuery other = (TermQuery) obj; + return term.equals(other.term); + } + } + + private class TermsQuery extends AbstractRuntimeQuery { + private final String[] terms; + + private TermsQuery(String fieldName, String[] terms) { + super(fieldName); + this.terms = terms.clone(); + Arrays.sort(terms); + } + + @Override + protected boolean matches() { + for (int i = 0; i < count; i++) { + if (Arrays.binarySearch(terms, values[i]) >= 0) { + return true; + } + } + return false; + } + + @Override + public void visit(QueryVisitor visitor) { + for (String term : terms) { + visitor.consumeTerms(this, new Term(fieldName, term)); + } + } + + @Override + protected String bareToString() { + return "{" + Arrays.toString(terms) + "}"; + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), Arrays.hashCode(terms)); + } + + @Override + public boolean equals(Object obj) { + if (false == super.equals(obj)) { + return false; + } + TermsQuery other = (TermsQuery) obj; + return Arrays.equals(terms, other.terms); + } + } + + private class RangeQuery extends AbstractRuntimeQuery { + private final String lowerValue; + private final String upperValue; + private final boolean includeLower; + private final boolean includeUpper; + + private RangeQuery(String fieldName, String lowerValue, String upperValue, boolean includeLower, boolean includeUpper) { + super(fieldName); + this.lowerValue = lowerValue; + this.upperValue = upperValue; + this.includeLower = includeLower; + this.includeUpper = includeUpper; + assert lowerValue.compareTo(upperValue) <= 0; + } + + @Override + protected boolean matches() { + for (int i = 0; i < count; i++) { + int lct = lowerValue.compareTo(values[i]); + boolean lowerOk = includeLower ? lct <= 0 : lct < 0; + if (lowerOk) { + int uct = upperValue.compareTo(values[i]); + boolean upperOk = includeUpper ? uct >= 0 : uct > 0; + if (upperOk) { + return true; + } + } + } + return false; + } + + @Override + protected String bareToString() { + return "[" + lowerValue + "," + upperValue + "]"; + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), lowerValue, upperValue); + } + + @Override + public boolean equals(Object obj) { + if (false == super.equals(obj)) { + return false; + } + RangeQuery other = (RangeQuery) obj; + return lowerValue.equals(other.lowerValue) && upperValue.equals(other.upperValue); + } + } + + private class RegexpQuery extends AbstractAutomatonQuery { + private final String pattern; + private final int flags; + + private RegexpQuery(String fieldName, String pattern, int flags, int maxDeterminizedStates) { + super(fieldName, new RegExp(pattern, flags).toAutomaton(maxDeterminizedStates)); + this.pattern = pattern; + this.flags = flags; + } + + @Override + protected String bareToString() { + return "/" + pattern + "/"; + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), pattern, flags); + } + + @Override + public boolean equals(Object obj) { + if (false == super.equals(obj)) { + return false; + } + RegexpQuery other = (RegexpQuery) obj; + return pattern.equals(other.pattern) && flags == other.flags; + } + } + + private class WildcardQuery extends AbstractAutomatonQuery { + private final String pattern; + + private WildcardQuery(String fieldName, String pattern) { + super(fieldName, org.apache.lucene.search.WildcardQuery.toAutomaton(new Term(fieldName, pattern))); + this.pattern = pattern; + } + + @Override + protected String bareToString() { + return pattern; + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), pattern); + } + + @Override + public boolean equals(Object obj) { + if (false == super.equals(obj)) { + return false; + } + WildcardQuery other = (WildcardQuery) obj; + return pattern.equals(other.pattern); + } + } + + private abstract class AbstractAutomatonQuery extends AbstractRuntimeQuery { + private final BytesRefBuilder scratch = new BytesRefBuilder(); + private final ByteRunAutomaton automaton; + + private AbstractAutomatonQuery(String fieldName, Automaton automaton) { + super(fieldName); + this.automaton = new ByteRunAutomaton(automaton); + } + + @Override + protected final boolean matches() { + for (int i = 0; i < count; i++) { + scratch.copyChars(values[i]); + if (automaton.run(scratch.bytes(), 0, scratch.length())) { + return true; + } + } + return false; + } + + @Override + public final void visit(QueryVisitor visitor) { + if (visitor.acceptField(fieldName)) { + visitor.consumeTermsMatching(this, fieldName, () -> automaton); + } + } + } + } +} 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 391be84b8d964..a93d4b70060a1 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 @@ -18,6 +18,9 @@ import java.util.Map; import java.util.function.Consumer; +/** + * Script for building {@link String} values at runtime. + */ public abstract class StringScriptFieldScript extends AbstractScriptFieldScript { public static final ScriptContext CONTEXT = new ScriptContext<>("string_script_field", Factory.class); @@ -25,14 +28,41 @@ static List whitelist() { return List.of(WhitelistLoader.loadFromResourceFiles(RuntimeFieldsPainlessExtension.class, "string_whitelist.txt")); } + /** + * Magic constant that painless needs to name the parameters. There aren't any so it is empty. + */ public static final String[] PARAMETERS = {}; + /** + * Factory for building instances of the script for a particular search context. + */ public interface Factory extends ScriptFactory { LeafFactory newFactory(Map params, SearchLookup searchLookup); } + /** + * Factory for building the script for a particular leaf or for building + * runtime values which manages the creation of doc values and queries. + */ public interface LeafFactory { + /** + * Build a new script. + */ StringScriptFieldScript newInstance(LeafReaderContext ctx, Consumer sync) throws IOException; + + /** + * Build an {@link StringRuntimeFieldHelper} to manage creation of doc + * values and queries using the script. + */ + default StringRuntimeFieldHelper runtimeFieldHelper() { + return new StringRuntimeFieldHelper((ctx, sync) -> { + StringScriptFieldScript script = newInstance(ctx, sync); + return docId -> { + script.setDocId(docId); + script.execute(); + }; + }); + } } private final Consumer sync; diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/fielddata/ScriptBinaryDocValues.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/fielddata/ScriptBinaryDocValues.java deleted file mode 100644 index ae0edeae46b4e..0000000000000 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/fielddata/ScriptBinaryDocValues.java +++ /dev/null @@ -1,41 +0,0 @@ -/* - * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one - * or more contributor license agreements. Licensed under the Elastic License; - * you may not use this file except in compliance with the Elastic License. - */ - -package org.elasticsearch.xpack.runtimefields.fielddata; - -import org.elasticsearch.index.fielddata.SortingBinaryDocValues; -import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript; - -public final class ScriptBinaryDocValues extends SortingBinaryDocValues { - - private final StringScriptFieldScript script; - private final ScriptBinaryFieldData.ScriptBinaryResult scriptBinaryResult; - - ScriptBinaryDocValues(StringScriptFieldScript script, ScriptBinaryFieldData.ScriptBinaryResult scriptBinaryResult) { - this.script = script; - this.scriptBinaryResult = scriptBinaryResult; - } - - @Override - public boolean advanceExact(int doc) { - script.setDocument(doc); - script.execute(); - - count = scriptBinaryResult.getResult().size(); - if (count == 0) { - grow(); - return false; - } - - int i = 0; - for (String value : scriptBinaryResult.getResult()) { - grow(); - values[i++].copyChars(value); - } - sort(); - return true; - } -} diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/fielddata/ScriptBinaryFieldData.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/fielddata/ScriptBinaryFieldData.java index 42ca35731274b..10999c17de694 100644 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/fielddata/ScriptBinaryFieldData.java +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/fielddata/ScriptBinaryFieldData.java @@ -10,6 +10,7 @@ import org.apache.lucene.search.SortField; import org.apache.lucene.util.SetOnce; import org.elasticsearch.ElasticsearchException; +import org.elasticsearch.common.CheckedFunction; import org.elasticsearch.common.util.BigArrays; import org.elasticsearch.index.AbstractIndexComponent; import org.elasticsearch.index.IndexSettings; @@ -33,11 +34,9 @@ import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript; import java.io.IOException; -import java.util.ArrayList; import java.util.Collections; -import java.util.List; -public final class ScriptBinaryFieldData extends AbstractIndexComponent +public final class ScriptBinaryFieldData extends AbstractIndexComponent // TODO do we need to extends AbstractIndexComponent? implements IndexFieldData, SearchLookupAware { @@ -64,7 +63,7 @@ public IndexFieldData build( private final String fieldName; private final StringScriptFieldScript.Factory scriptFactory; - private final SetOnce leafFactory = new SetOnce<>(); + private final SetOnce> docValuesBuilder = new SetOnce<>(); private ScriptBinaryFieldData(IndexSettings indexSettings, String fieldName, StringScriptFieldScript.Factory scriptFactory) { super(indexSettings); @@ -72,9 +71,11 @@ private ScriptBinaryFieldData(IndexSettings indexSettings, String fieldName, Str this.scriptFactory = scriptFactory; } + @Override public void setSearchLookup(SearchLookup searchLookup) { // TODO wire the params from the mappings definition, we don't parse them yet - this.leafFactory.set(scriptFactory.newFactory(Collections.emptyMap(), searchLookup)); + // TODO it'd be nice if we could stuff `runtimeValues` some place into the search context so we could reuse it + this.docValuesBuilder.set(scriptFactory.newFactory(Collections.emptyMap(), searchLookup).runtimeFieldHelper().docValues()); } @Override @@ -102,10 +103,7 @@ public ScriptBinaryLeafFieldData load(LeafReaderContext context) { @Override public ScriptBinaryLeafFieldData loadDirect(LeafReaderContext context) throws IOException { - ScriptBinaryResult scriptBinaryResult = new ScriptBinaryResult(); - return new ScriptBinaryLeafFieldData( - new ScriptBinaryDocValues(leafFactory.get().newInstance(context, scriptBinaryResult::accept), scriptBinaryResult) - ); + return new ScriptBinaryLeafFieldData(docValuesBuilder.get().apply(context)); } @Override @@ -134,10 +132,10 @@ public void clear() { } static class ScriptBinaryLeafFieldData implements LeafFieldData { - private final ScriptBinaryDocValues scriptBinaryDocValues; + private final SortedBinaryDocValues docValues; - ScriptBinaryLeafFieldData(ScriptBinaryDocValues scriptBinaryDocValues) { - this.scriptBinaryDocValues = scriptBinaryDocValues; + ScriptBinaryLeafFieldData(SortedBinaryDocValues docValues) { + this.docValues = docValues; } @Override @@ -147,7 +145,7 @@ public ScriptDocValues getScriptValues() { @Override public SortedBinaryDocValues getBytesValues() { - return scriptBinaryDocValues; + return docValues; } @Override @@ -160,16 +158,4 @@ public void close() { } } - - static class ScriptBinaryResult { - private final List result = new ArrayList<>(); - - void accept(String value) { - this.result.add(value); - } - - List getResult() { - return result; - } - } } diff --git a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/mapper/RuntimeKeywordMappedFieldType.java b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/mapper/RuntimeKeywordMappedFieldType.java index 5aef8e857408a..0a5eceef35b4c 100644 --- a/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/mapper/RuntimeKeywordMappedFieldType.java +++ b/x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/mapper/RuntimeKeywordMappedFieldType.java @@ -6,19 +6,27 @@ package org.elasticsearch.xpack.runtimefields.mapper; +import org.apache.lucene.search.MultiTermQuery.RewriteMethod; import org.apache.lucene.search.Query; import org.apache.lucene.util.BytesRef; -import org.elasticsearch.common.xcontent.XContentBuilder; +import org.elasticsearch.common.geo.ShapeRelation; +import org.elasticsearch.common.lucene.BytesRefs; +import org.elasticsearch.common.time.DateMathParser; +import org.elasticsearch.common.unit.Fuzziness; import org.elasticsearch.common.xcontent.ToXContent.Params; +import org.elasticsearch.common.xcontent.XContentBuilder; import org.elasticsearch.index.fielddata.IndexFieldData; import org.elasticsearch.index.mapper.MappedFieldType; import org.elasticsearch.index.mapper.TextSearchInfo; import org.elasticsearch.index.query.QueryShardContext; import org.elasticsearch.script.Script; +import org.elasticsearch.xpack.runtimefields.StringRuntimeFieldHelper; import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript; import org.elasticsearch.xpack.runtimefields.fielddata.ScriptBinaryFieldData; import java.io.IOException; +import java.time.ZoneId; +import java.util.List; import java.util.Map; public final class RuntimeKeywordMappedFieldType extends MappedFieldType { @@ -27,7 +35,7 @@ public final class RuntimeKeywordMappedFieldType extends MappedFieldType { private final StringScriptFieldScript.Factory scriptFactory; RuntimeKeywordMappedFieldType(String name, Script script, StringScriptFieldScript.Factory scriptFactory, Map meta) { - super(name, false, false, TextSearchInfo.NONE, meta); + super(name, false, false, TextSearchInfo.SIMPLE_MATCH_ONLY, meta); this.script = script; this.scriptFactory = scriptFactory; } @@ -60,20 +68,71 @@ public String typeName() { return ScriptFieldMapper.CONTENT_TYPE; } + private StringRuntimeFieldHelper helper(QueryShardContext ctx) { + return scriptFactory.newFactory(script.getParams(), ctx.lookup()).runtimeFieldHelper(); + } + @Override public IndexFieldData.Builder fielddataBuilder(String fullyQualifiedIndexName) { // TODO once we get SearchLookup as an argument, we can already call scriptFactory.newFactory here and pass through the result return new ScriptBinaryFieldData.Builder(scriptFactory); } + @Override + public Query existsQuery(QueryShardContext context) { + return helper(context).existsQuery(name()); + } + + @Override + public Query fuzzyQuery( + Object value, + Fuzziness fuzziness, + int prefixLength, + int maxExpansions, + boolean transpositions, + QueryShardContext context + ) { + String term = BytesRefs.toString(value); + return helper(context).fuzzyQuery(name(), term, fuzziness.asDistance(term), prefixLength, maxExpansions, transpositions); + } + + @Override + public Query prefixQuery(String value, RewriteMethod method, QueryShardContext context) { + return helper(context).prefixQuery(name(), value); + } + + @Override + public Query rangeQuery( + Object lowerTerm, + Object upperTerm, + boolean includeLower, + boolean includeUpper, + ShapeRelation relation, + ZoneId timeZone, + DateMathParser parser, + QueryShardContext context + ) { + return helper(context).rangeQuery(name(), BytesRefs.toString(lowerTerm), BytesRefs.toString(upperTerm), includeLower, includeUpper); + } + + @Override + public Query regexpQuery(String value, int flags, int maxDeterminizedStates, RewriteMethod method, QueryShardContext context) { + return helper(context).regexpQuery(name(), value, flags, maxDeterminizedStates); + } + @Override public Query termQuery(Object value, QueryShardContext context) { - return null; + return helper(context).termQuery(name(), BytesRefs.toString(value)); } @Override - public Query existsQuery(QueryShardContext context) { - return null; + public Query termsQuery(List values, QueryShardContext context) { + return helper(context).termsQuery(name(), values.stream().map(BytesRefs::toString).toArray(String[]::new)); + } + + @Override + public Query wildcardQuery(String value, RewriteMethod method, QueryShardContext context) { + return helper(context).wildcardQuery(name(), value); } void doXContentBody(XContentBuilder builder, boolean includeDefaults, Params params) throws IOException { diff --git a/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/DoubleScriptFieldScriptTests.java b/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/DoubleScriptFieldScriptTests.java index 23081d1ff3427..e47ce41040f11 100644 --- a/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/DoubleScriptFieldScriptTests.java +++ b/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/DoubleScriptFieldScriptTests.java @@ -8,115 +8,240 @@ import org.apache.lucene.document.SortedNumericDocValuesField; import org.apache.lucene.document.StoredField; +import org.apache.lucene.index.IndexableField; import org.apache.lucene.index.LeafReaderContext; -import org.apache.lucene.index.RandomIndexWriter; import org.apache.lucene.util.BytesRef; import org.apache.lucene.util.NumericUtils; -import org.elasticsearch.common.CheckedConsumer; +import org.elasticsearch.common.CheckedFunction; +import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.index.fielddata.ScriptDocValues; +import org.elasticsearch.index.fielddata.SortedNumericDoubleValues; +import org.elasticsearch.index.mapper.MappedFieldType; import org.elasticsearch.index.mapper.NumberFieldMapper.NumberFieldType; import org.elasticsearch.index.mapper.NumberFieldMapper.NumberType; +import org.elasticsearch.plugins.ScriptPlugin; import org.elasticsearch.script.ScriptContext; +import org.elasticsearch.script.ScriptEngine; import org.elasticsearch.search.lookup.SearchLookup; +import org.elasticsearch.xpack.runtimefields.DoubleScriptFieldScript.Factory; import java.io.IOException; +import java.util.Collection; import java.util.List; import java.util.Map; +import java.util.Set; +import java.util.function.BiConsumer; +import java.util.function.Consumer; +import java.util.function.DoubleConsumer; import static org.hamcrest.Matchers.equalTo; public class DoubleScriptFieldScriptTests extends ScriptFieldScriptTestCase< - DoubleScriptFieldScript, DoubleScriptFieldScript.Factory, - DoubleScriptFieldScript.LeafFactory, + DoubleRuntimeFieldHelper, + SortedNumericDoubleValues, Double> { public void testConstant() throws IOException { - CheckedConsumer indexBuilder = iw -> { - iw.addDocument(List.of(new SortedNumericDocValuesField("foo", NumericUtils.doubleToSortableLong(randomDouble())))); - iw.addDocument(List.of(new SortedNumericDocValuesField("foo", NumericUtils.doubleToSortableLong(randomDouble())))); - }; - assertThat(execute(indexBuilder, "value(3.14)"), equalTo(List.of(3.14, 3.14))); + assertThat(randomDoubles().collect("value(3.14)"), equalTo(List.of(3.14, 3.14))); } public void testTwoConstants() throws IOException { - CheckedConsumer indexBuilder = iw -> { - iw.addDocument(List.of(new SortedNumericDocValuesField("foo", randomLong()))); - iw.addDocument(List.of(new SortedNumericDocValuesField("foo", randomLong()))); - }; - assertThat(execute(indexBuilder, "value(3.14); value(2.72)"), equalTo(List.of(3.14, 2.72, 3.14, 2.72))); + assertThat(randomDoubles().collect("value(3.14); value(2.72)"), equalTo(List.of(2.72, 3.14, 2.72, 3.14))); } public void testSource() throws IOException { - CheckedConsumer indexBuilder = iw -> { - iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": 1}")))); - iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": 10}")))); - }; - assertThat(execute(indexBuilder, "value(source['foo'] * 10.1)"), equalTo(List.of(10.1, 101.0))); + assertThat(singleValueInSource().collect("value(source['foo'] * 10.1)"), equalTo(List.of(10.1, 101.0))); } - public void testTwoSourceFields() throws IOException { - CheckedConsumer indexBuilder = iw -> { - iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": 1, \"bar\": 2}")))); - iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": 10, \"bar\": 20}")))); - }; + public void testTwoSourceValues() throws IOException { assertThat( - execute(indexBuilder, "value(source['foo'] * 10.1); value(source['bar'] * 10.2)"), + multiValueInSource().collect("value(source['foo'][0] * 10.1); value(source['foo'][1] * 10.2)"), equalTo(List.of(10.1, 20.4, 101.0, 204.0)) ); } public void testDocValues() throws IOException { - CheckedConsumer indexBuilder = iw -> { - iw.addDocument(List.of(new SortedNumericDocValuesField("foo", NumericUtils.doubleToSortableLong(1.1)))); - iw.addDocument(List.of(new SortedNumericDocValuesField("foo", NumericUtils.doubleToSortableLong(10.1)))); - }; - assertThat( - execute(indexBuilder, "value(doc['foo'].value * 9.9)", new NumberFieldType("foo", NumberType.DOUBLE)), - equalTo(List.of(10.89, 99.99)) - ); + assertThat(singleValueInDocValues().collect("value(doc['foo'].value * 9.9)"), equalTo(List.of(10.89, 99.99))); } - public void testTwoDocValuesValues() throws IOException { - CheckedConsumer indexBuilder = iw -> { - iw.addDocument( - List.of( - new SortedNumericDocValuesField("foo", NumericUtils.doubleToSortableLong(1.1)), - new SortedNumericDocValuesField("foo", NumericUtils.doubleToSortableLong(2.2)) - ) - ); - iw.addDocument( - List.of( - new SortedNumericDocValuesField("foo", NumericUtils.doubleToSortableLong(10.1)), - new SortedNumericDocValuesField("foo", NumericUtils.doubleToSortableLong(20.1)) - ) - ); - }; + public void testMultipleDocValuesValues() throws IOException { assertThat( - execute(indexBuilder, "for (double d : doc['foo']) {value(d * 9.9)}", new NumberFieldType("foo", NumberType.DOUBLE)), + multipleValuesInDocValues().collect("for (double d : doc['foo']) {value(d * 9.9)}"), equalTo(List.of(10.89, 21.78, 99.99, 198.99)) ); } + public void testExistsQuery() throws IOException { + TestCase c = multipleValuesInDocValues(); + DoubleRuntimeFieldHelper isOnePointOne = c.testScript("is_one_point_one"); + assertThat(c.collect(isOnePointOne.existsQuery("foo"), isOnePointOne), equalTo(List.of(1.1))); + } + + public void testTermQuery() throws IOException { + TestCase c = multipleValuesInDocValues(); + DoubleRuntimeFieldHelper timesTen = c.testScript("times_nine_point_nine"); + assertThat(c.collect(timesTen.termQuery("foo", 1), timesTen), equalTo(List.of())); + assertThat(c.collect(timesTen.termQuery("foo", 10.89), timesTen), equalTo(List.of(10.89, 21.78))); + assertThat(c.collect(timesTen.termQuery("foo", 21.78), timesTen), equalTo(List.of(10.89, 21.78))); + assertThat(c.collect(timesTen.termQuery("foo", 99.99), timesTen), equalTo(List.of(99.99, 198.99))); + } + + public void testTermsQuery() throws IOException { + TestCase c = multipleValuesInDocValues(); + DoubleRuntimeFieldHelper timesTen = c.testScript("times_nine_point_nine"); + assertThat(c.collect(timesTen.termsQuery("foo", 1, 2), timesTen), equalTo(List.of())); + assertThat(c.collect(timesTen.termsQuery("foo", 10.89, 11), timesTen), equalTo(List.of(10.89, 21.78))); + assertThat(c.collect(timesTen.termsQuery("foo", 21.78, 22), timesTen), equalTo(List.of(10.89, 21.78))); + assertThat(c.collect(timesTen.termsQuery("foo", 20, 21.78), timesTen), equalTo(List.of(10.89, 21.78))); + assertThat(c.collect(timesTen.termsQuery("foo", 99.99, 100), timesTen), equalTo(List.of(99.99, 198.99))); + } + + public void testRangeQuery() throws IOException { + TestCase c = multipleValuesInDocValues(); + DoubleRuntimeFieldHelper timesTen = c.testScript("times_nine_point_nine"); + assertThat(c.collect(timesTen.rangeQuery("foo", 1, 2), timesTen), equalTo(List.of())); + assertThat(c.collect(timesTen.rangeQuery("foo", 9, 11), timesTen), equalTo(List.of(10.89, 21.78))); + assertThat(c.collect(timesTen.rangeQuery("foo", 10.89, 11), timesTen), equalTo(List.of(10.89, 21.78))); + assertThat(c.collect(timesTen.rangeQuery("foo", 21.78, 22), timesTen), equalTo(List.of(10.89, 21.78))); + assertThat(c.collect(timesTen.rangeQuery("foo", 21, 21.78), timesTen), equalTo(List.of(10.89, 21.78))); + assertThat(c.collect(timesTen.rangeQuery("foo", 99, 100), timesTen), equalTo(List.of(99.99, 198.99))); + } + + private TestCase randomDoubles() throws IOException { + return testCase(iw -> { + iw.addDocument(List.of(doubleDocValue(randomDouble()))); + iw.addDocument(List.of(doubleDocValue(randomDouble()))); + }); + } + + private TestCase singleValueInSource() throws IOException { + return testCase(iw -> { + iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": 1}")))); + iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": 10}")))); + }); + } + + private TestCase multiValueInSource() throws IOException { + return testCase(iw -> { + iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [1, 2]}")))); + iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [10, 20]}")))); + }); + } + + private TestCase singleValueInDocValues() throws IOException { + return testCase(iw -> { + iw.addDocument(List.of(doubleDocValue(1.1))); + iw.addDocument(List.of(doubleDocValue(10.1))); + }); + } + + private TestCase multipleValuesInDocValues() throws IOException { + return testCase(iw -> { + iw.addDocument(List.of(doubleDocValue(1.1), doubleDocValue(2.2))); + iw.addDocument(List.of(doubleDocValue(10.1), doubleDocValue(20.1))); + }); + } + + private IndexableField doubleDocValue(double value) { + return new SortedNumericDocValuesField("foo", NumericUtils.doubleToSortableLong(value)); + } + + @Override + protected MappedFieldType[] fieldTypes() { + return new MappedFieldType[] { new NumberFieldType("foo", NumberType.DOUBLE) }; + } + @Override protected ScriptContext scriptContext() { return DoubleScriptFieldScript.CONTEXT; } @Override - protected DoubleScriptFieldScript.LeafFactory newLeafFactory( - DoubleScriptFieldScript.Factory factory, - Map params, - SearchLookup searchLookup - ) { - return factory.newFactory(params, searchLookup); + protected DoubleRuntimeFieldHelper newHelper(Factory factory, Map params, SearchLookup searchLookup) + throws IOException { + return factory.newFactory(params, searchLookup).runtimeFieldHelper(); } @Override - protected DoubleScriptFieldScript newInstance( - DoubleScriptFieldScript.LeafFactory leafFactory, - LeafReaderContext context, - List result - ) throws IOException { - return leafFactory.newInstance(context, result::add); + protected CheckedFunction docValuesBuilder(DoubleRuntimeFieldHelper values) { + return values.docValues(); + } + + @Override + protected void readAllDocValues(SortedNumericDoubleValues docValues, int docId, Consumer sync) throws IOException { + assertTrue(docValues.advanceExact(docId)); + int count = docValues.docValueCount(); + for (int i = 0; i < count; i++) { + sync.accept(docValues.nextValue()); + } + } + + @Override + protected List extraScriptPlugins() { + return List.of(new ScriptPlugin() { + @Override + public ScriptEngine getScriptEngine(Settings settings, Collection> contexts) { + return new ScriptEngine() { + @Override + public String getType() { + return "test"; + } + + @Override + public Set> getSupportedContexts() { + return Set.of(DoubleScriptFieldScript.CONTEXT); + } + + @Override + public FactoryType compile( + String name, + String code, + ScriptContext context, + Map params + ) { + assert context == DoubleScriptFieldScript.CONTEXT; + @SuppressWarnings("unchecked") + FactoryType result = (FactoryType) compile(name); + return result; + } + + private DoubleScriptFieldScript.Factory compile(String name) { + if (name.equals("times_nine_point_nine")) { + return assertingScript((fieldData, sync) -> { + for (Object v : fieldData.get("foo")) { + sync.accept(((double) v) * 9.9); + } + }); + } + if (name.equals("is_one_point_one")) { + return assertingScript((fieldData, sync) -> { + for (Object v : fieldData.get("foo")) { + double d = (double) v; + if (d == 1.1) { + sync.accept(1.1); + } + } + }); + } + throw new IllegalArgumentException(); + } + }; + } + }); + } + + private DoubleScriptFieldScript.Factory assertingScript(BiConsumer>, DoubleConsumer> impl) { + return (params, searchLookup) -> { + DoubleScriptFieldScript.LeafFactory leafFactory = (ctx, sync) -> { + return new DoubleScriptFieldScript(params, searchLookup, ctx, sync) { + @Override + public void execute() { + impl.accept(getDoc(), sync); + } + }; + }; + return leafFactory; + }; } } diff --git a/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/LongScriptFieldScriptTests.java b/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/LongScriptFieldScriptTests.java index ef0c5cd304d0e..4fa4472c37c9f 100644 --- a/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/LongScriptFieldScriptTests.java +++ b/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/LongScriptFieldScriptTests.java @@ -9,78 +9,161 @@ import org.apache.lucene.document.SortedNumericDocValuesField; import org.apache.lucene.document.StoredField; import org.apache.lucene.index.LeafReaderContext; -import org.apache.lucene.index.RandomIndexWriter; +import org.apache.lucene.index.SortedNumericDocValues; +import org.apache.lucene.search.BooleanClause.Occur; +import org.apache.lucene.search.BooleanQuery; import org.apache.lucene.util.BytesRef; -import org.elasticsearch.common.CheckedConsumer; +import org.elasticsearch.common.CheckedFunction; +import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.index.fielddata.ScriptDocValues; +import org.elasticsearch.index.mapper.MappedFieldType; import org.elasticsearch.index.mapper.NumberFieldMapper.NumberFieldType; import org.elasticsearch.index.mapper.NumberFieldMapper.NumberType; +import org.elasticsearch.plugins.ScriptPlugin; import org.elasticsearch.script.ScriptContext; +import org.elasticsearch.script.ScriptEngine; import org.elasticsearch.search.lookup.SearchLookup; +import org.elasticsearch.xpack.runtimefields.LongScriptFieldScript.Factory; import java.io.IOException; +import java.util.Collection; import java.util.List; import java.util.Map; +import java.util.Set; +import java.util.function.BiConsumer; +import java.util.function.Consumer; +import java.util.function.LongConsumer; import static org.hamcrest.Matchers.equalTo; public class LongScriptFieldScriptTests extends ScriptFieldScriptTestCase< - LongScriptFieldScript, LongScriptFieldScript.Factory, - LongScriptFieldScript.LeafFactory, + LongRuntimeFieldHelper, + SortedNumericDocValues, Long> { public void testConstant() throws IOException { - CheckedConsumer indexBuilder = iw -> { - iw.addDocument(List.of(new SortedNumericDocValuesField("foo", randomLong()))); - iw.addDocument(List.of(new SortedNumericDocValuesField("foo", randomLong()))); - }; - assertThat(execute(indexBuilder, "value(10)"), equalTo(List.of(10L, 10L))); + assertThat(randomLongs().collect("value(10)"), equalTo(List.of(10L, 10L))); } public void testTwoConstants() throws IOException { - CheckedConsumer indexBuilder = iw -> { + assertThat(randomLongs().collect("value(10); value(20)"), equalTo(List.of(10L, 20L, 10L, 20L))); + } + + public void testSource() throws IOException { + assertThat(singleValueInSource().collect("value(source['foo'] * 10)"), equalTo(List.of(10L, 100L))); + } + + public void testMultileSourceValues() throws IOException { + assertThat(multiValueInSource().collect("for (long l : source['foo']){value(l * 10)}"), equalTo(List.of(10L, 20L, 100L, 200L))); + } + + public void testDocValues() throws IOException { + assertThat(singleValueInDocValues().collect("value(doc['foo'].value * 10)"), equalTo(List.of(10L, 100L))); + } + + public void testMultipleDocValuesValues() throws IOException { + TestCase c = multipleValuesInDocValues(); + assertThat(multipleValuesInDocValues().collect(c.testScript("times_ten")), equalTo(List.of(10L, 20L, 100L, 200L))); + } + + public void testExistsQuery() throws IOException { + TestCase c = multipleValuesInDocValues(); + LongRuntimeFieldHelper isOne = c.testScript("is_one"); + assertThat(c.collect(isOne.existsQuery("foo"), isOne), equalTo(List.of(1L))); + } + + public void testTermQuery() throws IOException { + TestCase c = multipleValuesInDocValues(); + LongRuntimeFieldHelper timesTen = c.testScript("times_ten"); + assertThat(c.collect(timesTen.termQuery("foo", 1), timesTen), equalTo(List.of())); + assertThat(c.collect(timesTen.termQuery("foo", 10), timesTen), equalTo(List.of(10L, 20L))); + assertThat(c.collect(timesTen.termQuery("foo", 20), timesTen), equalTo(List.of(10L, 20L))); + assertThat(c.collect(timesTen.termQuery("foo", 100), timesTen), equalTo(List.of(100L, 200L))); + } + + public void testTermsQuery() throws IOException { + TestCase c = multipleValuesInDocValues(); + LongRuntimeFieldHelper timesTen = c.testScript("times_ten"); + assertThat(c.collect(timesTen.termsQuery("foo", 1, 2), timesTen), equalTo(List.of())); + assertThat(c.collect(timesTen.termsQuery("foo", 10, 11), timesTen), equalTo(List.of(10L, 20L))); + assertThat(c.collect(timesTen.termsQuery("foo", 20, 21), timesTen), equalTo(List.of(10L, 20L))); + assertThat(c.collect(timesTen.termsQuery("foo", 19, 20), timesTen), equalTo(List.of(10L, 20L))); + assertThat(c.collect(timesTen.termsQuery("foo", 100, 11), timesTen), equalTo(List.of(100L, 200L))); + } + + public void testRangeQuery() throws IOException { + TestCase c = multipleValuesInDocValues(); + LongRuntimeFieldHelper timesTen = c.testScript("times_ten"); + assertThat(c.collect(timesTen.rangeQuery("foo", 1, 2), timesTen), equalTo(List.of())); + assertThat(c.collect(timesTen.rangeQuery("foo", 9, 11), timesTen), equalTo(List.of(10L, 20L))); + assertThat(c.collect(timesTen.rangeQuery("foo", 10, 11), timesTen), equalTo(List.of(10L, 20L))); + assertThat(c.collect(timesTen.rangeQuery("foo", 19, 21), timesTen), equalTo(List.of(10L, 20L))); + assertThat(c.collect(timesTen.rangeQuery("foo", 99, 101), timesTen), equalTo(List.of(100L, 200L))); + } + + public void testInsideBoolTermQuery() throws IOException { + /* + * Its required that bool queries that contain more our runtime fields queries + * be wrapped in ForceNoBulkScoringQuery. Exactly what queries in the tree need + * to be wrapped and when isn't super clear but it is safest to wrap the whole + * query tree when there are *any* of these queries in it. We might be able to + * skip some of them eventually, when we're more comfortable with this. + */ + TestCase c = multipleValuesInDocValues(); + LongRuntimeFieldHelper timesTen = c.testScript("times_ten"); + assertThat( + c.collect( + new ForceNoBulkScoringQuery( + new BooleanQuery.Builder().add(timesTen.termQuery("foo", 1), Occur.SHOULD) + .add(timesTen.termQuery("foo", 10), Occur.SHOULD) + .add(timesTen.termQuery("foo", 100), Occur.SHOULD) + .build() + ), + timesTen + ), + equalTo(List.of(10L, 20L, 100L, 200L)) + ); + } + + private TestCase randomLongs() throws IOException { + return testCase(iw -> { iw.addDocument(List.of(new SortedNumericDocValuesField("foo", randomLong()))); iw.addDocument(List.of(new SortedNumericDocValuesField("foo", randomLong()))); - }; - assertThat(execute(indexBuilder, "value(10); value(20)"), equalTo(List.of(10L, 20L, 10L, 20L))); + }); } - public void testSource() throws IOException { - CheckedConsumer indexBuilder = iw -> { + private TestCase singleValueInSource() throws IOException { + return testCase(iw -> { iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": 1}")))); iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": 10}")))); - }; - assertThat(execute(indexBuilder, "value(source['foo'] * 10)"), equalTo(List.of(10L, 100L))); + }); } - public void testTwoSourceFields() throws IOException { - CheckedConsumer indexBuilder = iw -> { - iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": 1, \"bar\": 2}")))); - iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": 10, \"bar\": 20}")))); - }; - assertThat(execute(indexBuilder, "value(source['foo'] * 10); value(source['bar'] * 10)"), equalTo(List.of(10L, 20L, 100L, 200L))); + private TestCase multiValueInSource() throws IOException { + return testCase(iw -> { + iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [1, 2]}")))); + iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [10, 20]}")))); + }); } - public void testDocValues() throws IOException { - CheckedConsumer indexBuilder = iw -> { + private TestCase singleValueInDocValues() throws IOException { + return testCase(iw -> { iw.addDocument(List.of(new SortedNumericDocValuesField("foo", 1))); iw.addDocument(List.of(new SortedNumericDocValuesField("foo", 10))); - }; - assertThat( - execute(indexBuilder, "value(doc['foo'].value * 10)", new NumberFieldType("foo", NumberType.LONG)), - equalTo(List.of(10L, 100L)) - ); + }); } - public void testTwoDocValuesValues() throws IOException { - CheckedConsumer indexBuilder = iw -> { + private TestCase multipleValuesInDocValues() throws IOException { + return testCase(iw -> { iw.addDocument(List.of(new SortedNumericDocValuesField("foo", 1), new SortedNumericDocValuesField("foo", 2))); iw.addDocument(List.of(new SortedNumericDocValuesField("foo", 10), new SortedNumericDocValuesField("foo", 20))); - }; - assertThat( - execute(indexBuilder, "for (long l : doc['foo']) {value(l * 10)}", new NumberFieldType("foo", NumberType.LONG)), - equalTo(List.of(10L, 20L, 100L, 200L)) - ); + }); + } + + @Override + protected MappedFieldType[] fieldTypes() { + return new MappedFieldType[] { new NumberFieldType("foo", NumberType.LONG) }; } @Override @@ -89,18 +172,89 @@ protected ScriptContext scriptContext() { } @Override - protected LongScriptFieldScript.LeafFactory newLeafFactory( - LongScriptFieldScript.Factory factory, - Map params, - SearchLookup searchLookup - ) { - return factory.newFactory(params, searchLookup); + protected LongRuntimeFieldHelper newHelper(Factory factory, Map params, SearchLookup searchLookup) throws IOException { + return factory.newFactory(params, searchLookup).runtimeFieldHelper(); } @Override - protected LongScriptFieldScript newInstance(LongScriptFieldScript.LeafFactory leafFactory, LeafReaderContext context, List result) - throws IOException { + protected CheckedFunction docValuesBuilder(LongRuntimeFieldHelper values) { + return values.docValues(); + } - return leafFactory.newInstance(context, result::add); + @Override + protected void readAllDocValues(SortedNumericDocValues docValues, int docId, Consumer sync) throws IOException { + assertTrue(docValues.advanceExact(docId)); + int count = docValues.docValueCount(); + for (int i = 0; i < count; i++) { + sync.accept(docValues.nextValue()); + } + } + + @Override + protected List extraScriptPlugins() { + return List.of(new ScriptPlugin() { + @Override + public ScriptEngine getScriptEngine(Settings settings, Collection> contexts) { + return new ScriptEngine() { + @Override + public String getType() { + return "test"; + } + + @Override + public Set> getSupportedContexts() { + return Set.of(LongScriptFieldScript.CONTEXT); + } + + @Override + public FactoryType compile( + String name, + String code, + ScriptContext context, + Map params + ) { + assert context == LongScriptFieldScript.CONTEXT; + @SuppressWarnings("unchecked") + FactoryType result = (FactoryType) compile(name); + return result; + } + + private LongScriptFieldScript.Factory compile(String name) { + if (name.equals("times_ten")) { + return assertingScript((fieldData, sync) -> { + for (Object v : fieldData.get("foo")) { + sync.accept(((long) v) * 10); + } + }); + } + if (name.equals("is_one")) { + return assertingScript((fieldData, sync) -> { + for (Object v : fieldData.get("foo")) { + long l = (long) v; + if (l == 1) { + sync.accept(1); + } + } + }); + } + throw new IllegalArgumentException(); + } + }; + } + }); + } + + private LongScriptFieldScript.Factory assertingScript(BiConsumer>, LongConsumer> impl) { + return (params, searchLookup) -> { + LongScriptFieldScript.LeafFactory leafFactory = (ctx, sync) -> { + return new LongScriptFieldScript(params, searchLookup, ctx, sync) { + @Override + public void execute() { + impl.accept(getDoc(), sync); + } + }; + }; + return leafFactory; + }; } } diff --git a/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/ScriptFieldScriptTestCase.java b/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/ScriptFieldScriptTestCase.java index 3b7a819d04b40..5bc869a36690e 100644 --- a/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/ScriptFieldScriptTestCase.java +++ b/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/ScriptFieldScriptTestCase.java @@ -13,13 +13,16 @@ import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.LeafCollector; import org.apache.lucene.search.MatchAllDocsQuery; +import org.apache.lucene.search.Query; import org.apache.lucene.search.Scorable; import org.apache.lucene.search.ScoreMode; import org.apache.lucene.store.Directory; import org.elasticsearch.Version; import org.elasticsearch.cluster.metadata.IndexMetadata; import org.elasticsearch.common.CheckedConsumer; +import org.elasticsearch.common.CheckedFunction; import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.core.internal.io.IOUtils; import org.elasticsearch.index.IndexSettings; import org.elasticsearch.index.fielddata.IndexFieldData; import org.elasticsearch.index.mapper.MappedFieldType; @@ -27,35 +30,33 @@ import org.elasticsearch.indices.breaker.NoneCircuitBreakerService; import org.elasticsearch.painless.PainlessPlugin; import org.elasticsearch.plugins.ExtensiblePlugin.ExtensionLoader; +import org.elasticsearch.plugins.ScriptPlugin; import org.elasticsearch.script.Script; import org.elasticsearch.script.ScriptContext; import org.elasticsearch.script.ScriptModule; import org.elasticsearch.script.ScriptService; +import org.elasticsearch.script.ScriptType; import org.elasticsearch.search.lookup.SearchLookup; import org.elasticsearch.test.ESTestCase; +import org.junit.After; +import java.io.Closeable; import java.io.IOException; -import java.security.AccessController; -import java.security.PrivilegedAction; import java.util.ArrayList; -import java.util.HashMap; +import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.function.Consumer; import java.util.function.Function; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; -public abstract class ScriptFieldScriptTestCase extends ESTestCase { - protected abstract ScriptContext scriptContext(); - - protected abstract LF newLeafFactory(F factory, Map params, SearchLookup searchLookup); - - protected abstract S newInstance(LF leafFactory, LeafReaderContext context, List results) throws IOException; - - protected final List execute(CheckedConsumer indexBuilder, String script, MappedFieldType... types) - throws IOException { +public abstract class ScriptFieldScriptTestCase extends ESTestCase { + private final List lazyClose = new ArrayList<>(); + private final ScriptService scriptService; + public ScriptFieldScriptTestCase() { PainlessPlugin painlessPlugin = new PainlessPlugin(); painlessPlugin.loadExtensions(new ExtensionLoader() { @Override @@ -64,50 +65,99 @@ public List loadExtensions(Class extensionPointType) { return (List) List.of(new RuntimeFieldsPainlessExtension()); } }); - ScriptModule scriptModule = new ScriptModule(Settings.EMPTY, List.of(painlessPlugin, new RuntimeFields())); - Map params = new HashMap<>(); - MapperService mapperService = mock(MapperService.class); - for (MappedFieldType type : types) { - when(mapperService.fieldType(type.name())).thenReturn(type); - } - Function> fieldDataLookup = ft -> ft.fielddataBuilder("test") - .build(indexSettings(), ft, null, new NoneCircuitBreakerService(), mapperService); - SearchLookup searchLookup = new SearchLookup(mapperService, fieldDataLookup); - try (ScriptService scriptService = new ScriptService(Settings.EMPTY, scriptModule.engines, scriptModule.contexts)) { - F factory = AccessController.doPrivileged( - (PrivilegedAction) () -> scriptService.compile(new Script(script), scriptContext()) - ); - - try (Directory directory = newDirectory(); RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directory)) { + List scriptPlugins = new ArrayList<>(); + scriptPlugins.add(painlessPlugin); // TODO move painless to integration tests + scriptPlugins.add(new RuntimeFields()); + scriptPlugins.addAll(extraScriptPlugins()); + ScriptModule scriptModule = new ScriptModule(Settings.EMPTY, scriptPlugins); + scriptService = new ScriptService(Settings.EMPTY, scriptModule.engines, scriptModule.contexts); + } + + protected List extraScriptPlugins() { + return List.of(); + } + + protected abstract MappedFieldType[] fieldTypes(); + + protected abstract ScriptContext scriptContext(); + + protected abstract H newHelper(F factory, Map params, SearchLookup searchLookup) throws IOException; + + protected abstract CheckedFunction docValuesBuilder(H values); + + protected abstract void readAllDocValues(DV docValues, int docId, Consumer sync) throws IOException; + + protected final TestCase testCase(CheckedConsumer indexBuilder) throws IOException { + return new TestCase(indexBuilder); + } + + protected class TestCase { + private final MapperService mapperService = mock(MapperService.class); + private final SearchLookup searchLookup; + private final IndexSearcher searcher; + + private TestCase(CheckedConsumer indexBuilder) throws IOException { + for (MappedFieldType type : fieldTypes()) { + when(mapperService.fieldType(type.name())).thenReturn(type); + } + Function> fieldDataLookup = ft -> ft.fielddataBuilder("test") + .build(indexSettings(), ft, null, new NoneCircuitBreakerService(), mapperService); + searchLookup = new SearchLookup(mapperService, fieldDataLookup); + Directory directory = newDirectory(); + lazyClose.add(directory); + try (RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directory)) { indexBuilder.accept(indexWriter); - try (DirectoryReader reader = indexWriter.getReader()) { - IndexSearcher searcher = newSearcher(reader); - LF leafFactory = newLeafFactory(factory, params, searchLookup); - List result = new ArrayList<>(); - searcher.search(new MatchAllDocsQuery(), new Collector() { + DirectoryReader reader = indexWriter.getReader(); + lazyClose.add(reader); + searcher = newSearcher(reader); + } + } + + protected H script(String script) throws IOException { + return script(new Script(script)); + } + + protected H testScript(String name) throws IOException { + return script(new Script(ScriptType.INLINE, "test", name, Map.of())); + } + + private H script(Script script) throws IOException { + return newHelper(scriptService.compile(script, scriptContext()), Map.of(), searchLookup); + } + + protected List collect(String script) throws IOException { + return collect(new MatchAllDocsQuery(), script(script)); + } + + protected List collect(H values) throws IOException { + return collect(new MatchAllDocsQuery(), values); + } + + protected List collect(Query query, H values) throws IOException { + // Now run the query and collect the results + List result = new ArrayList<>(); + CheckedFunction docValuesBuilder = docValuesBuilder(values); + searcher.search(query, new Collector() { + @Override + public ScoreMode scoreMode() { + return ScoreMode.COMPLETE_NO_SCORES; + } + + @Override + public LeafCollector getLeafCollector(LeafReaderContext ctx) throws IOException { + DV docValues = docValuesBuilder.apply(ctx); + return new LeafCollector() { @Override - public ScoreMode scoreMode() { - return ScoreMode.COMPLETE_NO_SCORES; - } + public void setScorer(Scorable scorer) throws IOException {} @Override - public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException { - S compiled = newInstance(leafFactory, context, result); - return new LeafCollector() { - @Override - public void setScorer(Scorable scorer) {} - - @Override - public void collect(int doc) { - compiled.setDocument(doc); - compiled.execute(); - } - }; + public void collect(int docId) throws IOException { + readAllDocValues(docValues, docId, result::add); } - }); - return result; + }; } - } + }); + return result; } } @@ -122,4 +172,10 @@ private IndexSettings indexSettings() { Settings.EMPTY ); } + + @After + public void closeAll() throws IOException { + Collections.reverse(lazyClose); // Close in the reverse order added so readers close before directory + IOUtils.close(lazyClose); + } } diff --git a/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/StringScriptFieldScriptTests.java b/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/StringScriptFieldScriptTests.java index 5ecedd1762647..6caaf2ae9b3f4 100644 --- a/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/StringScriptFieldScriptTests.java +++ b/x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/StringScriptFieldScriptTests.java @@ -8,85 +8,253 @@ import org.apache.lucene.document.SortedSetDocValuesField; import org.apache.lucene.document.StoredField; +import org.apache.lucene.index.IndexableField; import org.apache.lucene.index.LeafReaderContext; -import org.apache.lucene.index.RandomIndexWriter; import org.apache.lucene.util.BytesRef; -import org.elasticsearch.common.CheckedConsumer; +import org.apache.lucene.util.automaton.RegExp; +import org.elasticsearch.common.CheckedFunction; +import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.index.fielddata.ScriptDocValues; +import org.elasticsearch.index.fielddata.SortedBinaryDocValues; import org.elasticsearch.index.mapper.KeywordFieldMapper.KeywordFieldType; +import org.elasticsearch.index.mapper.MappedFieldType; +import org.elasticsearch.plugins.ScriptPlugin; import org.elasticsearch.script.ScriptContext; +import org.elasticsearch.script.ScriptEngine; import org.elasticsearch.search.lookup.SearchLookup; +import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript.Factory; import java.io.IOException; +import java.util.ArrayList; +import java.util.Collection; import java.util.List; import java.util.Map; +import java.util.Set; +import java.util.function.BiConsumer; +import java.util.function.Consumer; import static org.hamcrest.Matchers.equalTo; public class StringScriptFieldScriptTests extends ScriptFieldScriptTestCase< - StringScriptFieldScript, StringScriptFieldScript.Factory, - StringScriptFieldScript.LeafFactory, + StringRuntimeFieldHelper, + SortedBinaryDocValues, String> { public void testConstant() throws IOException { - CheckedConsumer indexBuilder = iw -> { - iw.addDocument(List.of(new SortedSetDocValuesField("foo", new BytesRef(randomAlphaOfLength(2))))); - iw.addDocument(List.of(new SortedSetDocValuesField("foo", new BytesRef(randomAlphaOfLength(2))))); - }; - assertThat(execute(indexBuilder, "value('cat')"), equalTo(List.of("cat", "cat"))); + assertThat(randomStrings().collect("value('cat')"), equalTo(List.of("cat", "cat"))); } public void testTwoConstants() throws IOException { - CheckedConsumer indexBuilder = iw -> { - iw.addDocument(List.of(new SortedSetDocValuesField("foo", new BytesRef(randomAlphaOfLength(2))))); - iw.addDocument(List.of(new SortedSetDocValuesField("foo", new BytesRef(randomAlphaOfLength(2))))); - }; - assertThat(execute(indexBuilder, "value('cat'); value('dog')"), equalTo(List.of("cat", "dog", "cat", "dog"))); + assertThat(randomStrings().collect("value('cat'); value('dog')"), equalTo(List.of("cat", "dog", "cat", "dog"))); } public void testSource() throws IOException { - CheckedConsumer indexBuilder = iw -> { - iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": \"cat\"}")))); - iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": \"dog\"}")))); - }; - assertThat(execute(indexBuilder, "value(source['foo'] + 'o')"), equalTo(List.of("cato", "dogo"))); + assertThat(singleValueInSource().collect("value(source['foo'] + 'o')"), equalTo(List.of("cato", "dogo"))); } - public void testTwoSourceFields() throws IOException { - CheckedConsumer indexBuilder = iw -> { - iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": \"cat\", \"bar\": \"chicken\"}")))); - iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": \"dog\", \"bar\": \"pig\"}")))); - }; + public void testMultipleSourceValues() throws IOException { assertThat( - execute(indexBuilder, "value(source['foo'] + 'o'); value(source['bar'] + 'ie')"), - equalTo(List.of("cato", "chickenie", "dogo", "pigie")) + multiValueInSource().collect("for (String v : source['foo']) {value(v + 'o')}"), + equalTo(List.of("cato", "chickeno", "dogo", "pigo")) ); } public void testDocValues() throws IOException { - CheckedConsumer indexBuilder = iw -> { + assertThat(singleValueInDocValues().collect(ADD_O), equalTo(List.of("cato", "dogo"))); + } + + public void testMultipleDocValuesValues() throws IOException { + assertThat(multipleValuesInDocValues().collect(ADD_O), equalTo(List.of("cato", "pigo", "chickeno", "dogo"))); + } + + public void testExistsQuery() throws IOException { + TestCase c = multipleValuesInDocValues(); + StringRuntimeFieldHelper isCat = c.testScript("is_cat"); + assertThat(c.collect(isCat.existsQuery("foo"), isCat), equalTo(List.of("cat"))); + } + + public void testFuzzyQuery() throws IOException { + TestCase c = multipleValuesInDocValues(); + StringRuntimeFieldHelper addO = c.testScript("add_o"); + assertThat(c.collect(addO.fuzzyQuery("foo", "caaaaat", 1, 1, 1, true), addO), equalTo(List.of())); + assertThat(c.collect(addO.fuzzyQuery("foo", "cat", 1, 1, 1, true), addO), equalTo(List.of("cato", "pigo"))); + assertThat(c.collect(addO.fuzzyQuery("foo", "pig", 1, 1, 1, true), addO), equalTo(List.of("cato", "pigo"))); + assertThat(c.collect(addO.fuzzyQuery("foo", "dog", 1, 1, 1, true), addO), equalTo(List.of("chickeno", "dogo"))); + } + + public void testTermQuery() throws IOException { + TestCase c = multipleValuesInDocValues(); + StringRuntimeFieldHelper addO = c.testScript("add_o"); + assertThat(c.collect(addO.termQuery("foo", "cat"), addO), equalTo(List.of())); + assertThat(c.collect(addO.termQuery("foo", "cato"), addO), equalTo(List.of("cato", "pigo"))); + assertThat(c.collect(addO.termQuery("foo", "pigo"), addO), equalTo(List.of("cato", "pigo"))); + assertThat(c.collect(addO.termQuery("foo", "dogo"), addO), equalTo(List.of("chickeno", "dogo"))); + } + + public void testTermsQuery() throws IOException { + TestCase c = multipleValuesInDocValues(); + StringRuntimeFieldHelper addO = c.testScript("add_o"); + assertThat(c.collect(addO.termsQuery("foo", "cat", "dog"), addO), equalTo(List.of())); + assertThat(c.collect(addO.termsQuery("foo", "cato", "piglet"), addO), equalTo(List.of("cato", "pigo"))); + assertThat(c.collect(addO.termsQuery("foo", "pigo", "catington"), addO), equalTo(List.of("cato", "pigo"))); + assertThat(c.collect(addO.termsQuery("foo", "dogo", "lightbulb"), addO), equalTo(List.of("chickeno", "dogo"))); + } + + public void testPrefixQuery() throws IOException { + TestCase c = multipleValuesInDocValues(); + StringRuntimeFieldHelper addO = c.testScript("add_o"); + assertThat(c.collect(addO.prefixQuery("foo", "catdog"), addO), equalTo(List.of())); + assertThat(c.collect(addO.prefixQuery("foo", "cat"), addO), equalTo(List.of("cato", "pigo"))); + assertThat(c.collect(addO.prefixQuery("foo", "pig"), addO), equalTo(List.of("cato", "pigo"))); + assertThat(c.collect(addO.prefixQuery("foo", "dogo"), addO), equalTo(List.of("chickeno", "dogo"))); + assertThat(c.collect(addO.prefixQuery("foo", "d"), addO), equalTo(List.of("chickeno", "dogo"))); + } + + public void testRangeQuery() throws IOException { + TestCase c = multipleValuesInDocValues(); + StringRuntimeFieldHelper addO = c.testScript("add_o"); + assertThat(c.collect(addO.rangeQuery("foo", "catz", "cbat", false, false), addO), equalTo(List.of())); + assertThat(c.collect(addO.rangeQuery("foo", "c", "cb", false, false), addO), equalTo(List.of("cato", "pigo"))); + assertThat(c.collect(addO.rangeQuery("foo", "p", "q", false, false), addO), equalTo(List.of("cato", "pigo"))); + assertThat(c.collect(addO.rangeQuery("foo", "doggie", "dogs", false, false), addO), equalTo(List.of("chickeno", "dogo"))); + assertThat(c.collect(addO.rangeQuery("foo", "dogo", "dogs", false, false), addO), equalTo(List.of())); + assertThat(c.collect(addO.rangeQuery("foo", "dogo", "dogs", true, false), addO), equalTo(List.of("chickeno", "dogo"))); + assertThat(c.collect(addO.rangeQuery("foo", "dog", "dogo", false, false), addO), equalTo(List.of())); + assertThat(c.collect(addO.rangeQuery("foo", "dog", "dogo", false, true), addO), equalTo(List.of("chickeno", "dogo"))); + } + + public void testRegexpQuery() throws IOException { + TestCase c = multipleValuesInDocValues(); + StringRuntimeFieldHelper addO = c.testScript("add_o"); + assertThat(c.collect(addO.regexpQuery("foo", "cat", RegExp.ALL, 100000), addO), equalTo(List.of())); + assertThat(c.collect(addO.regexpQuery("foo", "cat[aeiou]", RegExp.ALL, 100000), addO), equalTo(List.of("cato", "pigo"))); + assertThat(c.collect(addO.regexpQuery("foo", "p.*", RegExp.ALL, 100000), addO), equalTo(List.of("cato", "pigo"))); + assertThat(c.collect(addO.regexpQuery("foo", "dog?o", RegExp.ALL, 100000), addO), equalTo(List.of("chickeno", "dogo"))); + } + + public void testWildcardQuery() throws IOException { + TestCase c = multipleValuesInDocValues(); + StringRuntimeFieldHelper addO = c.testScript("add_o"); + assertThat(c.collect(addO.wildcardQuery("foo", "cat"), addO), equalTo(List.of())); + assertThat(c.collect(addO.wildcardQuery("foo", "cat?"), addO), equalTo(List.of("cato", "pigo"))); + assertThat(c.collect(addO.wildcardQuery("foo", "p*"), addO), equalTo(List.of("cato", "pigo"))); + assertThat(c.collect(addO.wildcardQuery("foo", "do?o"), addO), equalTo(List.of("chickeno", "dogo"))); + } + + private TestCase randomStrings() throws IOException { + return testCase(iw -> { + iw.addDocument(List.of(new SortedSetDocValuesField("foo", new BytesRef(randomAlphaOfLength(2))))); + iw.addDocument(List.of(new SortedSetDocValuesField("foo", new BytesRef(randomAlphaOfLength(2))))); + }); + } + + private TestCase singleValueInSource() throws IOException { + return testCase(iw -> { + iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": \"cat\"}")))); + iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": \"dog\"}")))); + }); + } + + private TestCase multiValueInSource() throws IOException { + return testCase(iw -> { + iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [\"cat\", \"chicken\"]}")))); + iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [\"dog\", \"pig\"]}")))); + }); + } + + private TestCase singleValueInDocValues() throws IOException { + return testCase(iw -> { iw.addDocument(List.of(new SortedSetDocValuesField("foo", new BytesRef("cat")))); iw.addDocument(List.of(new SortedSetDocValuesField("foo", new BytesRef("dog")))); + }); + } + + private TestCase multipleValuesInDocValues() throws IOException { + return testCase(iw -> { + List doc = new ArrayList<>(); + doc.add(new SortedSetDocValuesField("foo", new BytesRef("cat"))); + doc.add(new SortedSetDocValuesField("foo", new BytesRef("pig"))); + iw.addDocument(doc); + doc.clear(); + doc.add(new SortedSetDocValuesField("foo", new BytesRef("chicken"))); + doc.add(new SortedSetDocValuesField("foo", new BytesRef("dog"))); + iw.addDocument(doc); + }); + } + + private static final String ADD_O = "for (String s: doc['foo']) {value(s + 'o')}"; + + @Override + protected List extraScriptPlugins() { + return List.of(new ScriptPlugin() { + @Override + public ScriptEngine getScriptEngine(Settings settings, Collection> contexts) { + return new ScriptEngine() { + @Override + public String getType() { + return "test"; + } + + @Override + public Set> getSupportedContexts() { + return Set.of(StringScriptFieldScript.CONTEXT); + } + + @Override + public FactoryType compile( + String name, + String code, + ScriptContext context, + Map params + ) { + assert context == StringScriptFieldScript.CONTEXT; + @SuppressWarnings("unchecked") + FactoryType result = (FactoryType) compile(name); + return result; + } + + private StringScriptFieldScript.Factory compile(String name) { + if (name.equals("add_o")) { + return assertingScript((fieldData, sync) -> { + for (Object v : fieldData.get("foo")) { + sync.accept(v + "o"); + } + }); + } + if (name.equals("is_cat")) { + return assertingScript((fieldData, sync) -> { + for (Object v : fieldData.get("foo")) { + if (v.equals("cat")) { + sync.accept("cat"); + } + } + }); + } + throw new IllegalArgumentException(); + } + }; + } + }); + } + + private StringScriptFieldScript.Factory assertingScript(BiConsumer>, Consumer> impl) { + return (params, searchLookup) -> { + StringScriptFieldScript.LeafFactory leafFactory = (ctx, sync) -> { + return new StringScriptFieldScript(params, searchLookup, ctx, sync) { + @Override + public void execute() { + impl.accept(getDoc(), sync); + } + }; + }; + return leafFactory; }; - assertThat(execute(indexBuilder, "value(doc['foo'].value + 'o')", new KeywordFieldType("foo")), equalTo(List.of("cato", "dogo"))); - } - - public void testTwoDocValuesValues() throws IOException { - CheckedConsumer indexBuilder = iw -> { - iw.addDocument( - List.of( - new SortedSetDocValuesField("foo", new BytesRef("cat")), - new SortedSetDocValuesField("foo", new BytesRef("chicken")) - ) - ); - iw.addDocument( - List.of(new SortedSetDocValuesField("foo", new BytesRef("dog")), new SortedSetDocValuesField("foo", new BytesRef("pig"))) - ); - }; - assertThat( - execute(indexBuilder, "for (String s: doc['foo']) {value(s + 'o')}", new KeywordFieldType("foo")), - equalTo(List.of("cato", "chickeno", "dogo", "pigo")) - ); + } + + @Override + protected MappedFieldType[] fieldTypes() { + return new MappedFieldType[] { new KeywordFieldType("foo") }; } @Override @@ -95,20 +263,22 @@ protected ScriptContext scriptContext() { } @Override - protected StringScriptFieldScript.LeafFactory newLeafFactory( - StringScriptFieldScript.Factory factory, - Map params, - SearchLookup searchLookup - ) { - return factory.newFactory(params, searchLookup); + protected StringRuntimeFieldHelper newHelper(Factory factory, Map params, SearchLookup searchLookup) + throws IOException { + return factory.newFactory(params, searchLookup).runtimeFieldHelper(); + } + + @Override + protected CheckedFunction docValuesBuilder(StringRuntimeFieldHelper values) { + return values.docValues(); } @Override - protected StringScriptFieldScript newInstance( - StringScriptFieldScript.LeafFactory leafFactory, - LeafReaderContext context, - List result - ) throws IOException { - return leafFactory.newInstance(context, result::add); + protected void readAllDocValues(SortedBinaryDocValues docValues, int docId, Consumer sync) throws IOException { + assertTrue(docValues.advanceExact(docId)); + int count = docValues.docValueCount(); + for (int i = 0; i < count; i++) { + sync.accept(docValues.nextValue().utf8ToString()); + } } } diff --git a/x-pack/plugin/src/test/resources/rest-api-spec/test/runtime_fields/10_keyword_script.yml b/x-pack/plugin/src/test/resources/rest-api-spec/test/runtime_fields/10_keyword_script.yml new file mode 100644 index 0000000000000..e0c5ef0db1640 --- /dev/null +++ b/x-pack/plugin/src/test/resources/rest-api-spec/test/runtime_fields/10_keyword_script.yml @@ -0,0 +1,190 @@ +--- +setup: + - do: + indices.create: + index: sensor + body: + settings: + number_of_shards: 1 + number_of_replicas: 0 + mappings: + properties: + timestamp: + type: date + temperature: + type: long + voltage: + type: float + node: + type: keyword + day_of_week: + type: script + runtime_type: keyword + script: value(doc['timestamp'].value.dayOfWeekEnum.getDisplayName(TextStyle.FULL, Locale.ROOT)) + days_starting_with_s: + type: script + runtime_type: keyword + script: | + String dayOfWeek = doc['day_of_week'].value; + if (dayOfWeek.startsWith('S')) { + value(dayOfWeek); + } + + - do: + bulk: + index: sensor + refresh: true + body: | + {"index":{}} + {"timestamp": 1516729294000, "temperature": 200, "voltage": 5.2, "node": "a"} + {"index":{}} + {"timestamp": 1516642894000, "temperature": 201, "voltage": 5.8, "node": "b"} + {"index":{}} + {"timestamp": 1516556494000, "temperature": 202, "voltage": 5.1, "node": "a"} + {"index":{}} + {"timestamp": 1516470094000, "temperature": 198, "voltage": 5.6, "node": "b"} + {"index":{}} + {"timestamp": 1516383694000, "temperature": 200, "voltage": 4.2, "node": "c"} + {"index":{}} + {"timestamp": 1516297294000, "temperature": 202, "voltage": 4.0, "node": "c"} + +--- +"fetch": + - do: + search: + index: sensor + body: + docvalue_fields: + - day_of_week + - days_starting_with_s + sort: timestamp + - match: {hits.total.value: 6} + - match: {hits.hits.0.fields.day_of_week: [Thursday]} + - match: {hits.hits.1.fields.day_of_week: [Friday]} + - match: {hits.hits.2.fields.day_of_week: [Saturday]} + - match: {hits.hits.3.fields.day_of_week: [Sunday]} + - match: {hits.hits.4.fields.day_of_week: [Monday]} + - match: {hits.hits.5.fields.day_of_week: [Tuesday]} + - is_false: hits.hits.0.fields.days_starting_with_s + - is_false: hits.hits.1.fields.days_starting_with_s + - match: {hits.hits.2.fields.days_starting_with_s: [Saturday]} + - match: {hits.hits.3.fields.days_starting_with_s: [Sunday]} + - is_false: hits.hits.4.fields.days_starting_with_s + - is_false: hits.hits.5.fields.days_starting_with_s + +--- +"exists query": + - do: + search: + index: sensor + body: + query: + exists: + field: days_starting_with_s + - match: {hits.total.value: 2} + +--- +"fuzzy query": + - do: + search: + index: sensor + body: + query: + fuzzy: + day_of_week: Mondai + - match: {hits.total.value: 1} + - match: {hits.hits.0._source.voltage: 5.8} + +--- +"match query": + - do: + search: + index: sensor + body: + query: + match: + day_of_week: Monday + - match: {hits.total.value: 1} + - match: {hits.hits.0._source.voltage: 5.8} + +--- +"prefix query": + - do: + search: + index: sensor + body: + query: + prefix: + day_of_week: Mon + - match: {hits.total.value: 1} + - match: {hits.hits.0._source.voltage: 5.8} + +--- +"regexp query": + - do: + search: + index: sensor + body: + query: + regexp: + day_of_week: + value: Mon.+ + sort: timestamp + - match: {hits.total.value: 1} + - match: {hits.hits.0._source.voltage: 5.8} + +--- +"term query": + - do: + search: + index: sensor + body: + query: + term: + day_of_week: Monday + - match: {hits.total.value: 1} + - match: {hits.hits.0._source.voltage: 5.8} + +--- +"terms query": + - do: + search: + index: sensor + body: + query: + terms: + day_of_week: [Monday, Tuesday] + sort: timestamp + - match: {hits.total.value: 2} + - match: {hits.hits.0._source.voltage: 5.8} + - match: {hits.hits.1._source.voltage: 5.2} + +--- +"wildcard query": + - do: + search: + index: sensor + body: + query: + wildcard: + day_of_week: Mond?y + - match: {hits.total.value: 1} + - match: {hits.hits.0._source.voltage: 5.8} + +--- +"in a bool query": + - do: + search: + index: sensor + body: + query: + bool: + should: + - term: + day_of_week: Monday + - term: + day_of_week: Tuesday + sort: timestamp + - match: {hits.total.value: 2} + - match: {hits.hits.0._source.voltage: 5.8} + - match: {hits.hits.1._source.voltage: 5.2}