SCRIPTING: Move terms_set Context to its Own Class

server/src/main/java/org/elasticsearch/index/query/TermsSetQueryBuilder.java

@@ -40,14 +40,14 @@
 import org.elasticsearch.index.fielddata.IndexNumericFieldData;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.script.Script;
-import org.elasticsearch.script.SearchScript;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import org.elasticsearch.script.TermsSetQueryScript;
 
 public final class TermsSetQueryBuilder extends AbstractQueryBuilder<TermsSetQueryBuilder> {
 
@@ -262,13 +262,12 @@ private LongValuesSource createValuesSource(QueryShardContext context) {
             IndexNumericFieldData fieldData = context.getForField(msmFieldType);
             longValuesSource = new FieldValuesSource(fieldData);
         } else if (minimumShouldMatchScript != null) {
-            SearchScript.Factory factory = context.getScriptService().compile(minimumShouldMatchScript,
-                SearchScript.TERMS_SET_QUERY_CONTEXT);
+            TermsSetQueryScript.Factory factory = context.getScriptService().compile(minimumShouldMatchScript,
+                TermsSetQueryScript.CONTEXT);
             Map<String, Object> params = new HashMap<>();
             params.putAll(minimumShouldMatchScript.getParams());
             params.put("num_terms", values.size());
-            SearchScript.LeafFactory leafFactory = factory.newFactory(params, context.lookup());
-            longValuesSource = new ScriptLongValueSource(minimumShouldMatchScript, leafFactory);
+            longValuesSource = new ScriptLongValueSource(minimumShouldMatchScript, factory.newFactory(params, context.lookup()));
         } else {
             throw new IllegalStateException("No minimum should match has been specified");
         }
@@ -278,26 +277,26 @@ private LongValuesSource createValuesSource(QueryShardContext context) {
     static final class ScriptLongValueSource extends LongValuesSource {
 
         private final Script script;
-        private final SearchScript.LeafFactory leafFactory;
+        private final TermsSetQueryScript.LeafFactory leafFactory;
 
-        ScriptLongValueSource(Script script, SearchScript.LeafFactory leafFactory) {
+        ScriptLongValueSource(Script script, TermsSetQueryScript.LeafFactory leafFactory) {
             this.script = script;
             this.leafFactory = leafFactory;
         }
 
         @Override
         public LongValues getValues(LeafReaderContext ctx, DoubleValues scores) throws IOException {
-            SearchScript searchScript = leafFactory.newInstance(ctx);
+            TermsSetQueryScript script = leafFactory.newInstance(ctx);
             return new LongValues() {
                 @Override
                 public long longValue() throws IOException {
-                    return searchScript.runAsLong();
+                    return script.runAsLong();
                 }
 
                 @Override
                 public boolean advanceExact(int doc) throws IOException {
-                    searchScript.setDocument(doc);
-                    return searchScript.run() != null;
+                    script.setDocument(doc);
+                    return script.execute() != null;
                 }
             };
         }

server/src/main/java/org/elasticsearch/script/ParameterMap.java

@@ -0,0 +1,105 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.script;
+
+import java.util.Collection;
+import java.util.Map;
+import java.util.Set;
+import org.apache.logging.log4j.LogManager;
+import org.elasticsearch.common.logging.DeprecationLogger;
+
+public final class ParameterMap implements Map<String, Object> {
+
+    private static final DeprecationLogger DEPRECATION_LOGGER =
+        new DeprecationLogger(LogManager.getLogger(ParameterMap.class));
+
+    private final Map<String, Object> params;
+
+    private final Map<String, String> deprecations;
+
+    ParameterMap(Map<String, Object> params, Map<String, String> deprecations) {
+        this.params = params;
+        this.deprecations = deprecations;
+    }
+
+    @Override
+    public int size() {
+        return params.size();
+    }
+
+    @Override
+    public boolean isEmpty() {
+        return params.isEmpty();
+    }
+
+    @Override
+    public boolean containsKey(final Object key) {
+        return params.containsKey(key);
+    }
+
+    @Override
+    public boolean containsValue(final Object value) {
+        return params.containsValue(value);
+    }
+
+    @Override
+    public Object get(final Object key) {
+        String deprecationMessage = deprecations.get(key);
+        if (deprecationMessage != null) {
+            DEPRECATION_LOGGER.deprecated(deprecationMessage);
+        }
+        return params.get(key);
+    }
+
+    @Override
+    public Object put(final String key, final Object value) {
+        return params.put(key, value);
+    }
+
+    @Override
+    public Object remove(final Object key) {
+        return params.remove(key);
+    }
+
+    @Override
+    public void putAll(final Map<? extends String, ?> m) {
+        params.putAll(m);
+    }
+
+    @Override
+    public void clear() {
+        params.clear();
+    }
+
+    @Override
+    public Set<String> keySet() {
+        return params.keySet();
+    }
+
+    @Override
+    public Collection<Object> values() {
+        return params.values();
+    }
+
+    @Override
+    public Set<Entry<String, Object>> entrySet() {
+        return params.entrySet();
+    }
+}

server/src/main/java/org/elasticsearch/script/ScriptModule.java

@@ -44,7 +44,7 @@ public class ScriptModule {
             SearchScript.AGGS_CONTEXT,
             ScoreScript.CONTEXT,
             SearchScript.SCRIPT_SORT_CONTEXT,
-            SearchScript.TERMS_SET_QUERY_CONTEXT,
+            TermsSetQueryScript.CONTEXT,
             ExecutableScript.CONTEXT,
             UpdateScript.CONTEXT,
             BucketAggregationScript.CONTEXT,

server/src/main/java/org/elasticsearch/script/SearchScript.java

@@ -149,6 +149,4 @@ public interface Factory {
     public static final ScriptContext<Factory> AGGS_CONTEXT = new ScriptContext<>("aggs", Factory.class);
     // Can return a double. (For ScriptSortType#NUMBER only, for ScriptSortType#STRING normal CONTEXT should be used)
     public static final ScriptContext<Factory> SCRIPT_SORT_CONTEXT = new ScriptContext<>("sort", Factory.class);
-    // Can return a long
-    public static final ScriptContext<Factory> TERMS_SET_QUERY_CONTEXT = new ScriptContext<>("terms_set", Factory.class);
 }

server/src/main/java/org/elasticsearch/script/TermsSetQueryScript.java

@@ -0,0 +1,112 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.elasticsearch.script;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.lucene.index.LeafReaderContext;
+import org.elasticsearch.index.fielddata.ScriptDocValues;
+import org.elasticsearch.search.lookup.LeafSearchLookup;
+import org.elasticsearch.search.lookup.SearchLookup;
+
+public abstract class TermsSetQueryScript {
+
+    public static final String[] PARAMETERS = {};
+
+    public static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("terms_set", Factory.class);
+
+    private static final Map<String, String> DEPRECATIONS;
+
+    static {
+        Map<String, String> deprecations = new HashMap<>();
+        deprecations.put(
+            "doc",
+            "Accessing variable [doc] via [params.doc] from within a terms-set-query-script " +
+                "is deprecated in favor of directly accessing [doc]."
+        );
+        deprecations.put(
+            "_doc",
+            "Accessing variable [doc] via [params._doc] from within a terms-set-query-script " +
+                "is deprecated in favor of directly accessing [doc]."
+        );
+        DEPRECATIONS = Collections.unmodifiableMap(deprecations);
+    }
+
+    /**
+     * The generic runtime parameters for the script.
+     */
+    private final Map<String, Object> params;
+
+    /**
+     * A leaf lookup for the bound segment this script will operate on.
+     */
+    private final LeafSearchLookup leafLookup;
+
+    public TermsSetQueryScript(Map<String, Object> params, SearchLookup lookup, LeafReaderContext leafContext) {
+        this.params = new ParameterMap(params, DEPRECATIONS);
+        this.leafLookup = lookup.getLeafSearchLookup(leafContext);
+    }
+
+    /**
+     * Return the parameters for this script.
+     */
+    public Map<String, Object> getParams() {
+        this.params.putAll(leafLookup.asMap());
+        return params;
+    }
+
+    /**
+     * The doc lookup for the Lucene segment this script was created for.
+     */
+    public Map<String, ScriptDocValues<?>> getDoc() {
+        return leafLookup.doc();
+    }
+
+    /**
+     * Set the current document to run the script on next.
+     */
+    public void setDocument(int docid) {
+        leafLookup.setDocument(docid);
+    }
+
+    /**
+     * Return the result as a long. This is used by aggregation scripts over long fields.
+     */
+    public long runAsLong() {
+        return execute().longValue();
+    }
+
+    public abstract Number execute();
+
+    /**
+     * A factory to construct {@link TermsSetQueryScript} instances.
+     */
+    public interface LeafFactory {
+        TermsSetQueryScript newInstance(LeafReaderContext ctx) throws IOException;
+    }
+
+    /**
+     * A factory to construct stateful {@link TermsSetQueryScript} factories for a specific index.
+     */
+    public interface Factory {
+        LeafFactory newFactory(Map<String, Object> params, SearchLookup lookup);
+    }
+}

test/framework/src/main/java/org/elasticsearch/script/MockScriptEngine.java

@@ -85,6 +85,18 @@ public <T> T compile(String name, String source, ScriptContext<T> context, Map<S
         if (context.instanceClazz.equals(SearchScript.class)) {
             SearchScript.Factory factory = mockCompiled::createSearchScript;
             return context.factoryClazz.cast(factory);
+        } else if(context.instanceClazz.equals(TermsSetQueryScript.class)) {
+            TermsSetQueryScript.Factory factory = (parameters, lookup) -> (TermsSetQueryScript.LeafFactory) ctx
+                -> new TermsSetQueryScript(parameters, lookup, ctx) {
+                @Override
+                public Number execute() {
+                    Map<String, Object> vars = new HashMap<>(parameters);
+                    vars.put("params", parameters);
+                    vars.put("doc", getDoc());
+                    return (Number) script.apply(vars);
+                }
+            };
+            return context.factoryClazz.cast(factory);
         } else if (context.instanceClazz.equals(ExecutableScript.class)) {
             ExecutableScript.Factory factory = mockCompiled::createExecutableScript;
             return context.factoryClazz.cast(factory);