Add initial implementation of script_source

src/main/java/org/elasticsearch/action/search/SearchRequestBuilder.java

@@ -447,6 +447,40 @@ public SearchRequestBuilder addScriptField(String name, String lang, String scri
         return this;
     }
 
+    /**
+     * Sets a script based source to load.
+     *
+     * @param script The script to use
+     */
+    public SearchRequestBuilder setScriptSource(String script) {
+        sourceBuilder().scriptSource(script);
+        return this;
+    }
+
+    /**
+     * Sets a script based source to load.
+     *
+     * @param script The script to use
+     * @param params Parameters that the script can use.
+     */
+    public SearchRequestBuilder setScriptSource(String script, Map<String, Object> params) {
+        sourceBuilder().scriptSource(script, params);
+        return this;
+    }
+
+    /**
+     * Sets a script based source to load.
+     *
+     * @param lang   The language of the script
+     * @param script The script to use
+     * @param params Parameters that the script can use (can be <tt>null</tt>).
+     */
+    public SearchRequestBuilder setScriptSource(String lang, String script, Map<String, Object> params) {
+        sourceBuilder().scriptSource(lang, script, params);
+        return this;
+    }
+
+
     /**
      * Adds a sort against the given field name and the sort ordering.
      *

src/main/java/org/elasticsearch/search/SearchModule.java

@@ -29,9 +29,11 @@
 import org.elasticsearch.search.facet.FacetModule;
 import org.elasticsearch.search.fetch.FetchPhase;
 import org.elasticsearch.search.fetch.explain.ExplainFetchSubPhase;
+import org.elasticsearch.search.fetch.extractfields.ExtractFieldsFetchSubPhase;
 import org.elasticsearch.search.fetch.matchedfilters.MatchedFiltersFetchSubPhase;
 import org.elasticsearch.search.fetch.partial.PartialFieldsFetchSubPhase;
 import org.elasticsearch.search.fetch.script.ScriptFieldsFetchSubPhase;
+import org.elasticsearch.search.fetch.scriptsource.ScriptSourceFetchSubPhase;
 import org.elasticsearch.search.fetch.version.VersionFetchSubPhase;
 import org.elasticsearch.search.highlight.HighlightPhase;
 import org.elasticsearch.search.query.QueryPhase;
@@ -56,6 +58,8 @@ protected void configure() {
         bind(FetchPhase.class).asEagerSingleton();
         bind(ExplainFetchSubPhase.class).asEagerSingleton();
         bind(ScriptFieldsFetchSubPhase.class).asEagerSingleton();
+        bind(ScriptSourceFetchSubPhase.class).asEagerSingleton();
+        bind(ExtractFieldsFetchSubPhase.class).asEagerSingleton();
         bind(PartialFieldsFetchSubPhase.class).asEagerSingleton();
         bind(VersionFetchSubPhase.class).asEagerSingleton();
         bind(MatchedFiltersFetchSubPhase.class).asEagerSingleton();

src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java

@@ -95,6 +95,7 @@ public static HighlightBuilder highlight() {
 
     private List<String> fieldNames;
     private List<ScriptField> scriptFields;
+    private ScriptSource scriptSource;
     private List<PartialField> partialFields;
 
     private List<AbstractFacetBuilder> facets;
@@ -481,6 +482,38 @@ public SearchSourceBuilder scriptField(String name, String lang, String script,
         return this;
     }
 
+    /**
+     * Sets a script source.
+     *
+     * @param script The script
+     */
+    public SearchSourceBuilder scriptSource(String script) {
+        return scriptSource(null, script, null);
+    }
+
+    /**
+     * Sets a script source.
+     *
+     * @param script The script to execute
+     * @param params The script parameters
+     */
+    public SearchSourceBuilder scriptSource(String script, Map<String, Object> params) {
+        return scriptSource(null, script, params);
+    }
+
+    /**
+     * Sets a script source.
+     *
+     * @param lang   The language of the script
+     * @param script The script to execute
+     * @param params The script parameters (can be <tt>null</tt>)
+     * @return
+     */
+    public SearchSourceBuilder scriptSource(String lang, String script, Map<String, Object> params) {
+        scriptSource = new ScriptSource(lang, script, params);
+        return this;
+    }
+
     /**
      * Adds a partial field based on _source, with an "include" and/or "exclude" set which can include simple wildcard
      * elements.
@@ -666,6 +699,19 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             builder.endObject();
         }
 
+        if (scriptSource != null) {
+            builder.startObject("script_source");
+            builder.field("script", scriptSource.script());
+            if (scriptSource.lang() != null) {
+                builder.field("lang", scriptSource.lang());
+            }
+            if (scriptSource.params() != null) {
+                builder.field("params");
+                builder.map(scriptSource.params());
+            }
+            builder.endObject();
+        }
+
         if (sorts != null) {
             builder.startArray("sort");
             for (SortBuilder sort : sorts) {
@@ -780,4 +826,28 @@ public String[] excludes() {
             return excludes;
         }
     }
+
+    private static class ScriptSource {
+        private final String script;
+        private final String lang;
+        private final Map<String, Object> params;
+
+        private ScriptSource(String lang, String script, Map<String, Object> params) {
+            this.lang = lang;
+            this.script = script;
+            this.params = params;
+        }
+
+        public String script() {
+            return script;
+        }
+
+        public String lang() {
+            return this.lang;
+        }
+
+        public Map<String, Object> params() {
+            return params;
+        }
+    }
 }

src/main/java/org/elasticsearch/search/fetch/FetchPhase.java

@@ -20,7 +20,6 @@
 package org.elasticsearch.search.fetch;
 
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Lists;
 import org.apache.lucene.document.Document;
 import org.apache.lucene.document.Fieldable;
 import org.apache.lucene.index.IndexReader;
@@ -44,9 +43,11 @@
 import org.elasticsearch.search.SearchParseElement;
 import org.elasticsearch.search.SearchPhase;
 import org.elasticsearch.search.fetch.explain.ExplainFetchSubPhase;
+import org.elasticsearch.search.fetch.extractfields.ExtractFieldsFetchSubPhase;
 import org.elasticsearch.search.fetch.matchedfilters.MatchedFiltersFetchSubPhase;
 import org.elasticsearch.search.fetch.partial.PartialFieldsFetchSubPhase;
 import org.elasticsearch.search.fetch.script.ScriptFieldsFetchSubPhase;
+import org.elasticsearch.search.fetch.scriptsource.ScriptSourceFetchSubPhase;
 import org.elasticsearch.search.fetch.version.VersionFetchSubPhase;
 import org.elasticsearch.search.highlight.HighlightPhase;
 import org.elasticsearch.search.internal.InternalSearchHit;
@@ -68,9 +69,10 @@ public class FetchPhase implements SearchPhase {
     private final FetchSubPhase[] fetchSubPhases;
 
     @Inject
-    public FetchPhase(HighlightPhase highlightPhase, ScriptFieldsFetchSubPhase scriptFieldsPhase, PartialFieldsFetchSubPhase partialFieldsPhase,
+    public FetchPhase(HighlightPhase highlightPhase, ScriptSourceFetchSubPhase scriptSourcePhase, ScriptFieldsFetchSubPhase scriptFieldsPhase,
+                      ExtractFieldsFetchSubPhase extractFieldsPhase, PartialFieldsFetchSubPhase partialFieldsPhase,
                       MatchedFiltersFetchSubPhase matchFiltersPhase, ExplainFetchSubPhase explainPhase, VersionFetchSubPhase versionPhase) {
-        this.fetchSubPhases = new FetchSubPhase[]{scriptFieldsPhase, partialFieldsPhase, matchFiltersPhase, explainPhase, highlightPhase, versionPhase};
+        this.fetchSubPhases = new FetchSubPhase[]{scriptSourcePhase, extractFieldsPhase, scriptFieldsPhase, partialFieldsPhase, matchFiltersPhase, explainPhase, highlightPhase, versionPhase};
     }
 
     @Override
@@ -89,7 +91,6 @@ public void preProcess(SearchContext context) {
 
     public void execute(SearchContext context) {
         ResetFieldSelector fieldSelector;
-        List<String> extractFieldNames = null;
         boolean sourceRequested = false;
         if (!context.hasFieldNames()) {
             if (context.hasPartialFields()) {
@@ -126,27 +127,24 @@ public void execute(SearchContext context) {
                     }
                     fieldSelectorMapper.add(x);
                 } else {
-                    if (extractFieldNames == null) {
-                        extractFieldNames = Lists.newArrayList();
-                    }
-                    extractFieldNames.add(fieldName);
+                    context.extractFieldNames().add(fieldName);
                 }
             }
 
             if (loadAllStored) {
-                if (sourceRequested || extractFieldNames != null) {
+                if (sourceRequested || context.hasExtractFieldNames()) {
                     fieldSelector = null; // load everything, including _source
                 } else {
                     fieldSelector = AllButSourceFieldSelector.INSTANCE;
                 }
             } else if (fieldSelectorMapper != null) {
                 // we are asking specific stored fields, just add the UID and be done
                 fieldSelectorMapper.add(UidFieldMapper.NAME);
-                if (extractFieldNames != null || sourceRequested) {
+                if (context.hasExtractFieldNames() || sourceRequested) {
                     fieldSelectorMapper.add(SourceFieldMapper.NAME);
                 }
                 fieldSelector = fieldSelectorMapper;
-            } else if (extractFieldNames != null || sourceRequested) {
+            } else if (context.hasExtractFieldNames() || sourceRequested) {
                 fieldSelector = new UidAndSourceFieldSelector();
             } else {
                 fieldSelector = UidFieldSelector.INSTANCE;
@@ -169,7 +167,7 @@ public void execute(SearchContext context) {
 
             // get the version
 
-            InternalSearchHit searchHit = new InternalSearchHit(docId, uid.id(), uid.type(), sourceRequested ? source : null, null);
+            InternalSearchHit searchHit = new InternalSearchHit(docId, uid.id(), uid.type(), null);
             hits[index] = searchHit;
 
             for (Object oField : doc.getFields()) {
@@ -219,37 +217,24 @@ public void execute(SearchContext context) {
             IndexReader subReader = context.searcher().subReaders()[readerIndex];
             int subDoc = docId - context.searcher().docStarts()[readerIndex];
 
-            // go over and extract fields that are not mapped / stored
             context.lookup().setNextReader(subReader);
             context.lookup().setNextDocId(subDoc);
             if (source != null) {
                 context.lookup().source().setNextSource(new BytesArray(source));
             }
-            if (extractFieldNames != null) {
-                for (String extractFieldName : extractFieldNames) {
-                    Object value = context.lookup().source().extractValue(extractFieldName);
-                    if (value != null) {
-                        if (searchHit.fieldsOrNull() == null) {
-                            searchHit.fields(new HashMap<String, SearchHitField>(2));
-                        }
-
-                        SearchHitField hitField = searchHit.fields().get(extractFieldName);
-                        if (hitField == null) {
-                            hitField = new InternalSearchHitField(extractFieldName, new ArrayList<Object>(2));
-                            searchHit.fields().put(extractFieldName, hitField);
-                        }
-                        hitField.values().add(value);
-                    }
-                }
-            }
 
+            FetchSubPhase.HitContext hitContext = new FetchSubPhase.HitContext(source);
             for (FetchSubPhase fetchSubPhase : fetchSubPhases) {
-                FetchSubPhase.HitContext hitContext = new FetchSubPhase.HitContext();
                 if (fetchSubPhase.hitExecutionNeeded(context)) {
                     hitContext.reset(searchHit, subReader, subDoc, context.searcher().getIndexReader(), docId, doc);
                     fetchSubPhase.hitExecute(context, hitContext);
                 }
             }
+
+            // Store source if needed
+            if (sourceRequested && hitContext.source() != null) {
+                searchHit.source(hitContext.source());
+            }
         }
 
         for (FetchSubPhase fetchSubPhase : fetchSubPhases) {

src/main/java/org/elasticsearch/search/fetch/FetchSubPhase.java

@@ -23,10 +23,16 @@
 import org.apache.lucene.document.Document;
 import org.apache.lucene.index.IndexReader;
 import org.elasticsearch.ElasticSearchException;
+import org.elasticsearch.common.io.stream.BytesStreamOutput;
+import org.elasticsearch.common.io.stream.CachedStreamOutput;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.common.xcontent.XContentFactory;
 import org.elasticsearch.search.SearchParseElement;
 import org.elasticsearch.search.internal.InternalSearchHit;
 import org.elasticsearch.search.internal.SearchContext;
+import org.elasticsearch.search.lookup.SourceLookup;
 
+import java.io.IOException;
 import java.util.Map;
 
 /**
@@ -42,6 +48,13 @@ public static class HitContext {
         private int docId;
         private Document doc;
         private Map<String, Object> cache;
+        private byte[] source;
+        private Map<String, Object> sourceAsMap;
+
+        public HitContext(byte[] source) {
+            this.source = source;
+            this.sourceAsMap = null;
+        }
 
         public void reset(InternalSearchHit hit, IndexReader reader, int docId, IndexReader topLevelReader, int topLevelDocId, Document doc) {
             this.hit = hit;
@@ -82,6 +95,49 @@ public Map<String, Object> cache() {
             }
             return cache;
         }
+
+        public byte[] source() {
+            if (source != null) {
+                return source;
+            }
+            if (sourceAsMap != null) {
+                CachedStreamOutput.Entry cachedEntry = CachedStreamOutput.popEntry();
+                try {
+                    BytesStreamOutput streamOutput = cachedEntry.bytes();
+                    // TODO: Is there a better way to figure out Content Type?
+                    XContentBuilder builder = XContentFactory.jsonBuilder(streamOutput).map(sourceAsMap);
+                    builder.close();
+                    source = streamOutput.bytes().copyBytesArray().toBytes();
+                    return source;
+                } catch (IOException ex) {
+                    throw new ElasticSearchException("Cannot serialize map " + sourceAsMap, ex);
+                } finally {
+                    CachedStreamOutput.pushEntry(cachedEntry);
+                }
+            }
+            return null;
+        }
+
+        public Map<String, Object> sourceAsMap() {
+            if (sourceAsMap != null) {
+                return sourceAsMap;
+            }
+            if (source != null) {
+                sourceAsMap = SourceLookup.sourceAsMap(source, 0, source.length);
+                return sourceAsMap;
+            }
+            return null;
+        }
+
+        public void source(byte[] source) {
+            this.source = source;
+            this.sourceAsMap = null;
+        }
+
+        public void sourceAsMap(Map<String, Object>  sourceAsMap) {
+            this.source = null;
+            this.sourceAsMap = sourceAsMap;
+        }
     }
 
     Map<String, ? extends SearchParseElement> parseElements();