Two queries for keyword script field

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

@@ -17,11 +17,17 @@
 import org.elasticsearch.script.Script;
 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.StringScriptFieldTermQuery;
+import org.elasticsearch.xpack.runtimefields.query.StringScriptFieldTermsQuery;
 
 import java.io.IOException;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
+
+import static java.util.stream.Collectors.toSet;
 
 public final class RuntimeKeywordMappedFieldType extends MappedFieldType {
 
@@ -57,18 +63,24 @@ public ScriptBinaryFieldData.Builder fielddataBuilder(String fullyQualifiedIndex
         return new ScriptBinaryFieldData.Builder(scriptFactory);
     }
 
+    private StringScriptFieldScript.LeafFactory leafFactory(QueryShardContext context) {
+        return scriptFactory.newFactory(script.getParams(), context.lookup());
+    }
+
+    @Override
+    public Query existsQuery(QueryShardContext context) {
+        return new StringScriptFieldExistsQuery(leafFactory(context), name());
+    }
+
     @Override
     public Query termQuery(Object value, QueryShardContext context) {
-        return new StringScriptFieldTermQuery(
-            scriptFactory.newFactory(script.getParams(), context.lookup()),
-            name(),
-            BytesRefs.toString(Objects.requireNonNull(value))
-        );
+        return new StringScriptFieldTermQuery(leafFactory(context), name(), BytesRefs.toString(Objects.requireNonNull(value)));
     }
 
     @Override
-    public Query existsQuery(QueryShardContext context) {
-        return null;
+    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);
     }
 
     void doXContentBody(XContentBuilder builder, boolean includeDefaults, Params params) throws IOException {

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

@@ -0,0 +1,91 @@
+/*
+ * 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.LeafReaderContext;
+import org.apache.lucene.search.ConstantScoreScorer;
+import org.apache.lucene.search.ConstantScoreWeight;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.search.Scorer;
+import org.apache.lucene.search.TwoPhaseIterator;
+import org.apache.lucene.search.Weight;
+import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript;
+import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript.LeafFactory;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Objects;
+
+/**
+ * Abstract base class for building queries based on {@link StringScriptFieldScript}.
+ */
+abstract class AbstractStringScriptFieldQuery extends Query {
+    private final StringScriptFieldScript.LeafFactory leafFactory;
+    private final String fieldName;
+
+    AbstractStringScriptFieldQuery(LeafFactory leafFactory, String fieldName) {
+        this.leafFactory = Objects.requireNonNull(leafFactory);
+        this.fieldName = Objects.requireNonNull(fieldName);
+    }
+
+    /**
+     * Does the value match this query?
+     */
+    public abstract boolean matches(List<String> values);
+
+    @Override
+    public final Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
+        return new ConstantScoreWeight(this, boost) {
+            @Override
+            public boolean isCacheable(LeafReaderContext ctx) {
+                return false; // scripts aren't really cacheable at this point
+            }
+
+            @Override
+            public Scorer scorer(LeafReaderContext ctx) throws IOException {
+                StringScriptFieldScript script = leafFactory.newInstance(ctx);
+                DocIdSetIterator approximation = DocIdSetIterator.all(ctx.reader().maxDoc());
+                TwoPhaseIterator twoPhase = new TwoPhaseIterator(approximation) {
+                    @Override
+                    public boolean matches() throws IOException {
+                        return AbstractStringScriptFieldQuery.this.matches(script.resultsForDoc(approximation().docID()));
+                    }
+
+                    @Override
+                    public float matchCost() {
+                        // TODO we don't have a good way of estimating the complexity of the script so we just go with 9000
+                        return 9000f;
+                    }
+                };
+                return new ConstantScoreScorer(this, score(), scoreMode, twoPhase);
+            }
+        };
+    }
+
+    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);
+    }
+
+    @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);
+    }
+}

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

