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

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() {

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()));

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());
     }
 
     /**

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

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);

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

core/src/main/java/org/elasticsearch/script/LeafSearchScript.java

@@ -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();
-}

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;
+    }
 }

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);
 }

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()));
             }
         }

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) {

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 {

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);

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()));

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;
             }

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;
     }

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;
     }

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;
     }

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;
     }

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;
         }
 

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);
             }

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 {

core/src/test/java/org/elasticsearch/common/lucene/search/function/ScriptScoreFunctionTests.java

@@ -21,7 +21,6 @@ package org.elasticsearch.common.lucene.search.function;
 
 import org.apache.lucene.index.LeafReaderContext;
 import org.elasticsearch.script.GeneralScriptException;
-import org.elasticsearch.script.LeafSearchScript;
 import org.elasticsearch.script.SearchScript;
 import org.elasticsearch.test.ESTestCase;
 
@@ -33,10 +32,15 @@ public class ScriptScoreFunctionTests extends ESTestCase {
      */
     public void testScriptScoresReturnsNaN() throws IOException {
         // script that always returns NaN
-        ScoreFunction scoreFunction = new ScriptScoreFunction(mockScript("Double.NaN"), new SearchScript() {
+        ScoreFunction scoreFunction = new ScriptScoreFunction(mockScript("Double.NaN"), new SearchScript.LeafFactory() {
             @Override
-            public LeafSearchScript getLeafSearchScript(LeafReaderContext context) throws IOException {
-                return () -> Double.NaN;
+            public SearchScript newInstance(LeafReaderContext context) throws IOException {
+                return new SearchScript(null, null, null) {
+                    @Override
+                    public double runAsDouble() {
+                        return Double.NaN;
+                    }
+                };
             }
             
             @Override

core/src/test/java/org/elasticsearch/script/ScriptContextTests.java

@@ -26,6 +26,15 @@ public class ScriptContextTests extends ESTestCase {
     public interface TwoNewInstance {
         String newInstance(int foo, int bar);
         String newInstance(int foo);
+
+        interface StatefulFactory {
+            TwoNewInstance newFactory();
+        }
+    }
+
+    public interface TwoNewFactory {
+        String newFactory(int foo, int bar);
+        String newFactory(int foo);
     }
 
     public interface MissingNewInstance {
@@ -40,6 +49,16 @@ public class ScriptContextTests extends ESTestCase {
         }
     }
 
+    public interface DummyStatefulScript {
+        int execute(int foo);
+        interface StatefulFactory {
+            DummyStatefulScript newInstance();
+        }
+        interface Factory {
+            StatefulFactory newFactory();
+        }
+    }
+
     public void testTwoNewInstanceMethods() {
         IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () ->
             new ScriptContext<>("test", TwoNewInstance.class));
@@ -47,10 +66,24 @@ public class ScriptContextTests extends ESTestCase {
             + TwoNewInstance.class.getName() + "] for script context [test]", e.getMessage());
     }
 
+    public void testTwoNewFactoryMethods() {
+        IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () ->
+            new ScriptContext<>("test", TwoNewFactory.class));
+        assertEquals("Cannot have multiple newFactory methods on FactoryType class ["
+            + TwoNewFactory.class.getName() + "] for script context [test]", e.getMessage());
+    }
+
+    public void testTwoNewInstanceStatefulFactoryMethods() {
+        IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () ->
+            new ScriptContext<>("test", TwoNewInstance.StatefulFactory.class));
+        assertEquals("Cannot have multiple newInstance methods on StatefulFactoryType class ["
+            + TwoNewInstance.class.getName() + "] for script context [test]", e.getMessage());
+    }
+
     public void testMissingNewInstanceMethod() {
         IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () ->
             new ScriptContext<>("test", MissingNewInstance.class));
-        assertEquals("Could not find method newInstance on FactoryType class ["
+        assertEquals("Could not find method newInstance or method newFactory on FactoryType class ["
             + MissingNewInstance.class.getName() + "] for script context [test]", e.getMessage());
     }
 
@@ -58,6 +91,15 @@ public class ScriptContextTests extends ESTestCase {
         ScriptContext<?> context = new ScriptContext<>("test", DummyScript.Factory.class);
         assertEquals("test", context.name);
         assertEquals(DummyScript.class, context.instanceClazz);
+        assertNull(context.statefulFactoryClazz);
         assertEquals(DummyScript.Factory.class, context.factoryClazz);
     }
+
+    public void testStatefulFactoryReflection() {
+        ScriptContext<?> context = new ScriptContext<>("test", DummyStatefulScript.Factory.class);
+        assertEquals("test", context.name);
+        assertEquals(DummyStatefulScript.class, context.instanceClazz);
+        assertEquals(DummyStatefulScript.StatefulFactory.class, context.statefulFactoryClazz);
+        assertEquals(DummyStatefulScript.Factory.class, context.factoryClazz);
+    }
 }

core/src/test/java/org/elasticsearch/search/aggregations/support/ScriptValuesTests.java

@@ -22,7 +22,7 @@ package org.elasticsearch.search.aggregations.support;
 import com.carrotsearch.randomizedtesting.generators.RandomStrings;
 import org.apache.lucene.search.Scorer;
 import org.apache.lucene.util.BytesRef;
-import org.elasticsearch.script.LeafSearchScript;
+import org.elasticsearch.script.SearchScript;
 import org.elasticsearch.search.aggregations.support.values.ScriptBytesValues;
 import org.elasticsearch.search.aggregations.support.values.ScriptDoubleValues;
 import org.elasticsearch.search.aggregations.support.values.ScriptLongValues;
@@ -30,16 +30,16 @@ import org.elasticsearch.test.ESTestCase;
 
 import java.io.IOException;
 import java.util.Arrays;
-import java.util.Map;
 
 public class ScriptValuesTests extends ESTestCase {
 
-    private static class FakeSearchScript implements LeafSearchScript {
+    private static class FakeSearchScript extends SearchScript {
 
         private final Object[][] values;
         int index;
 
         FakeSearchScript(Object[][] values) {
+            super(null, null, null);
             this.values = values;
             index = -1;
         }
@@ -67,10 +67,6 @@ public class ScriptValuesTests extends ESTestCase {
             index = doc;
         }
 
-        @Override
-        public void setSource(Map<String, Object> source) {
-        }
-
         @Override
         public long runAsLong() {
             throw new UnsupportedOperationException();

core/src/test/java/org/elasticsearch/search/functionscore/ExplainableScriptIT.java

@@ -30,7 +30,6 @@ import org.elasticsearch.index.fielddata.ScriptDocValues;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.plugins.ScriptPlugin;
 import org.elasticsearch.script.ExplainableSearchScript;
-import org.elasticsearch.script.LeafSearchScript;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.ScriptEngine;
@@ -78,9 +77,9 @@ public class ExplainableScriptIT extends ESIntegTestCase {
                 public <T> T compile(String scriptName, String scriptSource, ScriptContext<T> context, Map<String, String> params) {
                     assert scriptSource.equals("explainable_script");
                     assert context == SearchScript.CONTEXT;
-                    SearchScript.Factory factory = (p, lookup) -> new SearchScript() {
+                    SearchScript.Factory factory = (p, lookup) -> new SearchScript.LeafFactory() {
                         @Override
-                        public LeafSearchScript getLeafSearchScript(LeafReaderContext context) throws IOException {
+                        public SearchScript newInstance(LeafReaderContext context) throws IOException {
                             return new MyScript(lookup.doc().getLeafDocLookup(context));
                         }
                         @Override
@@ -94,10 +93,11 @@ public class ExplainableScriptIT extends ESIntegTestCase {
         }
     }
 
-    static class MyScript implements ExplainableSearchScript {
+    static class MyScript extends SearchScript implements ExplainableSearchScript {
         LeafDocLookup docLookup;
 
         MyScript(LeafDocLookup docLookup) {
+            super(null, null, null);
             this.docLookup = docLookup;
         }
 

modules/lang-expression/src/main/java/org/elasticsearch/script/expression/ExpressionScriptEngine.java

@@ -111,7 +111,7 @@ public class ExpressionScriptEngine extends AbstractComponent implements ScriptE
         throw new IllegalArgumentException("painless does not know how to handle context [" + context.name + "]");
     }
 
-    private SearchScript newSearchScript(Expression expr, SearchLookup lookup, @Nullable Map<String, Object> vars) {
+    private SearchScript.LeafFactory newSearchScript(Expression expr, SearchLookup lookup, @Nullable Map<String, Object> vars) {
         MapperService mapper = lookup.doc().mapperService();
         // NOTE: if we need to do anything complicated with bindings in the future, we can just extend Bindings,
         // instead of complicating SimpleBindings (which should stay simple)

modules/lang-expression/src/main/java/org/elasticsearch/script/expression/ExpressionSearchScript.java

@@ -27,17 +27,15 @@ import org.apache.lucene.search.DoubleValues;
 import org.apache.lucene.search.DoubleValuesSource;
 import org.apache.lucene.search.Scorer;
 import org.elasticsearch.script.GeneralScriptException;
-import org.elasticsearch.script.LeafSearchScript;
 import org.elasticsearch.script.SearchScript;
 
 import java.io.IOException;
-import java.util.Map;
 
 /**
  * A bridge to evaluate an {@link Expression} against {@link Bindings} in the context
  * of a {@link SearchScript}.
  */
-class ExpressionSearchScript implements SearchScript {
+class ExpressionSearchScript implements SearchScript.LeafFactory {
 
     final Expression exprScript;
     final SimpleBindings bindings;
@@ -62,13 +60,13 @@ class ExpressionSearchScript implements SearchScript {
 
 
     @Override
-    public LeafSearchScript getLeafSearchScript(final LeafReaderContext leaf) throws IOException {
-        return new LeafSearchScript() {
+    public SearchScript newInstance(final LeafReaderContext leaf) throws IOException {
+        return new SearchScript(null, null, null) {
             // Fake the scorer until setScorer is called.
             DoubleValues values = source.getValues(leaf, new DoubleValues() {
                 @Override
                 public double doubleValue() throws IOException {
-                    return Double.NaN;
+                    return getScore();
                 }
 
                 @Override
@@ -76,59 +74,31 @@ class ExpressionSearchScript implements SearchScript {
                     return true;
                 }
             });
-            double evaluate() {
-                try {
-                    return values.doubleValue();
-                } catch (Exception exception) {
-                    throw new GeneralScriptException("Error evaluating " + exprScript, exception);
-                }
-            }
 
             @Override
-            public Object run() { return Double.valueOf(evaluate()); }
+            public Object run() { return Double.valueOf(runAsDouble()); }
 
             @Override
-            public long runAsLong() { return (long)evaluate(); }
+            public long runAsLong() { return (long)runAsDouble(); }
 
             @Override
-            public double runAsDouble() { return evaluate(); }
-
-            @Override
-            public void setDocument(int d) {
-                docid = d;
+            public double runAsDouble() {
                 try {
-                    values.advanceExact(d);
-                } catch (IOException e) {
-                    throw new IllegalStateException("Can't advance to doc using " + exprScript, e);
+                    return values.doubleValue();
+                } catch (Exception exception) {
+                    throw new GeneralScriptException("Error evaluating " + exprScript, exception);
                 }
             }
 
             @Override
-            public void setScorer(Scorer s) {
-                scorer = s;
+            public void setDocument(int d) {
                 try {
-                    // We have a new binding for the scorer so we need to reset the values
-                    values = source.getValues(leaf, new DoubleValues() {
-                        @Override
-                        public double doubleValue() throws IOException {
-                            return scorer.score();
-                        }
-
-                        @Override
-                        public boolean advanceExact(int doc) throws IOException {
-                            return true;
-                        }
-                    });
+                    values.advanceExact(d);
                 } catch (IOException e) {
-                    throw new IllegalStateException("Can't get values using " + exprScript, e);
+                    throw new IllegalStateException("Can't advance to doc using " + exprScript, e);
                 }
             }
 
-            @Override
-            public void setSource(Map<String, Object> source) {
-                // noop: expressions don't use source data
-            }
-
             @Override
             public void setNextAggregationValue(Object value) {
                 // _value isn't used in script if specialValue == null

modules/lang-expression/src/test/java/org/elasticsearch/script/expression/ExpressionTests.java

@@ -41,9 +41,9 @@ public class ExpressionTests extends ESSingleNodeTestCase {
         lookup = new SearchLookup(index.mapperService(), index.fieldData(), null);
     }
 
-    private SearchScript compile(String expression) {
+    private SearchScript.LeafFactory compile(String expression) {
         SearchScript.Factory factory = service.compile(null, expression, SearchScript.CONTEXT, Collections.emptyMap());
-        return factory.newInstance(Collections.emptyMap(), lookup);
+        return factory.newFactory(Collections.emptyMap(), lookup);
     }
 
     public void testNeedsScores() {

modules/lang-painless/src/main/java/org/elasticsearch/painless/PainlessScriptEngine.java

@@ -25,13 +25,11 @@ import org.elasticsearch.common.component.AbstractComponent;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.painless.Compiler.Loader;
 import org.elasticsearch.script.ExecutableScript;
-import org.elasticsearch.script.LeafSearchScript;
 import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.ScriptEngine;
 import org.elasticsearch.script.ScriptException;
 import org.elasticsearch.script.SearchScript;
 
-import java.io.IOException;
 import java.lang.reflect.Constructor;
 import java.security.AccessControlContext;
 import java.security.AccessController;
@@ -120,10 +118,10 @@ public final class PainlessScriptEngine extends AbstractComponent implements Scr
         GenericElasticsearchScript painlessScript =
             (GenericElasticsearchScript)compile(contextsToCompilers.get(context), scriptName, scriptSource, params);
         if (context.instanceClazz.equals(SearchScript.class)) {
-            SearchScript.Factory factory = (p, lookup) -> new SearchScript() {
+            SearchScript.Factory factory = (p, lookup) -> new SearchScript.LeafFactory() {
                 @Override
-                public LeafSearchScript getLeafSearchScript(final LeafReaderContext context) throws IOException {
-                    return new ScriptImpl(painlessScript, p, lookup.getLeafSearchLookup(context));
+                public SearchScript newInstance(final LeafReaderContext context) {
+                    return new ScriptImpl(painlessScript, p, lookup, context);
                 }
                 @Override
                 public boolean needsScores() {
@@ -132,7 +130,7 @@ public final class PainlessScriptEngine extends AbstractComponent implements Scr
             };
             return context.factoryClazz.cast(factory);
         } else if (context.instanceClazz.equals(ExecutableScript.class)) {
-            ExecutableScript.Factory factory = (p) -> new ScriptImpl(painlessScript, p, null);
+            ExecutableScript.Factory factory = (p) -> new ScriptImpl(painlessScript, p, null, null);
             return context.factoryClazz.cast(factory);
         }
         throw new IllegalArgumentException("painless does not know how to handle context [" + context.name + "]");

modules/lang-painless/src/main/java/org/elasticsearch/painless/ScriptImpl.java

@@ -19,23 +19,22 @@
 
 package org.elasticsearch.painless;
 
-import org.apache.lucene.search.Scorer;
-import org.elasticsearch.ElasticsearchException;
+import org.apache.lucene.index.LeafReaderContext;
 import org.elasticsearch.script.ExecutableScript;
-import org.elasticsearch.script.LeafSearchScript;
-import org.elasticsearch.search.lookup.LeafDocLookup;
+import org.elasticsearch.script.SearchScript;
 import org.elasticsearch.search.lookup.LeafSearchLookup;
+import org.elasticsearch.search.lookup.SearchLookup;
 
-import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.function.DoubleSupplier;
 import java.util.function.Function;
 
 /**
- * ScriptImpl can be used as either an {@link ExecutableScript} or a {@link LeafSearchScript}
+ * ScriptImpl can be used as either an {@link ExecutableScript} or a {@link SearchScript}
  * to run a previously compiled Painless script.
  */
-final class ScriptImpl implements ExecutableScript, LeafSearchScript {
+final class ScriptImpl extends SearchScript {
 
     /**
      * The Painless script that can be run.
@@ -47,32 +46,16 @@ final class ScriptImpl implements ExecutableScript, LeafSearchScript {
      */
     private final Map<String, Object> variables;
 
-    /**
-     * The lookup is used to access search field values at run-time.
-     */
-    private final LeafSearchLookup lookup;
-
-    /**
-     * the 'doc' object accessed by the script, if available.
-     */
-    private final LeafDocLookup doc;
-
     /**
      * Looks up the {@code _score} from {@link #scorer} if {@code _score} is used, otherwise returns {@code 0.0}.
      */
-    private final ScoreLookup scoreLookup;
+    private final DoubleSupplier scoreLookup;
 
     /**
      * Looks up the {@code ctx} from the {@link #variables} if {@code ctx} is used, otherwise return {@code null}.
      */
     private final Function<Map<String, Object>, Map<?, ?>> ctxLookup;
 
-    /**
-     * Current scorer being used
-     * @see #setScorer(Scorer)
-     */
-    private Scorer scorer;
-
     /**
      * Current _value for aggregation
      * @see #setNextAggregationValue(Object)
@@ -85,112 +68,50 @@ final class ScriptImpl implements ExecutableScript, LeafSearchScript {
      * @param vars The initial variables to run the script with.
      * @param lookup The lookup to allow search fields to be available if this is run as a search script.
      */
-    ScriptImpl(final GenericElasticsearchScript script, final Map<String, Object> vars, final LeafSearchLookup lookup) {
+    ScriptImpl(GenericElasticsearchScript script, Map<String, Object> vars, SearchLookup lookup, LeafReaderContext leafContext) {
+        super(null, lookup, leafContext);
         this.script = script;
-        this.lookup = lookup;
         this.variables = new HashMap<>();
 
         if (vars != null) {
             variables.putAll(vars);
         }
-
-        if (lookup != null) {
-            variables.putAll(lookup.asMap());
-            doc = lookup.doc();
-        } else {
-            doc = null;
+        LeafSearchLookup leafLookup = getLeafLookup();
+        if (leafLookup != null) {
+            variables.putAll(leafLookup.asMap());
         }
 
-        scoreLookup = script.uses$_score() ? ScriptImpl::getScore : scorer -> 0.0;
+        scoreLookup = script.uses$_score() ? this::getScore : () -> 0.0;
         ctxLookup = script.uses$ctx() ? variables -> (Map<?, ?>) variables.get("ctx") : variables -> null;
     }
 
-    /**
-     * Set a variable for the script to be run against.
-     * @param name The variable name.
-     * @param value The variable value.
-     */
+    @Override
+    public Map<String, Object> getParams() {
+        return variables;
+    }
+
     @Override
     public void setNextVar(final String name, final Object value) {
         variables.put(name, value);
     }
 
-    /**
-     * Set the next aggregation value.
-     * @param value Per-document value, typically a String, Long, or Double.
-     */
     @Override
     public void setNextAggregationValue(Object value) {
         this.aggregationValue = value;
     }
 
-    /**
-     * Run the script.
-     * @return The script result.
-     */
     @Override
     public Object run() {
-        return script.execute(variables, scoreLookup.apply(scorer), doc, aggregationValue, ctxLookup.apply(variables));
+        return script.execute(variables, scoreLookup.getAsDouble(), getDoc(), aggregationValue, ctxLookup.apply(variables));
     }
 
-    /**
-     * Run the script.
-     * @return The script result as a double.
-     */
     @Override
     public double runAsDouble() {
         return ((Number)run()).doubleValue();
     }
 
-    /**
-     * Run the script.
-     * @return The script result as a long.
-     */
     @Override
     public long runAsLong() {
         return ((Number)run()).longValue();
     }
-
-    /**
-     * Sets the scorer to be accessible within a script.
-     * @param scorer The scorer used for a search.
-     */
-    @Override
-    public void setScorer(final Scorer scorer) {
-        this.scorer = scorer;
-    }
-
-    /**
-     * Sets the current document.
-     * @param doc The current document.
-     */
-    @Override
-    public void setDocument(final int doc) {
-        if (lookup != null) {
-            lookup.setDocument(doc);
-        }
-    }
-
-    /**
-     * Sets the current source.
-     * @param source The current source.
-     */
-    @Override
-    public void setSource(final Map<String, Object> source) {
-        if (lookup != null) {
-            lookup.source().setSource(source);
-        }
-    }
-
-    private static double getScore(Scorer scorer) {
-        try {
-            return scorer.score();
-        } catch (IOException e) {
-            throw new ElasticsearchException("couldn't lookup score", e);
-        }
-    }
-
-    interface ScoreLookup {
-        double apply(Scorer scorer);
-    }
 }

modules/lang-painless/src/test/java/org/elasticsearch/painless/NeedsScoreTests.java

@@ -43,19 +43,19 @@ public class NeedsScoreTests extends ESSingleNodeTestCase {
         SearchLookup lookup = new SearchLookup(index.mapperService(), index.fieldData(), null);
 
         SearchScript.Factory factory = service.compile(null, "1.2", SearchScript.CONTEXT, Collections.emptyMap());
-        SearchScript ss = factory.newInstance(Collections.emptyMap(), lookup);
+        SearchScript.LeafFactory ss = factory.newFactory(Collections.emptyMap(), lookup);
         assertFalse(ss.needsScores());
 
         factory = service.compile(null, "doc['d'].value", SearchScript.CONTEXT, Collections.emptyMap());
-        ss = factory.newInstance(Collections.emptyMap(), lookup);
+        ss = factory.newFactory(Collections.emptyMap(), lookup);
         assertFalse(ss.needsScores());
 
         factory = service.compile(null, "1/_score", SearchScript.CONTEXT, Collections.emptyMap());
-        ss = factory.newInstance(Collections.emptyMap(), lookup);
+        ss = factory.newFactory(Collections.emptyMap(), lookup);
         assertTrue(ss.needsScores());
 
         factory = service.compile(null, "doc['d'].value * _score", SearchScript.CONTEXT, Collections.emptyMap());
-        ss = factory.newInstance(Collections.emptyMap(), lookup);
+        ss = factory.newFactory(Collections.emptyMap(), lookup);
         assertTrue(ss.needsScores());
     }
 

plugins/examples/script-expert-scoring/src/main/java/org/elasticsearch/example/expertscript/ExpertScriptPlugin.java

@@ -30,7 +30,6 @@ import org.apache.lucene.index.Term;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.plugins.ScriptPlugin;
-import org.elasticsearch.script.LeafSearchScript;
 import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.ScriptEngine;
 import org.elasticsearch.script.SearchScript;
@@ -60,7 +59,7 @@ public class ExpertScriptPlugin extends Plugin implements ScriptPlugin {
             }
             // we use the script "source" as the script identifier
             if ("pure_df".equals(scriptSource)) {
-                SearchScript.Factory factory = (p, lookup) -> new SearchScript() {
+                SearchScript.Factory factory = (p, lookup) -> new SearchScript.LeafFactory() {
                     final String field;
                     final String term;
                     {
@@ -75,13 +74,18 @@ public class ExpertScriptPlugin extends Plugin implements ScriptPlugin {
                     }
 
                     @Override
-                    public LeafSearchScript getLeafSearchScript(LeafReaderContext context) throws IOException {
+                    public SearchScript newInstance(LeafReaderContext context) throws IOException {
                         PostingsEnum postings = context.reader().postings(new Term(field, term));
                         if (postings == null) {
                             // the field and/or term don't exist in this segment, so always return 0
-                            return () -> 0.0d;
+                            return new SearchScript(p, lookup, context) {
+                                @Override
+                                public double runAsDouble() {
+                                    return 0.0d;
+                                }
+                            };
                         }
-                        return new LeafSearchScript() {
+                        return new SearchScript(p, lookup, context) {
                             int currentDocid = -1;
                             @Override
                             public void setDocument(int docid) {

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

@@ -116,7 +116,7 @@ public class MockScriptEngine implements ScriptEngine {
             return new MockExecutableScript(context, script != null ? script : ctx -> source);
         }
 
-        public SearchScript createSearchScript(Map<String, Object> params, SearchLookup lookup) {
+        public SearchScript.LeafFactory createSearchScript(Map<String, Object> params, SearchLookup lookup) {
             Map<String, Object> context = new HashMap<>();
             if (options != null) {
                 context.putAll(options); // TODO: remove this once scripts know to look for options under options key
@@ -151,7 +151,7 @@ public class MockScriptEngine implements ScriptEngine {
         }
     }
 
-    public class MockSearchScript implements SearchScript {
+    public class MockSearchScript implements SearchScript.LeafFactory {
 
         private final Function<Map<String, Object>, Object> script;
         private final Map<String, Object> vars;
@@ -164,7 +164,7 @@ public class MockScriptEngine implements ScriptEngine {
         }
 
         @Override
-        public LeafSearchScript getLeafSearchScript(LeafReaderContext context) throws IOException {
+        public SearchScript newInstance(LeafReaderContext context) throws IOException {
             LeafSearchLookup leafLookup = lookup.getLeafSearchLookup(context);
 
             Map<String, Object> ctx = new HashMap<>(leafLookup.asMap());
@@ -172,7 +172,7 @@ public class MockScriptEngine implements ScriptEngine {
                 ctx.putAll(vars);
             }
 
-            return new LeafSearchScript() {
+            return new SearchScript(vars, lookup, context) {
                 @Override
                 public Object run() {
                     return script.apply(ctx);
@@ -202,12 +202,6 @@ public class MockScriptEngine implements ScriptEngine {
                 public void setDocument(int doc) {
                     leafLookup.setDocument(doc);
                 }
-
-                @Override
-                public void setSource(Map<String, Object> source) {
-                    leafLookup.source().setSource(source);
-                }
-
             };
         }