Remaining queries for script keyword fields

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/mapper/RuntimeKeywordMappedFieldType.java

@@ -6,9 +6,13 @@
 
 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.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.mapper.MappedFieldType;
@@ -18,10 +22,16 @@
 import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript;
 import org.elasticsearch.xpack.runtimefields.fielddata.ScriptBinaryFieldData;
 import org.elasticsearch.xpack.runtimefields.query.StringScriptFieldExistsQuery;
+import org.elasticsearch.xpack.runtimefields.query.StringScriptFieldFuzzyQuery;
+import org.elasticsearch.xpack.runtimefields.query.StringScriptFieldPrefixQuery;
+import org.elasticsearch.xpack.runtimefields.query.StringScriptFieldRangeQuery;
+import org.elasticsearch.xpack.runtimefields.query.StringScriptFieldRegexpQuery;
 import org.elasticsearch.xpack.runtimefields.query.StringScriptFieldTermQuery;
 import org.elasticsearch.xpack.runtimefields.query.StringScriptFieldTermsQuery;
+import org.elasticsearch.xpack.runtimefields.query.StringScriptFieldWildcardQuery;
 
 import java.io.IOException;
+import java.time.ZoneId;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -69,18 +79,75 @@ private StringScriptFieldScript.LeafFactory leafFactory(QueryShardContext contex
 
  @Override
  public Query existsQuery(QueryShardContext context) {
- return new StringScriptFieldExistsQuery(leafFactory(context), name());
+ return new StringScriptFieldExistsQuery(script, leafFactory(context), name());
+ }
+
+ @Override
+ public Query fuzzyQuery(
+ Object value,
+ Fuzziness fuzziness,
+ int prefixLength,
+ int maxExpansions,
+ boolean transpositions,
+ QueryShardContext context
+ ) {
+ return StringScriptFieldFuzzyQuery.build(
+ script,
+ leafFactory(context),
+ name(),
+ BytesRefs.toString(Objects.requireNonNull(value)),
+ fuzziness.asDistance(BytesRefs.toString(value)),
+ prefixLength,
+ transpositions
+ );
+ }
+
+ @Override
+ public Query prefixQuery(String value, RewriteMethod method, org.elasticsearch.index.query.QueryShardContext context) {
+ return new StringScriptFieldPrefixQuery(script, leafFactory(context), name(), value);
+ }
+
+ @Override
+ public Query rangeQuery(
+ Object lowerTerm,
+ Object upperTerm,
+ boolean includeLower,
+ boolean includeUpper,
+ ShapeRelation relation,
+ ZoneId timeZone,
+ DateMathParser parser,
+ QueryShardContext context
+ ) {
+ return new StringScriptFieldRangeQuery(
+ script,
+ leafFactory(context),
+ name(),
+ BytesRefs.toString(Objects.requireNonNull(lowerTerm)),
+ BytesRefs.toString(Objects.requireNonNull(upperTerm)),
+ includeLower,
+ includeUpper
+ );
+ }
+
+ @Override
+ public Query regexpQuery(String value, int flags, int maxDeterminizedStates, RewriteMethod method, QueryShardContext context) {
+ return new StringScriptFieldRegexpQuery(script, leafFactory(context), name(), value, flags, maxDeterminizedStates);
  }
 
  @Override
  public Query termQuery(Object value, QueryShardContext context) {
- return new StringScriptFieldTermQuery(leafFactory(context), name(), BytesRefs.toString(Objects.requireNonNull(value)));
+ return new StringScriptFieldTermQuery(script, leafFactory(context), name(), BytesRefs.toString(Objects.requireNonNull(value)));
  }
 
  @Override
  public Query termsQuery(List<?> values, QueryShardContext context) {
  Set<String> terms = values.stream().map(v -> BytesRefs.toString(Objects.requireNonNull(v))).collect(toSet());
- return new StringScriptFieldTermsQuery(leafFactory(context), name(), terms);
+ return new StringScriptFieldTermsQuery(script, leafFactory(context), name(), terms);
+ }
+
+ @Override
+ public Query wildcardQuery(String value, RewriteMethod method, QueryShardContext context) {
+ return new StringScriptFieldWildcardQuery(script, leafFactory(context), name(), value);
  }
 
  void doXContentBody(XContentBuilder builder, boolean includeDefaults, Params params) throws IOException {

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractStringScriptFieldAutomatonQuery.java

@@ -0,0 +1,48 @@
+/*
+ * 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.query;
+
+import org.apache.lucene.search.QueryVisitor;
+import org.apache.lucene.util.BytesRefBuilder;
+import org.apache.lucene.util.automaton.ByteRunAutomaton;
+import org.elasticsearch.script.Script;
+import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript;
+
+import java.util.List;
+
+public abstract class AbstractStringScriptFieldAutomatonQuery extends AbstractStringScriptFieldQuery {
+ private final BytesRefBuilder scratch = new BytesRefBuilder();
+ private final ByteRunAutomaton automaton;
+
+ public AbstractStringScriptFieldAutomatonQuery(
+ Script script,
+ StringScriptFieldScript.LeafFactory leafFactory,
+ String fieldName,
+ ByteRunAutomaton automaton
+ ) {
+ super(script, leafFactory, fieldName);
+ this.automaton = automaton;
+ }
+
+ @Override
+ protected final boolean matches(List<String> values) {
+ for (String value : values) {
+ scratch.copyChars(value);
+ 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);
+ }
+ }
+}

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/AbstractStringScriptFieldQuery.java

@@ -16,6 +16,7 @@
 import org.apache.lucene.search.Scorer;
 import org.apache.lucene.search.TwoPhaseIterator;
 import org.apache.lucene.search.Weight;
+import org.elasticsearch.script.Script;
 import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript;
 import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript.LeafFactory;
 
@@ -27,18 +28,20 @@
  * Abstract base class for building queries based on {@link StringScriptFieldScript}.
  */
 abstract class AbstractStringScriptFieldQuery extends Query {
+ private final Script script;
  private final StringScriptFieldScript.LeafFactory leafFactory;
  private final String fieldName;
 
- AbstractStringScriptFieldQuery(LeafFactory leafFactory, String fieldName) {
+ AbstractStringScriptFieldQuery(Script script, LeafFactory leafFactory, String fieldName) {
+ this.script = script;
  this.leafFactory = Objects.requireNonNull(leafFactory);
  this.fieldName = Objects.requireNonNull(fieldName);
  }
 
  /**
  * Does the value match this query?
  */
- public abstract boolean matches(List<String> values);
+ protected abstract boolean matches(List<String> values);
 
  @Override
  public final Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
@@ -69,23 +72,25 @@ public float matchCost() {
  };
  }
 
+ final Script script() {
+ return script;
+ }
+
  protected final String fieldName() {
  return fieldName;
  }
 
  @Override
  public int hashCode() {
- // TODO should leafFactory be here? Something about the script probably should be!
- return Objects.hash(getClass(), fieldName);
+ return Objects.hash(getClass(), script, fieldName);
  }
 
  @Override
  public boolean equals(Object obj) {
  if (obj == null || getClass() != obj.getClass()) {
  return false;
  }
- // TODO should leafFactory be here? Something about the script probably should be!
  AbstractStringScriptFieldQuery other = (AbstractStringScriptFieldQuery) obj;
- return fieldName.equals(other.fieldName);
+ return script.equals(other.script) && fieldName.equals(other.fieldName);
  }
 }

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/StringScriptFieldExistsQuery.java

@@ -6,26 +6,27 @@
 
 package org.elasticsearch.xpack.runtimefields.query;
 
+import org.elasticsearch.script.Script;
 import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript;
 
 import java.util.List;
 
 public class StringScriptFieldExistsQuery extends AbstractStringScriptFieldQuery {
- public StringScriptFieldExistsQuery(StringScriptFieldScript.LeafFactory leafFactory, String fieldName) {
- super(leafFactory, fieldName);
+ public StringScriptFieldExistsQuery(Script script, StringScriptFieldScript.LeafFactory leafFactory, String fieldName) {
+ super(script, leafFactory, fieldName);
  }
 
  @Override
- public boolean matches(List<String> values) {
+ protected boolean matches(List<String> values) {
  return false == values.isEmpty();
  }
 
  @Override
  public final String toString(String field) {
  if (fieldName().contentEquals(field)) {
- return "*";
+ return "ScriptFieldExists";
  }
- return fieldName() + ":*";
+ return fieldName() + ":ScriptFieldExists";
  }
 
  // Superclass's equals and hashCode are great for this class

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/StringScriptFieldFuzzyQuery.java

@@ -0,0 +1,68 @@
+/*
+ * 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.query;
+
+import org.apache.lucene.index.Term;
+import org.apache.lucene.search.FuzzyQuery;
+import org.apache.lucene.util.automaton.ByteRunAutomaton;
+import org.elasticsearch.script.Script;
+import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript;
+
+import java.util.Objects;
+
+public class StringScriptFieldFuzzyQuery extends AbstractStringScriptFieldAutomatonQuery {
+ public static StringScriptFieldFuzzyQuery build(
+ Script script,
+ StringScriptFieldScript.LeafFactory leafFactory,
+ String fieldName,
+ String term,
+ int maxEdits,
+ int prefixLength,
+ boolean transpositions
+ ) {
+ int maxExpansions = 1; // We don't actually expand anything so the value here doesn't matter
+ FuzzyQuery delegate = new FuzzyQuery(new Term(fieldName, term), maxEdits, prefixLength, maxExpansions, transpositions);
+ ByteRunAutomaton automaton = delegate.getAutomata().runAutomaton;
+ return new StringScriptFieldFuzzyQuery(script, leafFactory, fieldName, automaton, delegate);
+ }
+
+ private final FuzzyQuery delegate;
+
+ private StringScriptFieldFuzzyQuery(
+ Script script,
+ StringScriptFieldScript.LeafFactory leafFactory,
+ String fieldName,
+ ByteRunAutomaton automaton,
+ FuzzyQuery delegate
+ ) {
+ super(script, leafFactory, fieldName, automaton);
+ this.delegate = delegate;
+ }
+
+ @Override
+ public final String toString(String field) {
+ return delegate.toString(field);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(super.hashCode(), delegate);
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (false == super.equals(obj)) {
+ return false;
+ }
+ StringScriptFieldFuzzyQuery other = (StringScriptFieldFuzzyQuery) obj;
+ return delegate.equals(other.delegate);
+ }
+
+ FuzzyQuery delegate() {
+ return delegate;
+ }
+}

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/StringScriptFieldPrefixQuery.java

@@ -0,0 +1,69 @@
+/*
+ * 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.query;
+
+import org.apache.lucene.search.PrefixQuery;
+import org.apache.lucene.search.QueryVisitor;
+import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.automaton.ByteRunAutomaton;
+import org.elasticsearch.script.Script;
+import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript;
+
+import java.util.List;
+import java.util.Objects;
+
+public class StringScriptFieldPrefixQuery extends AbstractStringScriptFieldQuery {
+ private final String prefix;
+
+ public StringScriptFieldPrefixQuery(Script script, StringScriptFieldScript.LeafFactory leafFactory, String fieldName, String prefix) {
+ super(script, leafFactory, fieldName);
+ this.prefix = Objects.requireNonNull(prefix);
+ }
+
+ @Override
+ protected boolean matches(List<String> values) {
+ for (String value : values) {
+ if (value != null && value.startsWith(prefix)) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ @Override
+ public void visit(QueryVisitor visitor) {
+ if (visitor.acceptField(fieldName())) {
+ visitor.consumeTermsMatching(this, fieldName(), () -> new ByteRunAutomaton(PrefixQuery.toAutomaton(new BytesRef(prefix))));
+ }
+ }
+
+ @Override
+ public final String toString(String field) {
+ if (fieldName().contentEquals(field)) {
+ return prefix + "*";
+ }
+ return fieldName() + ":" + prefix + "*";
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(super.hashCode(), prefix);
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (false == super.equals(obj)) {
+ return false;
+ }
+ StringScriptFieldPrefixQuery other = (StringScriptFieldPrefixQuery) obj;
+ return prefix.equals(other.prefix);
+ }
+
+ String prefix() {
+ return prefix;
+ }
+}