@@ -0,0 +1,32 @@
+/*
+ * 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.elasticsearch.xpack.runtimefields.StringScriptFieldScript;
+
+import java.util.List;
+
+public class StringScriptFieldExistsQuery extends AbstractStringScriptFieldQuery {
+    public StringScriptFieldExistsQuery(StringScriptFieldScript.LeafFactory leafFactory, String fieldName) {
+        super(leafFactory, fieldName);
+    }
+
+    @Override
+    public boolean matches(List<String> values) {
+        return false == values.isEmpty();
+    }
+
+    @Override
+    public final String toString(String field) {
+        if (fieldName().contentEquals(field)) {
+            return "*";
+        }
+        return fieldName() + ":*";
+    }
+
+    // Superclass's equals and hashCode are great for this class
+}

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

@@ -6,92 +6,59 @@
 
 package org.elasticsearch.xpack.runtimefields.query;
 
-import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.Term;
-import org.apache.lucene.search.ConstantScoreScorer;
-import org.apache.lucene.search.ConstantScoreWeight;
-import org.apache.lucene.search.DocIdSetIterator;
-import org.apache.lucene.search.IndexSearcher;
-import org.apache.lucene.search.Query;
 import org.apache.lucene.search.QueryVisitor;
-import org.apache.lucene.search.ScoreMode;
-import org.apache.lucene.search.Scorer;
-import org.apache.lucene.search.TwoPhaseIterator;
-import org.apache.lucene.search.Weight;
 import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript;
 
-import java.io.IOException;
+import java.util.List;
 import java.util.Objects;
 
-public class StringScriptFieldTermQuery extends Query {
-    private final StringScriptFieldScript.LeafFactory leafFactory;
-    private final String fieldName;
+public class StringScriptFieldTermQuery extends AbstractStringScriptFieldQuery {
     private final String term;
 
     public StringScriptFieldTermQuery(StringScriptFieldScript.LeafFactory leafFactory, String fieldName, String term) {
-        this.leafFactory = leafFactory;
-        this.fieldName = fieldName;
-        this.term = term;
+        super(leafFactory, fieldName);
+        this.term = Objects.requireNonNull(term);
     }
 
     @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
+    public boolean matches(List<String> values) {
+        for (String value : values) {
+            if (term.equals(value)) {
+                return true;
             }
-
-            @Override
-            public Scorer scorer(LeafReaderContext ctx) throws IOException {
-                StringScriptFieldScript script = leafFactory.newInstance(ctx);
-                DocIdSetIterator approximation = DocIdSetIterator.all(ctx.reader().maxDoc());
-                TwoPhaseIterator twoPhase = new TwoPhaseIterator(approximation) {
-                    @Override
-                    public boolean matches() throws IOException {
-                        for (String result : script.resultsForDoc(approximation().docID())) {
-                            if (term.equals(result)) {
-                                return true;
-                            }
-                        }
-                        return false;
-                    }
-
-                    @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 9000f;
-                    }
-                };
-                return new ConstantScoreScorer(this, score(), scoreMode, twoPhase);
-            }
-        };
+        }
+        return false;
     }
 
     @Override
     public void visit(QueryVisitor visitor) {
-        visitor.consumeTerms(this, new Term(fieldName, term));
+        visitor.consumeTerms(this, new Term(fieldName(), term));
     }
 
     @Override
     public final String toString(String field) {
-        if (fieldName.contentEquals(field)) {
+        if (fieldName().contentEquals(field)) {
             return term;
         }
-        return fieldName + ":" + term;
+        return fieldName() + ":" + term;
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(fieldName, term);
+        return Objects.hash(super.hashCode(), term);
     }
 
     @Override
     public boolean equals(Object obj) {
-        if (obj == null || getClass() != obj.getClass()) {
+        if (false == super.equals(obj)) {
             return false;
         }
         StringScriptFieldTermQuery other = (StringScriptFieldTermQuery) obj;
-        return fieldName.equals(other.fieldName) && term.equals(other.term);
+        return term.equals(other.term);
+    }
+
+    String term() {
+        return term;
     }
 }

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/StringScriptFieldTermsQuery.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.index.Term;
+import org.apache.lucene.search.QueryVisitor;
+import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript;
+
+import java.util.List;
+import java.util.Objects;
+import java.util.Set;
+
+public class StringScriptFieldTermsQuery extends AbstractStringScriptFieldQuery {
+    private final Set<String> terms;
+
+    public StringScriptFieldTermsQuery(StringScriptFieldScript.LeafFactory leafFactory, String fieldName, Set<String> terms) {
+        super(leafFactory, fieldName);
+        this.terms = terms;
+    }
+
+    @Override
+    public boolean matches(List<String> values) {
+        for (String value : values) {
+            if (terms.contains(value)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    @Override
+    public void visit(QueryVisitor visitor) {
+        if (visitor.acceptField(fieldName())) {
+            for (String term : terms) {
+                visitor.consumeTerms(this, new Term(fieldName(), term));
+            }
+        }
+    }
+
+    @Override
+    public final String toString(String field) {
+        if (fieldName().contentEquals(field)) {
+            return terms.toString();
+        }
+        return fieldName() + ":" + terms;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(super.hashCode(), terms);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (false == super.equals(obj)) {
+            return false;
+        }
+        StringScriptFieldTermsQuery other = (StringScriptFieldTermsQuery) obj;
+        return terms.equals(other.terms);
+    }
+
+    Set<String> terms() {
+        return terms;
+    }
+}

x-pack/plugin/runtime-fields/src/test/java/org/elasticsearch/xpack/runtimefields/mapper/RuntimeKeywordMappedFieldTypeTests.java

@@ -91,6 +91,17 @@ public void collect(int doc) throws IOException {
         }
     }
 
+    public void testExistsQuery() throws IOException {
+        try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
+            iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": [1]}"))));
+            iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": []}"))));
+            try (DirectoryReader reader = iw.getReader()) {
+                IndexSearcher searcher = newSearcher(reader);
+                assertThat(searcher.count(build("for (def v : source.foo) { value(v.toString())}").existsQuery(mockContext())), equalTo(1));
+            }
+        }
+    }
+
     public void testTermQuery() throws IOException {
         try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
             iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": 1}"))));
@@ -102,6 +113,19 @@ public void testTermQuery() throws IOException {
         }
     }
 
+    public void testTermsQuery() throws IOException {
+        try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
+            iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": 1}"))));
+            iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": 2}"))));
+            iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": 3}"))));
+            iw.addDocument(List.of(new StoredField("_source", new BytesRef("{\"foo\": 4}"))));
+            try (DirectoryReader reader = iw.getReader()) {
+                IndexSearcher searcher = newSearcher(reader);
+                assertThat(searcher.count(build("value(source.foo.toString())").termsQuery(List.of("1", "2"), mockContext())), equalTo(2));
+            }
+        }
+    }
+
     private RuntimeKeywordMappedFieldType build(String code) throws IOException {
         Script script = new Script(code);
         PainlessPlugin painlessPlugin = new PainlessPlugin();