Add ShardId and Index to SuggestionContext

Suggesters might need access to the shard they run on as well as the index they operate on. This patch adds indexname and shard ID to the SuggestionContext Closes #3199
elastic · Sep 27, 2013 · fe260a4 · fe260a4
1 parent bcf7a0e
commit fe260a4
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 7 deletions.
diff --git a/src/main/java/org/elasticsearch/action/suggest/TransportSuggestAction.java b/src/main/java/org/elasticsearch/action/suggest/TransportSuggestAction.java
@@ -158,7 +158,7 @@ protected ShardSuggestResponse shardOperation(ShardSuggestRequest request) throw
                 if (parser.nextToken() != XContentParser.Token.START_OBJECT) {
                     throw new ElasticSearchIllegalArgumentException("suggest content missing");
                 }
-                final SuggestionSearchContext context = suggestPhase.parseElement().parseInternal(parser, indexService.mapperService());
+                final SuggestionSearchContext context = suggestPhase.parseElement().parseInternal(parser, indexService.mapperService(), request.index(), request.shardId());
                 final Suggest result = suggestPhase.execute(context, searcher.reader());
                 return new ShardSuggestResponse(request.index(), request.shardId(), result);
             }

diff --git a/src/main/java/org/elasticsearch/search/suggest/SuggestParseElement.java b/src/main/java/org/elasticsearch/search/suggest/SuggestParseElement.java
@@ -18,8 +18,6 @@
  */
 package org.elasticsearch.search.suggest;
 
-import java.io.IOException;
-
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.ElasticSearchIllegalArgumentException;
 import org.elasticsearch.common.inject.Inject;
@@ -29,14 +27,15 @@
 import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.search.suggest.SuggestionSearchContext.SuggestionContext;
 
+import java.io.IOException;
 import java.util.Map;
 
 import static com.google.common.collect.Maps.newHashMap;
 
 /**
  *
  */
-public class SuggestParseElement implements SearchParseElement {
+public final class SuggestParseElement implements SearchParseElement {
     private Suggesters suggesters;
 
     @Inject
@@ -46,11 +45,11 @@ public SuggestParseElement(Suggesters suggesters) {
 
     @Override
     public void parse(XContentParser parser, SearchContext context) throws Exception {
-        SuggestionSearchContext suggestionSearchContext = parseInternal(parser, context.mapperService());
+        SuggestionSearchContext suggestionSearchContext = parseInternal(parser, context.mapperService(), context.shardTarget().index(), context.shardTarget().shardId());
         context.suggest(suggestionSearchContext);
     }
 
-    public SuggestionSearchContext parseInternal(XContentParser parser, MapperService mapperService) throws IOException {
+    public SuggestionSearchContext parseInternal(XContentParser parser, MapperService mapperService, String index, int shardId) throws IOException {
         SuggestionSearchContext suggestionSearchContext = new SuggestionSearchContext();
         BytesRef globalText = null;
         String fieldName = null;
@@ -100,7 +99,8 @@ public SuggestionSearchContext parseInternal(XContentParser parser, MapperServic
             for (Map.Entry<String, SuggestionContext> entry : suggestionContexts.entrySet()) {
                 String suggestionName = entry.getKey();
                 SuggestionContext suggestionContext = entry.getValue();
-
+                suggestionContext.setShard(shardId);
+                suggestionContext.setIndex(index);
                 SuggestUtils.verifySuggestion(mapperService, globalText, suggestionContext);
                 suggestionSearchContext.addSuggestion(suggestionName, suggestionContext);
             }

diff --git a/src/main/java/org/elasticsearch/search/suggest/SuggestPhase.java b/src/main/java/org/elasticsearch/search/suggest/SuggestPhase.java
@@ -95,3 +95,4 @@ public Suggest execute(SuggestionSearchContext suggest, IndexReader reader) {
         }
     }
 }
+
diff --git a/src/main/java/org/elasticsearch/search/suggest/SuggestionSearchContext.java b/src/main/java/org/elasticsearch/search/suggest/SuggestionSearchContext.java
@@ -46,6 +46,8 @@ public static class SuggestionContext {
         private Analyzer analyzer;
         private int size = 5;
         private int shardSize = -1;
+        private int shardId;
+        private String index;
 
         public BytesRef getText() {
             return text;
@@ -100,6 +102,24 @@ public void setShardSize(int shardSize) {
             }
             this.shardSize = shardSize;
         }
+
+        public void setShard(int shardId) {
+            this.shardId = shardId;
+        }
+
+        public void setIndex(String index) {
+            this.index = index;
+        }
+
+        public String getIndex() {
+            return index;
+        }
+
+        public int getShard() {
+            return shardId;
+        }
     }
 
+
+
 }