Scripting: Add StatefulFactoryType as optional intermediate factory in script contexts (#24974)

ScriptContexts currently understand a FactoryType that can produce
instances of the script InstanceType. However, for search scripts, this
does not work as we have the concept of LeafSearchScript that is created
per lucene segment. This commit effectively renames the existing
SearchScript class into SearchScript.LeafFactory, which is a new,
optional, class that can be defined within a ScriptContext.
LeafSearchScript is effectively renamed back into SearchScript. This
change allows the model of stateless factory -> stateful factory ->
script instance to continue, but in a generic way that any script
context may take advantage of.

relates #20426

22 files changed, 243 insertions, 184 deletions

diff --git a/core/src/main/java/org/elasticsearch/common/lucene/search/function/ScriptScoreFunction.java b/core/src/main/java/org/elasticsearch/common/lucene/search/function/ScriptScoreFunction.java
index bee7087c1d..112bf271c4 100644
--- a/core/src/main/java/org/elasticsearch/common/lucene/search/function/ScriptScoreFunction.java
+++ b/core/src/main/java/org/elasticsearch/common/lucene/search/function/ScriptScoreFunction.java
@@ -24,7 +24,6 @@ import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.Explanation;
 import org.apache.lucene.search.Scorer;
 import org.elasticsearch.script.ExplainableSearchScript;
-import org.elasticsearch.script.LeafSearchScript;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.GeneralScriptException;
 import org.elasticsearch.script.SearchScript;
@@ -65,10 +64,10 @@ public class ScriptScoreFunction extends ScoreFunction {
 
     private final Script sScript;
 
-    private final SearchScript script;
+    private final SearchScript.LeafFactory script;
 
 
-    public ScriptScoreFunction(Script sScript, SearchScript script) {
+    public ScriptScoreFunction(Script sScript, SearchScript.LeafFactory script) {
         super(CombineFunction.REPLACE);
         this.sScript = sScript;
         this.script = script;
@@ -76,7 +75,7 @@ public class ScriptScoreFunction extends ScoreFunction {
 
     @Override
     public LeafScoreFunction getLeafScoreFunction(LeafReaderContext ctx) throws IOException {
-        final LeafSearchScript leafScript = script.getLeafSearchScript(ctx);
+        final SearchScript leafScript = script.newInstance(ctx);
         final CannedScorer scorer = new CannedScorer();
         leafScript.setScorer(scorer);
         return new LeafScoreFunction() {
diff --git a/core/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java b/core/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java
index 74c8055410..68eb749e11 100644
--- a/core/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java
+++ b/core/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java
@@ -74,7 +74,7 @@ public abstract class InnerHitContextBuilder {
         }
         if (innerHitBuilder.getScriptFields() != null) {
             for (SearchSourceBuilder.ScriptField field : innerHitBuilder.getScriptFields()) {
-                SearchScript searchScript = innerHitsContext.getQueryShardContext().getSearchScript(field.script(),
+                SearchScript.LeafFactory searchScript = innerHitsContext.getQueryShardContext().getSearchScript(field.script(),
                     SearchScript.CONTEXT);
                 innerHitsContext.scriptFields().add(new org.elasticsearch.search.fetch.subphase.ScriptFieldsContext.ScriptField(
                     field.fieldName(), searchScript, field.ignoreFailure()));
diff --git a/core/src/main/java/org/elasticsearch/index/query/QueryShardContext.java b/core/src/main/java/org/elasticsearch/index/query/QueryShardContext.java
index 1ab89a0fe0..8e54af2bd6 100644
--- a/core/src/main/java/org/elasticsearch/index/query/QueryShardContext.java
+++ b/core/src/main/java/org/elasticsearch/index/query/QueryShardContext.java
@@ -329,21 +329,21 @@ public class QueryShardContext extends QueryRewriteContext {
      * Compiles (or retrieves from cache) and binds the parameters to the
      * provided script
      */
-    public final SearchScript getSearchScript(Script script, ScriptContext<SearchScript.Factory> context) {
+    public final SearchScript.LeafFactory getSearchScript(Script script, ScriptContext<SearchScript.Factory> context) {
         failIfFrozen();
         SearchScript.Factory factory = scriptService.compile(script, context);
-        return factory.newInstance(script.getParams(), lookup());
+        return factory.newFactory(script.getParams(), lookup());
     }
     /**
      * Returns a lazily created {@link SearchScript} that is compiled immediately but can be pulled later once all
      * parameters are available.
      */
-    public final Function<Map<String, Object>, SearchScript> getLazySearchScript(
+    public final Function<Map<String, Object>, SearchScript.LeafFactory> getLazySearchScript(
         Script script, ScriptContext<SearchScript.Factory> context) {
         // TODO: this "lazy" binding can be removed once scripted metric aggs have their own contexts, which take _agg/_aggs as a parameter
         failIfFrozen();
         SearchScript.Factory factory = scriptService.compile(script, context);
-        return (p) -> factory.newInstance(p, lookup());
+        return (p) -> factory.newFactory(p, lookup());
     }
 
     /**
diff --git a/core/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java b/core/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java
index 88fde50eb1..c86ff73d0f 100644
--- a/core/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java
+++ b/core/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java
@@ -34,9 +34,7 @@ import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
-import org.elasticsearch.script.LeafSearchScript;
 import org.elasticsearch.script.Script;
-import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.SearchScript;
 
 import java.io.IOException;
@@ -137,9 +135,9 @@ public class ScriptQueryBuilder extends AbstractQueryBuilder<ScriptQueryBuilder>
     static class ScriptQuery extends Query {
 
         final Script script;
-        final SearchScript searchScript;
+        final SearchScript.LeafFactory searchScript;
 
-        ScriptQuery(Script script, SearchScript searchScript) {
+        ScriptQuery(Script script, SearchScript.LeafFactory searchScript) {
             this.script = script;
             this.searchScript = searchScript;
         }
@@ -181,7 +179,7 @@ public class ScriptQueryBuilder extends AbstractQueryBuilder<ScriptQueryBuilder>
                 @Override
                 public Scorer scorer(LeafReaderContext context) throws IOException {
                     DocIdSetIterator approximation = DocIdSetIterator.all(context.reader().maxDoc());
-                    final LeafSearchScript leafScript = searchScript.getLeafSearchScript(context);
+                    final SearchScript leafScript = searchScript.newInstance(context);
                     TwoPhaseIterator twoPhase = new TwoPhaseIterator(approximation) {
 
                         @Override
diff --git a/core/src/main/java/org/elasticsearch/index/query/functionscore/ScriptScoreFunctionBuilder.java b/core/src/main/java/org/elasticsearch/index/query/functionscore/ScriptScoreFunctionBuilder.java
index 2af4c9dd75..4c9ac0f452 100644
--- a/core/src/main/java/org/elasticsearch/index/query/functionscore/ScriptScoreFunctionBuilder.java
+++ b/core/src/main/java/org/elasticsearch/index/query/functionscore/ScriptScoreFunctionBuilder.java
@@ -94,7 +94,7 @@ public class ScriptScoreFunctionBuilder extends ScoreFunctionBuilder<ScriptScore
     @Override
     protected ScoreFunction doToFunction(QueryShardContext context) {
         try {
-            SearchScript searchScript = context.getSearchScript(script, SearchScript.CONTEXT);
+            SearchScript.LeafFactory searchScript = context.getSearchScript(script, SearchScript.CONTEXT);
             return new ScriptScoreFunction(script, searchScript);
         } catch (Exception e) {
             throw new QueryShardException(context, "script_score: the script could not be loaded", e);
diff --git a/core/src/main/java/org/elasticsearch/script/ExplainableSearchScript.java b/core/src/main/java/org/elasticsearch/script/ExplainableSearchScript.java
index 9284b36c38..c397021275 100644
--- a/core/src/main/java/org/elasticsearch/script/ExplainableSearchScript.java
+++ b/core/src/main/java/org/elasticsearch/script/ExplainableSearchScript.java
@@ -47,7 +47,7 @@ import java.io.IOException;
  * This is currently not used inside elasticsearch but it is used, see for example here:
  * https://github.com/elastic/elasticsearch/issues/8561
  */
-public interface ExplainableSearchScript extends LeafSearchScript {
+public interface ExplainableSearchScript {
 
     /**
      * Build the explanation of the current document being scored
diff --git a/core/src/main/java/org/elasticsearch/script/LeafSearchScript.java b/core/src/main/java/org/elasticsearch/script/LeafSearchScript.java
deleted file mode 100644
index 762168d3c9..0000000000
--- a/core/src/main/java/org/elasticsearch/script/LeafSearchScript.java
+++ /dev/null
@@ -1,78 +0,0 @@
-/*
- * 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 org.apache.lucene.search.Scorer;
-import org.elasticsearch.common.lucene.ScorerAware;
-
-import java.util.Map;
-
-/**
- * A per-segment {@link SearchScript}.
- *
- * This is effectively a functional interface, requiring at least implementing {@link #runAsDouble()}.
- */
-public interface LeafSearchScript extends ScorerAware, ExecutableScript {
-
-    /**
-     * Set the document this script will process next.
-     */
-    default void setDocument(int doc) {}
-
-    @Override
-    default void setScorer(Scorer scorer) {}
-
-    /**
-     * Set the source for the current document.
-     */
-    default void setSource(Map<String, Object> source) {}
-    
-    /**
-     * Sets per-document aggregation {@code _value}.
-     * <p>
-     * The default implementation just calls {@code setNextVar("_value", value)} but
-     * some engines might want to handle this differently for better performance.
-     * <p>
-     * @param value per-document value, typically a String, Long, or Double
-     */
-    default void setNextAggregationValue(Object value) {
-        setNextVar("_value", value);
-    }
-
-    @Override
-    default void setNextVar(String field, Object value) {}
-
-    /**
-     * Return the result as a long. This is used by aggregation scripts over long fields.
-     */
-    default long runAsLong() {
-        throw new UnsupportedOperationException("runAsLong is not implemented");
-    }
-
-    @Override
-    default Object run() {
-        return runAsDouble();
-    }
-
-    /**
-     * Return the result as a double. This is the main use case of search script, used for document scoring.
-     */
-    double runAsDouble();
-}
diff --git a/core/src/main/java/org/elasticsearch/script/ScriptContext.java b/core/src/main/java/org/elasticsearch/script/ScriptContext.java
index 420257e0e5..3f931f659e 100644
--- a/core/src/main/java/org/elasticsearch/script/ScriptContext.java
+++ b/core/src/main/java/org/elasticsearch/script/ScriptContext.java
@@ -27,27 +27,37 @@ import java.lang.reflect.Method;
  * A {@link ScriptContext} contains the information related to a single use case and the interfaces
  * and methods necessary for a {@link ScriptEngine} to implement.
  * <p>
- * There are two related classes which must be supplied to construct a {@link ScriptContext}.
+ * There are at least two (and optionally a third) related classes which must be defined.
  * <p>
- * The <i>FactoryType</i> is a factory class for constructing instances of a script. The
- * {@link ScriptService} returns an instance of <i>FactoryType</i> when compiling a script. This class
- * must be stateless so it is cacheable by the {@link ScriptService}. It must have an abstract method
- * named {@code newInstance} which {@link ScriptEngine} implementations will define.
- * <p>
- * The <i>InstanceType</i> is a class returned by the {@code newInstance} method of the
- * <i>FactoryType</i>. It is an instance of a script and may be stateful. Instances of
+ * The <i>InstanceType</i> is a class which users of the script api call to execute a script. It
+ * may be stateful. Instances of
  * the <i>InstanceType</i> may be executed multiple times by a caller with different arguments. This
  * class must have an abstract method named {@code execute} which {@link ScriptEngine} implementations
  * will define.
+ * <p>
+ * The <i>FactoryType</i> is a factory class returned by the {@link ScriptService} when compiling
+ * a script. This class must be stateless so it is cacheable by the {@link ScriptService}. It must
+ * have one of the following:
+ * <ul>
+ *     <li>An abstract method named {@code newInstance} which returns an instance of <i>InstanceType</i></li>
+ *     <li>An abstract method named {@code newFactory} which returns an instance of <i>StatefulFactoryType</i></li>
+ * </ul>
+ * <p>
+ * The <i>StatefulFactoryType</i> is an optional class which allows a stateful factory from the
+ * stateless factory type required by the {@link ScriptService}. If defined, the <i>StatefulFactoryType</i>
+ * must have a method named {@code newInstance} which returns an instance of <i>InstanceType</i>.
  */
 public final class ScriptContext<FactoryType> {
 
     /** A unique identifier for this context. */
     public final String name;
 
-    /** A factory class for constructing instances of a script. */
+    /** A factory class for constructing script or stateful factory instances. */
     public final Class<FactoryType> factoryClazz;
 
+    /** A factory class for construct script instances. */
+    public final Class<?> statefulFactoryClazz;
+
     /** A class that is an instance of a script. */
     public final Class<?> instanceClazz;
 
@@ -55,20 +65,38 @@ public final class ScriptContext<FactoryType> {
     public ScriptContext(String name, Class<FactoryType> factoryClazz) {
         this.name = name;
         this.factoryClazz = factoryClazz;
-        Method newInstanceMethod = null;
-        for (Method method : factoryClazz.getMethods()) {
-            if (method.getName().equals("newInstance")) {
-                if (newInstanceMethod != null) {
-                    throw new IllegalArgumentException("Cannot have multiple newInstance methods on FactoryType class ["
-                        + factoryClazz.getName() + "] for script context [" + name + "]");
-                }
-                newInstanceMethod = method;
+        Method newInstanceMethod = findMethod("FactoryType", factoryClazz, "newInstance");
+        Method newFactoryMethod = findMethod("FactoryType", factoryClazz, "newFactory");
+        if (newFactoryMethod != null) {
+            assert newInstanceMethod == null;
+            statefulFactoryClazz = newFactoryMethod.getReturnType();
+            newInstanceMethod = findMethod("StatefulFactoryType", statefulFactoryClazz, "newInstance");
+            if (newInstanceMethod == null) {
+                throw new IllegalArgumentException("Could not find method newInstance StatefulFactoryType class ["
+                    + statefulFactoryClazz.getName() + "] for script context [" + name + "]");
             }
-        }
-        if (newInstanceMethod == null) {
-            throw new IllegalArgumentException("Could not find method newInstance on FactoryType class ["
+        } else if (newInstanceMethod != null) {
+            assert newFactoryMethod == null;
+            statefulFactoryClazz = null;
+        } else {
+            throw new IllegalArgumentException("Could not find method newInstance or method newFactory on FactoryType class ["
                 + factoryClazz.getName() + "] for script context [" + name + "]");
         }
         instanceClazz = newInstanceMethod.getReturnType();
     }
+
+    /** Returns a method with the given name, or throws an exception if multiple are found. */
+    private Method findMethod(String type, Class<?> clazz, String methodName) {
+        Method foundMethod = null;
+        for (Method method : clazz.getMethods()) {
+            if (method.getName().equals(methodName)) {
+                if (foundMethod != null) {
+                    throw new IllegalArgumentException("Cannot have multiple " + methodName + " methods on " + type + " class ["
+                        + clazz.getName() + "] for script context [" + name + "]");
+                }
+                foundMethod = method;
+            }
+        }
+        return foundMethod;
+    }
 }
diff --git a/core/src/main/java/org/elasticsearch/script/SearchScript.java b/core/src/main/java/org/elasticsearch/script/SearchScript.java
index bbee2910c8..4c50147b22 100644
--- a/core/src/main/java/org/elasticsearch/script/SearchScript.java
+++ b/core/src/main/java/org/elasticsearch/script/SearchScript.java
@@ -19,30 +19,146 @@
 package org.elasticsearch.script;
 
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.Scorer;
+import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.common.lucene.ScorerAware;
+import org.elasticsearch.search.lookup.LeafDocLookup;
+import org.elasticsearch.search.lookup.LeafSearchLookup;
 import org.elasticsearch.search.lookup.SearchLookup;
 
 import java.io.IOException;
 import java.util.Map;
 
 /**
- * A search script.
+ * A generic script used for per document use cases.
+ *
+ * Using a {@link SearchScript} works as follows:
+ * <ol>
+ *     <li>Construct a {@link Factory} using {@link ScriptService#compile(Script, ScriptContext)}</li>
+ *     <li>Construct a {@link LeafFactory} for a an index using {@link Factory#newFactory(Map, SearchLookup)}</li>
+ *     <li>Construct a {@link SearchScript} for a Lucene segment using {@link LeafFactory#newInstance(LeafReaderContext)}</li>
+ *     <li>Call {@link #setDocument(int)} to indicate which document in the segment the script should be run for next</li>
+ *     <li>Call one of the {@code run} methods: {@link #run()}, {@link #runAsDouble()}, or {@link #runAsLong()}</li>
+ * </ol>
  */
-public interface SearchScript {
+public abstract class SearchScript implements ScorerAware, ExecutableScript {
+
+    /** The generic runtime parameters for the script. */
+    private final Map<String, Object> params;
+
+    /** A lookup for the index this script will operate on. */
+    private final SearchLookup lookup;
+
+    /** A leaf lookup for the bound segment this script will operate on. */
+    private final LeafReaderContext leafContext;
+
+    /** A leaf lookup for the bound segment this script will operate on. */
+    private final LeafSearchLookup leafLookup;
+
+    /** A scorer that will return the score for the current document when the script is run. */
+    private Scorer scorer;
+
+    public SearchScript(Map<String, Object> params, SearchLookup lookup, LeafReaderContext leafContext) {
+        this.params = params;
+        this.lookup = lookup;
+        this.leafContext = leafContext;
+        // TODO: remove leniency when painless does not implement SearchScript for executable script cases
+        this.leafLookup = leafContext == null ? null : lookup.getLeafSearchLookup(leafContext);
+    }
+
+    /** Return the parameters for this script. */
+    public Map<String, Object> getParams() {
+        return params;
+    }
+
+    /** The leaf lookup for the Lucene segment this script was created for. */
+    protected final LeafSearchLookup getLeafLookup() {
+        return leafLookup;
+    }
+
+    /** The leaf context for the Lucene segment this script was created for. */
+    protected final LeafReaderContext getLeafContext() {
+        return leafContext;
+    }
+
+    /** The doc lookup for the Lucene segment this script was created for. */
+    public final LeafDocLookup getDoc() {
+        // TODO: remove leniency when painless does not implement SearchScript for executable script cases
+        return leafLookup == null ? null : leafLookup.doc();
+    }
+
+    /** Set the current document to run the script on next. */
+    public void setDocument(int docid) {
+        // TODO: remove leniency when painless does not implement SearchScript for executable script cases
+        if (leafLookup != null) {
+            leafLookup.setDocument(docid);
+        }
+    }
+
+    @Override
+    public void setScorer(Scorer scorer) {
+        this.scorer = scorer;
+    }
 
-    LeafSearchScript getLeafSearchScript(LeafReaderContext context) throws IOException;
+    /** Return the score of the current document. */
+    public double getScore() {
+        // TODO: remove leniency when painless does not implement SearchScript for executable script cases
+        if (scorer == null) {
+            return 0.0d;
+        }
+        try {
+            return scorer.score();
+        } catch (IOException e) {
+            throw new ElasticsearchException("couldn't lookup score", e);
+        }
+    }
 
     /**
-     * Indicates if document scores may be needed by this {@link SearchScript}.
-     * 
-     * @return {@code true} if scores are needed.
+     * Sets per-document aggregation {@code _value}.
+     * <p>
+     * The default implementation just calls {@code setNextVar("_value", value)} but
+     * some engines might want to handle this differently for better performance.
+     * <p>
+     * @param value per-document value, typically a String, Long, or Double
      */
-    boolean needsScores();
+    public void setNextAggregationValue(Object value) {
+        setNextVar("_value", value);
+    }
+
+    @Override
+    public void setNextVar(String field, Object value) {}
+
+    /** Return the result as a long. This is used by aggregation scripts over long fields. */
+    public long runAsLong() {
+        throw new UnsupportedOperationException("runAsLong is not implemented");
+    }
+
+    @Override
+    public Object run() {
+        return runAsDouble();
+    }
+
+    /** Return the result as a double. This is the main use case of search script, used for document scoring. */
+    public abstract double runAsDouble();
+
+    /** A factory to construct {@link SearchScript} instances. */
+    public interface LeafFactory {
+        SearchScript newInstance(LeafReaderContext ctx) throws IOException;
+        /**
+         * Indicates if document scores may be needed by this {@link SearchScript}.
+         *
+         * @return {@code true} if scores are needed.
+         */
+        boolean needsScores();
+    }
 
-    interface Factory {
-        SearchScript newInstance(Map<String, Object> params, SearchLookup lookup);
+    /** A factory to construct stateful {@link SearchScript} factories for a specific index. */
+    public interface Factory {
+        LeafFactory newFactory(Map<String, Object> params, SearchLookup lookup);
     }
 
-    ScriptContext<Factory> CONTEXT = new ScriptContext<>("search", Factory.class);
+    /** The context used to compile {@link SearchScript} factories. */
+    public static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("search", Factory.class);
     // TODO: remove aggs context when it has its own interface
-    ScriptContext<Factory> AGGS_CONTEXT = new ScriptContext<>("aggs", Factory.class);
+    public static final ScriptContext<Factory> AGGS_CONTEXT = new ScriptContext<>("aggs", Factory.class);
 }
\ No newline at end of file
diff --git a/core/src/main/java/org/elasticsearch/search/SearchService.java b/core/src/main/java/org/elasticsearch/search/SearchService.java
index fef20f44f5..93217819ab 100644
--- a/core/src/main/java/org/elasticsearch/search/SearchService.java
+++ b/core/src/main/java/org/elasticsearch/search/SearchService.java
@@ -689,7 +689,7 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
         if (source.scriptFields() != null) {
             for (org.elasticsearch.search.builder.SearchSourceBuilder.ScriptField field : source.scriptFields()) {
                 SearchScript.Factory factory = scriptService.compile(field.script(), SearchScript.CONTEXT);
-                SearchScript searchScript = factory.newInstance(field.script().getParams(), context.lookup());
+                SearchScript.LeafFactory searchScript = factory.newFactory(field.script().getParams(), context.lookup());
                 context.scriptFields().add(new ScriptField(field.fieldName(), searchScript, field.ignoreFailure()));
             }
         }
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregationBuilder.java b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregationBuilder.java
index 34c06ec3c5..5b9e7e1dbd 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregationBuilder.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregationBuilder.java
@@ -190,7 +190,7 @@ public class ScriptedMetricAggregationBuilder extends AbstractAggregationBuilder
         } else {
             executableInitScript = (p) -> null;
         }
-        Function<Map<String, Object>, SearchScript> searchMapScript =
+        Function<Map<String, Object>, SearchScript.LeafFactory> searchMapScript =
             queryShardContext.getLazySearchScript(mapScript, SearchScript.AGGS_CONTEXT);
         Function<Map<String, Object>, ExecutableScript> executableCombineScript;
         if (combineScript != null) {
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregator.java b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregator.java
index cee7b3402f..bebe9f892b 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregator.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregator.java
@@ -21,7 +21,6 @@ package org.elasticsearch.search.aggregations.metrics.scripted;
 
 import org.apache.lucene.index.LeafReaderContext;
 import org.elasticsearch.script.ExecutableScript;
-import org.elasticsearch.script.LeafSearchScript;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.SearchScript;
 import org.elasticsearch.search.aggregations.Aggregator;
@@ -38,12 +37,12 @@ import java.util.Map;
 
 public class ScriptedMetricAggregator extends MetricsAggregator {
 
-    private final SearchScript mapScript;
+    private final SearchScript.LeafFactory mapScript;
     private final ExecutableScript combineScript;
     private final Script reduceScript;
     private Map<String, Object> params;
 
-    protected ScriptedMetricAggregator(String name, SearchScript mapScript, ExecutableScript combineScript,
+    protected ScriptedMetricAggregator(String name, SearchScript.LeafFactory mapScript, ExecutableScript combineScript,
                                        Script reduceScript,
             Map<String, Object> params, SearchContext context, Aggregator parent, List<PipelineAggregator> pipelineAggregators, Map<String, Object> metaData)
             throws IOException {
@@ -62,7 +61,7 @@ public class ScriptedMetricAggregator extends MetricsAggregator {
     @Override
     public LeafBucketCollector getLeafCollector(LeafReaderContext ctx,
             final LeafBucketCollector sub) throws IOException {
-        final LeafSearchScript leafMapScript = mapScript.getLeafSearchScript(ctx);
+        final SearchScript leafMapScript = mapScript.newInstance(ctx);
         return new LeafBucketCollectorBase(sub, leafMapScript) {
             @Override
             public void collect(int doc, long bucket) throws IOException {
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregatorFactory.java b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregatorFactory.java
index bac2becc8e..6e75cba7b7 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregatorFactory.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregatorFactory.java
@@ -38,13 +38,13 @@ import java.util.function.Function;
 
 public class ScriptedMetricAggregatorFactory extends AggregatorFactory<ScriptedMetricAggregatorFactory> {
 
-    private final Function<Map<String, Object>, SearchScript> mapScript;
+    private final Function<Map<String, Object>, SearchScript.LeafFactory> mapScript;
     private final Function<Map<String, Object>, ExecutableScript> combineScript;
     private final Script reduceScript;
     private final Map<String, Object> params;
     private final Function<Map<String, Object>, ExecutableScript> initScript;
 
-    public ScriptedMetricAggregatorFactory(String name, Function<Map<String, Object>, SearchScript> mapScript,
+    public ScriptedMetricAggregatorFactory(String name, Function<Map<String, Object>, SearchScript.LeafFactory> mapScript,
             Function<Map<String, Object>, ExecutableScript> initScript, Function<Map<String, Object>, ExecutableScript> combineScript,
             Script reduceScript, Map<String, Object> params, SearchContext context, AggregatorFactory<?> parent,
             AggregatorFactories.Builder subFactories, Map<String, Object> metaData) throws IOException {
@@ -71,7 +71,7 @@ public class ScriptedMetricAggregatorFactory extends AggregatorFactory<ScriptedM
         }
 
         final ExecutableScript initScript = this.initScript.apply(params);
-        final SearchScript mapScript = this.mapScript.apply(params);
+        final SearchScript.LeafFactory mapScript = this.mapScript.apply(params);
         final ExecutableScript combineScript = this.combineScript.apply(params);
 
         final Script reduceScript = deepCopyScript(this.reduceScript, context);
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregationBuilder.java b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregationBuilder.java
index 94467add51..31cd1681f8 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregationBuilder.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregationBuilder.java
@@ -533,7 +533,7 @@ public class TopHitsAggregationBuilder extends AbstractAggregationBuilder<TopHit
         List<ScriptFieldsContext.ScriptField> fields = new ArrayList<>();
         if (scriptFields != null) {
             for (ScriptField field : scriptFields) {
-                SearchScript searchScript = context.getQueryShardContext().getSearchScript(field.script(),
+                SearchScript.LeafFactory searchScript = context.getQueryShardContext().getSearchScript(field.script(),
                     SearchScript.CONTEXT);
                 fields.add(new org.elasticsearch.search.fetch.subphase.ScriptFieldsContext.ScriptField(
                     field.fieldName(), searchScript, field.ignoreFailure()));
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSource.java b/core/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSource.java
index a8aaa25940..5fca34beff 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSource.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSource.java
@@ -40,7 +40,6 @@ import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
 import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
 import org.elasticsearch.index.fielddata.SortingBinaryDocValues;
 import org.elasticsearch.index.fielddata.SortingNumericDoubleValues;
-import org.elasticsearch.script.LeafSearchScript;
 import org.elasticsearch.script.SearchScript;
 import org.elasticsearch.search.aggregations.support.ValuesSource.WithScript.BytesValues;
 import org.elasticsearch.search.aggregations.support.values.ScriptBytesValues;
@@ -162,15 +161,15 @@ public abstract class ValuesSource {
 
         public static class Script extends Bytes {
 
-            private final SearchScript script;
+            private final SearchScript.LeafFactory script;
 
-            public Script(SearchScript script) {
+            public Script(SearchScript.LeafFactory script) {
                 this.script = script;
             }
 
             @Override
             public SortedBinaryDocValues bytesValues(LeafReaderContext context) throws IOException {
-                return new ScriptBytesValues(script.getLeafSearchScript(context));
+                return new ScriptBytesValues(script.newInstance(context));
             }
 
             @Override
@@ -233,9 +232,9 @@ public abstract class ValuesSource {
         public static class WithScript extends Numeric {
 
             private final Numeric delegate;
-            private final SearchScript script;
+            private final SearchScript.LeafFactory script;
 
-            public WithScript(Numeric delegate, SearchScript script) {
+            public WithScript(Numeric delegate, SearchScript.LeafFactory script) {
                 this.delegate = delegate;
                 this.script = script;
             }
@@ -252,25 +251,25 @@ public abstract class ValuesSource {
 
             @Override
             public SortedBinaryDocValues bytesValues(LeafReaderContext context) throws IOException {
-                return new ValuesSource.WithScript.BytesValues(delegate.bytesValues(context), script.getLeafSearchScript(context));
+                return new ValuesSource.WithScript.BytesValues(delegate.bytesValues(context), script.newInstance(context));
             }
 
             @Override
             public SortedNumericDocValues longValues(LeafReaderContext context) throws IOException {
-                return new LongValues(delegate.longValues(context), script.getLeafSearchScript(context));
+                return new LongValues(delegate.longValues(context), script.newInstance(context));
             }
 
             @Override
             public SortedNumericDoubleValues doubleValues(LeafReaderContext context) throws IOException {
-                return new DoubleValues(delegate.doubleValues(context), script.getLeafSearchScript(context));
+                return new DoubleValues(delegate.doubleValues(context), script.newInstance(context));
             }
 
             static class LongValues extends AbstractSortingNumericDocValues implements ScorerAware {
 
                 private final SortedNumericDocValues longValues;
-                private final LeafSearchScript script;
+                private final SearchScript script;
 
-                LongValues(SortedNumericDocValues values, LeafSearchScript script) {
+                LongValues(SortedNumericDocValues values, SearchScript script) {
                     this.longValues = values;
                     this.script = script;
                 }
@@ -299,9 +298,9 @@ public abstract class ValuesSource {
             static class DoubleValues extends SortingNumericDoubleValues implements ScorerAware {
 
                 private final SortedNumericDoubleValues doubleValues;
-                private final LeafSearchScript script;
+                private final SearchScript script;
 
-                DoubleValues(SortedNumericDoubleValues values, LeafSearchScript script) {
+                DoubleValues(SortedNumericDoubleValues values, SearchScript script) {
                     this.doubleValues = values;
                     this.script = script;
                 }
@@ -358,10 +357,10 @@ public abstract class ValuesSource {
         }
 
         public static class Script extends Numeric {
-            private final SearchScript script;
+            private final SearchScript.LeafFactory script;
             private final ValueType scriptValueType;
 
-            public Script(SearchScript script, ValueType scriptValueType) {
+            public Script(SearchScript.LeafFactory script, ValueType scriptValueType) {
                 this.script = script;
                 this.scriptValueType = scriptValueType;
             }
@@ -373,17 +372,17 @@ public abstract class ValuesSource {
 
             @Override
             public SortedNumericDocValues longValues(LeafReaderContext context) throws IOException {
-                return new ScriptLongValues(script.getLeafSearchScript(context));
+                return new ScriptLongValues(script.newInstance(context));
             }
 
             @Override
             public SortedNumericDoubleValues doubleValues(LeafReaderContext context) throws IOException {
-                return new ScriptDoubleValues(script.getLeafSearchScript(context));
+                return new ScriptDoubleValues(script.newInstance(context));
             }
 
             @Override
             public SortedBinaryDocValues bytesValues(LeafReaderContext context) throws IOException {
-                return new ScriptBytesValues(script.getLeafSearchScript(context));
+                return new ScriptBytesValues(script.newInstance(context));
             }
 
             @Override
@@ -398,9 +397,9 @@ public abstract class ValuesSource {
     public static class WithScript extends Bytes {
 
         private final ValuesSource delegate;
-        private final SearchScript script;
+        private final SearchScript.LeafFactory script;
 
-        public WithScript(ValuesSource delegate, SearchScript script) {
+        public WithScript(ValuesSource delegate, SearchScript.LeafFactory script) {
             this.delegate = delegate;
             this.script = script;
         }
@@ -412,15 +411,15 @@ public abstract class ValuesSource {
 
         @Override
         public SortedBinaryDocValues bytesValues(LeafReaderContext context) throws IOException {
-            return new BytesValues(delegate.bytesValues(context), script.getLeafSearchScript(context));
+            return new BytesValues(delegate.bytesValues(context), script.newInstance(context));
         }
 
         static class BytesValues extends SortingBinaryDocValues implements ScorerAware {
 
             private final SortedBinaryDocValues bytesValues;
-            private final LeafSearchScript script;
+            private final SearchScript script;
 
-            BytesValues(SortedBinaryDocValues bytesValues, LeafSearchScript script) {
+            BytesValues(SortedBinaryDocValues bytesValues, SearchScript script) {
                 this.bytesValues = bytesValues;
                 this.script = script;
             }
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceConfig.java b/core/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceConfig.java
index 6404cae0e4..d7b6272308 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceConfig.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceConfig.java
@@ -116,7 +116,7 @@ public class ValuesSourceConfig<VS extends ValuesSource> {
         return config;
     }
 
-    private static SearchScript createScript(Script script, QueryShardContext context) {
+    private static SearchScript.LeafFactory createScript(Script script, QueryShardContext context) {
         if (script == null) {
             return null;
         } else {
@@ -137,7 +137,7 @@ public class ValuesSourceConfig<VS extends ValuesSource> {
 
     private final ValuesSourceType valueSourceType;
     private FieldContext fieldContext;
-    private SearchScript script;
+    private SearchScript.LeafFactory script;
     private ValueType scriptValueType;
     private boolean unmapped = false;
     private DocValueFormat format = DocValueFormat.RAW;
@@ -156,7 +156,7 @@ public class ValuesSourceConfig<VS extends ValuesSource> {
         return fieldContext;
     }
 
-    public SearchScript script() {
+    public SearchScript.LeafFactory script() {
         return script;
     }
 
@@ -173,7 +173,7 @@ public class ValuesSourceConfig<VS extends ValuesSource> {
         return this;
     }
 
-    public ValuesSourceConfig<VS> script(SearchScript script) {
+    public ValuesSourceConfig<VS> script(SearchScript.LeafFactory script) {
         this.script = script;
         return this;
     }
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/support/values/ScriptBytesValues.java b/core/src/main/java/org/elasticsearch/search/aggregations/support/values/ScriptBytesValues.java
index 78685ed0e8..38950325da 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/support/values/ScriptBytesValues.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/support/values/ScriptBytesValues.java
@@ -22,7 +22,7 @@ import org.apache.lucene.search.Scorer;
 import org.elasticsearch.common.lucene.ScorerAware;
 import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
 import org.elasticsearch.index.fielddata.SortingBinaryDocValues;
-import org.elasticsearch.script.LeafSearchScript;
+import org.elasticsearch.script.SearchScript;
 
 import java.io.IOException;
 import java.lang.reflect.Array;
@@ -33,9 +33,9 @@ import java.util.Collection;
  */
 public class ScriptBytesValues extends SortingBinaryDocValues implements ScorerAware {
 
-    private final LeafSearchScript script;
+    private final SearchScript script;
 
-    public ScriptBytesValues(LeafSearchScript script) {
+    public ScriptBytesValues(SearchScript script) {
         super();
         this.script = script;
     }
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/support/values/ScriptDoubleValues.java b/core/src/main/java/org/elasticsearch/search/aggregations/support/values/ScriptDoubleValues.java
index 619ffde0a1..ac3c8f682b 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/support/values/ScriptDoubleValues.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/support/values/ScriptDoubleValues.java
@@ -21,7 +21,7 @@ package org.elasticsearch.search.aggregations.support.values;
 import org.apache.lucene.search.Scorer;
 import org.elasticsearch.common.lucene.ScorerAware;
 import org.elasticsearch.index.fielddata.SortingNumericDoubleValues;
-import org.elasticsearch.script.LeafSearchScript;
+import org.elasticsearch.script.SearchScript;
 import org.elasticsearch.search.aggregations.AggregationExecutionException;
 import org.joda.time.ReadableInstant;
 
@@ -34,9 +34,9 @@ import java.util.Collection;
  */
 public class ScriptDoubleValues extends SortingNumericDoubleValues implements ScorerAware {
 
-    final LeafSearchScript script;
+    final SearchScript script;
 
-    public ScriptDoubleValues(LeafSearchScript script) {
+    public ScriptDoubleValues(SearchScript script) {
         super();
         this.script = script;
     }
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/support/values/ScriptLongValues.java b/core/src/main/java/org/elasticsearch/search/aggregations/support/values/ScriptLongValues.java
index 6247e96c7e..818a9d9fd8 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/support/values/ScriptLongValues.java
+++ b/core/src/main/java/org/elasticsearch/search/aggregations/support/values/ScriptLongValues.java
@@ -22,7 +22,7 @@ import org.apache.lucene.search.Scorer;
 import org.apache.lucene.util.LongValues;
 import org.elasticsearch.common.lucene.ScorerAware;
 import org.elasticsearch.index.fielddata.AbstractSortingNumericDocValues;
-import org.elasticsearch.script.LeafSearchScript;
+import org.elasticsearch.script.SearchScript;
 import org.elasticsearch.search.aggregations.AggregationExecutionException;
 import org.joda.time.ReadableInstant;
 
@@ -36,9 +36,9 @@ import java.util.Iterator;
  */
 public class ScriptLongValues extends AbstractSortingNumericDocValues implements ScorerAware {
 
-    final LeafSearchScript script;
+    final SearchScript script;
 
-    public ScriptLongValues(LeafSearchScript script) {
+    public ScriptLongValues(SearchScript script) {
         super();
         this.script = script;
     }
diff --git a/core/src/main/java/org/elasticsearch/search/fetch/subphase/ScriptFieldsContext.java b/core/src/main/java/org/elasticsearch/search/fetch/subphase/ScriptFieldsContext.java
index 79bacd7f93..9e43b0bdd3 100644
--- a/core/src/main/java/org/elasticsearch/search/fetch/subphase/ScriptFieldsContext.java
+++ b/core/src/main/java/org/elasticsearch/search/fetch/subphase/ScriptFieldsContext.java
@@ -28,10 +28,10 @@ public class ScriptFieldsContext {
 
     public static class ScriptField {
         private final String name;
-        private final SearchScript script;
+        private final SearchScript.LeafFactory script;
         private final boolean ignoreException;
 
-        public ScriptField(String name, SearchScript script, boolean ignoreException) {
+        public ScriptField(String name, SearchScript.LeafFactory script, boolean ignoreException) {
             this.name = name;
             this.script = script;
             this.ignoreException = ignoreException;
@@ -41,7 +41,7 @@ public class ScriptFieldsContext {
             return name;
         }
 
-        public SearchScript script() {
+        public SearchScript.LeafFactory script() {
             return this.script;
         }
 
diff --git a/core/src/main/java/org/elasticsearch/search/fetch/subphase/ScriptFieldsFetchSubPhase.java b/core/src/main/java/org/elasticsearch/search/fetch/subphase/ScriptFieldsFetchSubPhase.java
index 6bed20e6b3..61c1c802de 100644
--- a/core/src/main/java/org/elasticsearch/search/fetch/subphase/ScriptFieldsFetchSubPhase.java
+++ b/core/src/main/java/org/elasticsearch/search/fetch/subphase/ScriptFieldsFetchSubPhase.java
@@ -18,7 +18,7 @@
  */
 package org.elasticsearch.search.fetch.subphase;
 
-import org.elasticsearch.script.LeafSearchScript;
+import org.elasticsearch.script.SearchScript;
 import org.elasticsearch.search.SearchHitField;
 import org.elasticsearch.search.fetch.FetchSubPhase;
 import org.elasticsearch.search.internal.SearchContext;
@@ -40,9 +40,9 @@ public final class ScriptFieldsFetchSubPhase implements FetchSubPhase {
         for (ScriptFieldsContext.ScriptField scriptField : context.scriptFields().fields()) {
             /* Because this is called once per document we end up creating new ScriptDocValues for every document which is important because
              * the values inside ScriptDocValues might be reused for different documents (Dates do this). */
-            LeafSearchScript leafScript;
+            SearchScript leafScript;
             try {
-                leafScript = scriptField.script().getLeafSearchScript(hitContext.readerContext());
+                leafScript = scriptField.script().newInstance(hitContext.readerContext());
             } catch (IOException e1) {
                 throw new IllegalStateException("Failed to load script", e1);
             }
diff --git a/core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java b/core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java
index 4ac8b023d7..18b15bfec2 100644
--- a/core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java
+++ b/core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java
@@ -45,9 +45,7 @@ import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryParseContext;
 import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.index.query.QueryShardException;
-import org.elasticsearch.script.LeafSearchScript;
 import org.elasticsearch.script.Script;
-import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.SearchScript;
 import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.MultiValueMode;
@@ -242,7 +240,7 @@ public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> {
 
     @Override
     public SortFieldAndFormat build(QueryShardContext context) throws IOException {
-        final SearchScript searchScript = context.getSearchScript(script, SearchScript.CONTEXT);
+        final SearchScript.LeafFactory searchScript = context.getSearchScript(script, SearchScript.CONTEXT);
 
         MultiValueMode valueMode = null;
         if (sortMode != null) {
@@ -258,10 +256,10 @@ public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> {
         switch (type) {
             case STRING:
                 fieldComparatorSource = new BytesRefFieldComparatorSource(null, null, valueMode, nested) {
-                    LeafSearchScript leafScript;
+                    SearchScript leafScript;
                     @Override
                     protected SortedBinaryDocValues getValues(LeafReaderContext context) throws IOException {
-                        leafScript = searchScript.getLeafSearchScript(context);
+                        leafScript = searchScript.newInstance(context);
                         final BinaryDocValues values = new AbstractBinaryDocValues() {
                             final BytesRefBuilder spare = new BytesRefBuilder();
                             @Override
@@ -285,10 +283,10 @@ public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> {
                 break;
             case NUMBER:
                 fieldComparatorSource = new DoubleValuesComparatorSource(null, Double.MAX_VALUE, valueMode, nested) {
-                    LeafSearchScript leafScript;
+                    SearchScript leafScript;
                     @Override
                     protected SortedNumericDoubleValues getValues(LeafReaderContext context) throws IOException {
-                        leafScript = searchScript.getLeafSearchScript(context);
+                        leafScript = searchScript.newInstance(context);
                         final NumericDoubleValues values = new NumericDoubleValues() {
                             @Override
                             public boolean advanceExact(int doc) throws IOException {