add query context builders

5 files changed, 196 insertions, 156 deletions

diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryContextMapping.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryContextMapping.java
index 62fec54a33..23c9ca730b 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryContextMapping.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryContextMapping.java
@@ -153,11 +153,11 @@ public class CategoryContextMapping extends ContextMapping {
         Token token = parser.nextToken();
         if (token == Token.START_OBJECT || token == Token.VALUE_STRING) {
             CategoryQueryContext parse = CategoryQueryContext.parse(parser);
-            queryContexts.add(new QueryContext(parse.context.toString(), parse.boost, parse.isPrefix));
+            queryContexts.add(new QueryContext(parse.getCategory().toString(), parse.getBoost(), parse.isPrefix()));
         } else if (token == Token.START_ARRAY) {
             while (parser.nextToken() != Token.END_ARRAY) {
                 CategoryQueryContext parse = CategoryQueryContext.parse(parser);
-                queryContexts.add(new QueryContext(parse.context.toString(), parse.boost, parse.isPrefix));
+                queryContexts.add(new QueryContext(parse.getCategory().toString(), parse.getBoost(), parse.isPrefix()));
             }
         }
         return queryContexts;
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryQueryContext.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryQueryContext.java
index 977587ffed..ee2655ebdd 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryQueryContext.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryQueryContext.java
@@ -33,78 +33,105 @@ import static org.elasticsearch.search.suggest.completion.context.CategoryContex
  * Defines the query context for {@link CategoryContextMapping}
  */
 public final class CategoryQueryContext implements ToXContent {
+    private final CharSequence category;
+    private final boolean isPrefix;
+    private final int boost;
 
-    public CharSequence context;
-
-    public boolean isPrefix = false;
-
-    public int boost = 1;
+    private CategoryQueryContext(CharSequence category, int boost, boolean isPrefix) {
+        this.category = category;
+        this.boost = boost;
+        this.isPrefix = isPrefix;
+    }
 
     /**
-     * Creates a query context with a provided context and a
-     * boost of 1
+     * Returns the category of the context
      */
-    public CategoryQueryContext(CharSequence context) {
-        this(context, 1);
+    public CharSequence getCategory() {
+        return category;
     }
 
     /**
-     * Creates a query context with a provided context and boost
+     * Returns if the context should be treated as a prefix
      */
-    public CategoryQueryContext(CharSequence context, int boost) {
-        this(context, boost, false);
+    public boolean isPrefix() {
+        return isPrefix;
     }
 
     /**
-     * Creates a query context with a provided context and boost
-     * Allows specifying whether the context should be treated as
-     * a prefix or not
+     * Returns the query-time boost of the context
      */
-    public CategoryQueryContext(CharSequence context, int boost, boolean isPrefix) {
-        this.context = context;
-        this.boost = boost;
-        this.isPrefix = isPrefix;
+    public int getBoost() {
+        return boost;
     }
 
-    private CategoryQueryContext() {
+    public static Builder builder() {
+        return new Builder();
     }
 
-    void setContext(CharSequence context) {
-        this.context = context;
-    }
+    public static class Builder {
+        private CharSequence category;
+        private boolean isPrefix = false;
+        private int boost = 1;
 
-    void setIsPrefix(boolean isPrefix) {
-        this.isPrefix = isPrefix;
-    }
+        public Builder() {
+        }
 
-    void setBoost(int boost) {
-        this.boost = boost;
+        /**
+         * Sets the category of the context.
+         * This is a required field
+         */
+        public Builder setCategory(CharSequence context) {
+            this.category = context;
+            return this;
+        }
+
+        /**
+         * Sets if the context should be treated as a prefix or not.
+         * Defaults to false
+         */
+        public Builder setPrefix(boolean prefix) {
+            this.isPrefix = prefix;
+            return this;
+        }
+
+        /**
+         * Sets the query-time boost of the context.
+         * Defaults to 1.
+         */
+        public Builder setBoost(int boost) {
+            this.boost = boost;
+            return this;
+        }
+
+        public CategoryQueryContext build() {
+            return new CategoryQueryContext(category, boost, isPrefix);
+        }
     }
 
-    private static ObjectParser<CategoryQueryContext, CategoryContextMapping> CATEGORY_PARSER = new ObjectParser<>("category", null);
+    private static ObjectParser<Builder, Void> CATEGORY_PARSER = new ObjectParser<>("category", null);
     static {
-        CATEGORY_PARSER.declareString(CategoryQueryContext::setContext, new ParseField("context"));
-        CATEGORY_PARSER.declareInt(CategoryQueryContext::setBoost, new ParseField("boost"));
-        CATEGORY_PARSER.declareBoolean(CategoryQueryContext::setIsPrefix, new ParseField("prefix"));
+        CATEGORY_PARSER.declareString(Builder::setCategory, new ParseField("context"));
+        CATEGORY_PARSER.declareInt(Builder::setBoost, new ParseField("boost"));
+        CATEGORY_PARSER.declareBoolean(Builder::setPrefix, new ParseField("prefix"));
     }
 
     public static CategoryQueryContext parse(XContentParser parser) throws IOException {
         XContentParser.Token token = parser.currentToken();
-        CategoryQueryContext queryContext = new CategoryQueryContext();
+        Builder builder = builder();
         if (token == XContentParser.Token.START_OBJECT) {
-            CATEGORY_PARSER.parse(parser, queryContext);
+            CATEGORY_PARSER.parse(parser, builder);
         } else if (token == XContentParser.Token.VALUE_STRING) {
-            queryContext.setContext(parser.text());
+            builder.setCategory(parser.text());
         } else {
             throw new ElasticsearchParseException("category context must be an object or string");
         }
-        return queryContext;
+        return builder.build();
     }
 
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();
-        builder.field(CONTEXT_VALUE, context);
+        builder.field(CONTEXT_VALUE, category);
         builder.field(CONTEXT_BOOST, boost);
         builder.field(CONTEXT_PREFIX, isPrefix);
         builder.endObject();
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoContextMapping.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoContextMapping.java
index 149bfa665c..d4ff2f106b 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoContextMapping.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoContextMapping.java
@@ -245,27 +245,26 @@ public class GeoContextMapping extends ContextMapping {
         List<QueryContext> queryContextList = new ArrayList<>();
         for (GeoQueryContext queryContext : queryContexts) {
             int minPrecision = this.precision;
-            if (queryContext.precision != -1) {
-                minPrecision = Math.min(minPrecision, queryContext.precision);
+            if (queryContext.getPrecision() != -1) {
+                minPrecision = Math.min(minPrecision, queryContext.getPrecision());
             }
-            GeoPoint point = queryContext.geoPoint;
+            GeoPoint point = queryContext.getGeoPoint();
             final Collection<String> locations = new HashSet<>();
             String geoHash = GeoHashUtils.stringEncode(point.getLon(), point.getLat(), minPrecision);
             locations.add(geoHash);
-            if (queryContext.neighbours.isEmpty() && geoHash.length() == this.precision) {
+            if (queryContext.getNeighbours().isEmpty() && geoHash.length() == this.precision) {
                 GeoHashUtils.addNeighbors(geoHash, locations);
-            } else if (queryContext.neighbours.isEmpty() == false) {
-                for (Integer neighbourPrecision : queryContext.neighbours) {
+            } else if (queryContext.getNeighbours().isEmpty() == false) {
+                for (Integer neighbourPrecision : queryContext.getNeighbours()) {
                     if (neighbourPrecision < geoHash.length()) {
                         String truncatedGeoHash = geoHash.substring(0, neighbourPrecision);
                         locations.add(truncatedGeoHash);
                         GeoHashUtils.addNeighbors(truncatedGeoHash, locations);
                     }
-
                 }
             }
             for (String location : locations) {
-                queryContextList.add(new QueryContext(location, queryContext.boost, location.length() < this.precision));
+                queryContextList.add(new QueryContext(location, queryContext.getBoost(), location.length() < this.precision));
             }
         }
         return queryContextList;
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoQueryContext.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoQueryContext.java
index 291f65091e..75cab1e8e8 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoQueryContext.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoQueryContext.java
@@ -29,7 +29,6 @@ import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
 
 import java.io.IOException;
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
@@ -39,127 +38,140 @@ import static org.elasticsearch.search.suggest.completion.context.GeoContextMapp
  * Defines the query context for {@link GeoContextMapping}
  */
 public final class GeoQueryContext implements ToXContent {
-    public GeoPoint geoPoint;
-    public int boost = 1;
-    public int precision = -1;
-    public List<Integer> neighbours = new ArrayList<>(0);
+    private final GeoPoint geoPoint;
+    private final int boost;
+    private final int precision;
+    private final List<Integer> neighbours;
 
-    /**
-     * Creates a query context for a given geo point with a boost of 1
-     * and a precision of {@value GeoContextMapping#DEFAULT_PRECISION}
-     */
-    public GeoQueryContext(GeoPoint geoPoint) {
-        this(geoPoint.geohash());
+    private GeoQueryContext(GeoPoint geoPoint, int boost, int precision, List<Integer> neighbours) {
+        this.geoPoint = geoPoint;
+        this.boost = boost;
+        this.precision = precision;
+        this.neighbours = neighbours;
     }
 
     /**
-     * Creates a query context for a given geo point with a
-     * provided boost
+     * Returns the geo point of the context
      */
-    public GeoQueryContext(GeoPoint geoPoint, int boost) {
-        this(geoPoint.geohash(), boost);
+    public GeoPoint getGeoPoint() {
+        return geoPoint;
     }
 
     /**
-     * Creates a query context with a given geo hash with a boost of 1
-     * and a precision of {@value GeoContextMapping#DEFAULT_PRECISION}
+     * Returns the query-time boost of the context
      */
-    public GeoQueryContext(CharSequence geoHash) {
-        this(geoHash, 1);
+    public int getBoost() {
+        return boost;
     }
 
     /**
-     * Creates a query context for a given geo hash with a
-     * provided boost
+     * Returns the precision (length) for the geohash
      */
-    public GeoQueryContext(CharSequence geoHash, int boost) {
-        this(geoHash, boost, -1);
+    public int getPrecision() {
+        return precision;
     }
 
     /**
-     * Creates a query context for a geo point with
-     * a provided boost and enables generating neighbours
-     * at specified precisions
+     * Returns the precision levels at which geohash cells neighbours are considered
      */
-    public GeoQueryContext(CharSequence geoHash, int boost, int precision, Integer... neighbours) {
-        this(GeoPoint.fromGeohash(geoHash.toString()), boost, precision, neighbours);
+    public List<Integer> getNeighbours() {
+        return neighbours;
     }
 
-    /**
-     * Creates a query context for a geo hash with
-     * a provided boost and enables generating neighbours
-     * at specified precisions
-     */
-    public GeoQueryContext(GeoPoint geoPoint, int boost, int precision, Integer... neighbours) {
-        this.geoPoint = geoPoint;
-        this.boost = boost;
-        this.precision = precision;
-        Collections.addAll(this.neighbours, neighbours);
+    public static Builder builder() {
+        return new Builder();
     }
 
-    private GeoQueryContext() {
-    }
+    public static class Builder {
+        private GeoPoint geoPoint;
+        private int boost = 1;
+        private int precision = -1;
+        private List<Integer> neighbours = Collections.emptyList();
 
-    void setBoost(int boost) {
-        this.boost = boost;
-    }
+        public Builder() {
+        }
 
-    void setPrecision(int precision) {
-        this.precision = precision;
-    }
+        /**
+         * Sets the query-time boost for the context
+         * Defaults to 1
+         */
+        public Builder setBoost(int boost) {
+            this.boost = boost;
+            return this;
+        }
 
-    void setNeighbours(List<Integer> neighbours) {
-        this.neighbours = neighbours;
-    }
+        /**
+         * Sets the precision level for computing the geohash from the context geo point.
+         * Defaults to using index-time precision level
+         */
+        public Builder setPrecision(int precision) {
+            this.precision = precision;
+            return this;
+        }
 
-    void setGeoPoint(GeoPoint geoPoint) {
-        this.geoPoint = geoPoint;
-    }
+        /**
+         * Sets the precision levels at which geohash cells neighbours are considered.
+         * Defaults to only considering neighbours at the index-time precision level
+         */
+        public Builder setNeighbours(List<Integer> neighbours) {
+            this.neighbours = neighbours;
+            return this;
+        }
 
-    private double lat = Double.NaN;
-    void setLat(double lat) {
-        this.lat = lat;
-    }
+        /**
+         * Sets the geo point of the context.
+         * This is a required field
+         */
+        public Builder setGeoPoint(GeoPoint geoPoint) {
+            this.geoPoint = geoPoint;
+            return this;
+        }
 
-    private double lon = Double.NaN;
-    void setLon(double lon) {
-        this.lon = lon;
-    }
+        private double lat = Double.NaN;
+        void setLat(double lat) {
+            this.lat = lat;
+        }
+
+        private double lon = Double.NaN;
+        void setLon(double lon) {
+            this.lon = lon;
+        }
 
-    void finish() {
-        if (geoPoint == null) {
-            if (Double.isNaN(lat) == false && Double.isNaN(lon) == false) {
-                geoPoint = new GeoPoint(lat, lon);
-            } else {
-                throw new ElasticsearchParseException("no geohash or geo point provided");
+        public GeoQueryContext build() {
+            if (geoPoint == null) {
+                if (Double.isNaN(lat) == false && Double.isNaN(lon) == false) {
+                    geoPoint = new GeoPoint(lat, lon);
+                } else {
+                    throw new IllegalArgumentException("no geohash or geo point provided");
+                }
             }
+            return new GeoQueryContext(geoPoint, boost, precision, neighbours);
         }
     }
 
-    private static ObjectParser<GeoQueryContext, GeoContextMapping> GEO_CONTEXT_PARSER = new ObjectParser<>("geo", null);
+    private static ObjectParser<GeoQueryContext.Builder, Void> GEO_CONTEXT_PARSER = new ObjectParser<>("geo", null);
     static {
         GEO_CONTEXT_PARSER.declareField((parser, geoQueryContext, geoContextMapping) -> geoQueryContext.setGeoPoint(GeoUtils.parseGeoPoint(parser)), new ParseField("context"), ObjectParser.ValueType.OBJECT);
-        GEO_CONTEXT_PARSER.declareInt(GeoQueryContext::setBoost, new ParseField("boost"));
+        GEO_CONTEXT_PARSER.declareInt(GeoQueryContext.Builder::setBoost, new ParseField("boost"));
         // TODO : add string support for precision for GeoUtils.geoHashLevelsForPrecision()
-        GEO_CONTEXT_PARSER.declareInt(GeoQueryContext::setPrecision, new ParseField("precision"));
+        GEO_CONTEXT_PARSER.declareInt(GeoQueryContext.Builder::setPrecision, new ParseField("precision"));
         // TODO : add string array support for precision for GeoUtils.geoHashLevelsForPrecision()
-        GEO_CONTEXT_PARSER.declareIntArray(GeoQueryContext::setNeighbours, new ParseField("neighbours"));
-        GEO_CONTEXT_PARSER.declareDouble(GeoQueryContext::setLat, new ParseField("lat"));
-        GEO_CONTEXT_PARSER.declareDouble(GeoQueryContext::setLon, new ParseField("lon"));
+        GEO_CONTEXT_PARSER.declareIntArray(GeoQueryContext.Builder::setNeighbours, new ParseField("neighbours"));
+        GEO_CONTEXT_PARSER.declareDouble(GeoQueryContext.Builder::setLat, new ParseField("lat"));
+        GEO_CONTEXT_PARSER.declareDouble(GeoQueryContext.Builder::setLon, new ParseField("lon"));
     }
 
     public static GeoQueryContext parse(XContentParser parser) throws IOException {
         XContentParser.Token token = parser.currentToken();
-        GeoQueryContext queryContext = new GeoQueryContext();
+        GeoQueryContext.Builder builder = new Builder();
         if (token == XContentParser.Token.START_OBJECT) {
-            GEO_CONTEXT_PARSER.parse(parser, queryContext);
+            GEO_CONTEXT_PARSER.parse(parser, builder);
         } else if (token == XContentParser.Token.VALUE_STRING) {
-            queryContext.setGeoPoint(GeoPoint.fromGeohash(parser.text()));
+            builder.setGeoPoint(GeoPoint.fromGeohash(parser.text()));
         } else {
             throw new ElasticsearchParseException("geo context must be an object or string");
         }
-        queryContext.finish();
-        return queryContext;
+        return builder.build();
     }
 
     @Override
diff --git a/core/src/test/java/org/elasticsearch/search/suggest/ContextCompletionSuggestSearchIT.java b/core/src/test/java/org/elasticsearch/search/suggest/ContextCompletionSuggestSearchIT.java
index 91ef41e924..af336b0cc0 100644
--- a/core/src/test/java/org/elasticsearch/search/suggest/ContextCompletionSuggestSearchIT.java
+++ b/core/src/test/java/org/elasticsearch/search/suggest/ContextCompletionSuggestSearchIT.java
@@ -168,8 +168,7 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
         indexRandom(true, indexRequestBuilders);
         ensureYellow(INDEX);
         CompletionSuggestionBuilder prefix = SuggestBuilders.completionSuggestion("foo").field(FIELD).prefix("sugg")
-                .categoryContexts("cat",
-                        new CategoryQueryContext("cat0"));
+                .categoryContexts("cat", CategoryQueryContext.builder().setCategory("cat0").build());
 
         assertSuggestions("foo", prefix, "suggestion8", "suggestion6", "suggestion4", "suggestion2", "suggestion0");
     }
@@ -197,8 +196,9 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
         ensureYellow(INDEX);
         CompletionSuggestionBuilder prefix = SuggestBuilders.completionSuggestion("foo").field(FIELD).prefix("sugg")
                 .categoryContexts("cat",
-                        new CategoryQueryContext("cat0", 3),
-                        new CategoryQueryContext("cat1"));
+                        CategoryQueryContext.builder().setCategory("cat0").setBoost(3).build(),
+                        CategoryQueryContext.builder().setCategory("cat1").build()
+                );
         assertSuggestions("foo", prefix, "suggestion8", "suggestion6", "suggestion4", "suggestion9", "suggestion2");
     }
 
@@ -255,25 +255,25 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
 
         // filter only on context cat
         CompletionSuggestionBuilder catFilterSuggest = SuggestBuilders.completionSuggestion("foo").field(FIELD).prefix("sugg");
-        catFilterSuggest.categoryContexts("cat", new CategoryQueryContext("cat0"));
+        catFilterSuggest.categoryContexts("cat", CategoryQueryContext.builder().setCategory("cat0").build());
         assertSuggestions("foo", catFilterSuggest, "suggestion8", "suggestion6", "suggestion4", "suggestion2", "suggestion0");
 
         // filter only on context type
         CompletionSuggestionBuilder typeFilterSuggest = SuggestBuilders.completionSuggestion("foo").field(FIELD).prefix("sugg");
-        typeFilterSuggest.categoryContexts("type", new CategoryQueryContext("type2"), new CategoryQueryContext("type1"));
+        typeFilterSuggest.categoryContexts("type", CategoryQueryContext.builder().setCategory("type2").build(),
+                CategoryQueryContext.builder().setCategory("type1").build());
         assertSuggestions("foo", typeFilterSuggest, "suggestion9", "suggestion6", "suggestion5", "suggestion2", "suggestion1");
 
-        // filter on both contexts
         CompletionSuggestionBuilder multiContextFilterSuggest = SuggestBuilders.completionSuggestion("foo").field(FIELD).prefix("sugg");
         // query context order should never matter
         if (randomBoolean()) {
-            multiContextFilterSuggest.categoryContexts("type", new CategoryQueryContext("type2"), new CategoryQueryContext("type1"));
-            multiContextFilterSuggest.categoryContexts("cat", new CategoryQueryContext("cat0"));
+            multiContextFilterSuggest.categoryContexts("type", CategoryQueryContext.builder().setCategory("type2").build());
+            multiContextFilterSuggest.categoryContexts("cat", CategoryQueryContext.builder().setCategory("cat2").build());
         } else {
-            multiContextFilterSuggest.categoryContexts("cat", new CategoryQueryContext("cat0"));
-            multiContextFilterSuggest.categoryContexts("type", new CategoryQueryContext("type2"), new CategoryQueryContext("type1"));
+            multiContextFilterSuggest.categoryContexts("cat", CategoryQueryContext.builder().setCategory("cat2").build());
+            multiContextFilterSuggest.categoryContexts("type", CategoryQueryContext.builder().setCategory("type2").build());
         }
-        assertSuggestions("foo", multiContextFilterSuggest, "suggestion9", "suggestion8",  "suggestion6", "suggestion5", "suggestion4");
+        assertSuggestions("foo", multiContextFilterSuggest, "suggestion6", "suggestion2");
     }
 
     @AwaitsFix(bugUrl = "multiple context boosting is broken, as a suggestion, contexts pair is treated as (num(context) entries)")
@@ -304,15 +304,15 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
         // boost only on context cat
         CompletionSuggestionBuilder catBoostSuggest = SuggestBuilders.completionSuggestion("foo").field(FIELD).prefix("sugg");
         catBoostSuggest.categoryContexts("cat",
-                new CategoryQueryContext("cat0", 3),
-                new CategoryQueryContext("cat1"));
+                CategoryQueryContext.builder().setCategory("cat0").setBoost(3).build(),
+                CategoryQueryContext.builder().setCategory("cat1").build());
         assertSuggestions("foo", catBoostSuggest, "suggestion8", "suggestion6", "suggestion4", "suggestion9", "suggestion2");
 
         // boost only on context type
         CompletionSuggestionBuilder typeBoostSuggest = SuggestBuilders.completionSuggestion("foo").field(FIELD).prefix("sugg");
         typeBoostSuggest.categoryContexts("type",
-                new CategoryQueryContext("type2", 2),
-                new CategoryQueryContext("type1", 4));
+                CategoryQueryContext.builder().setCategory("type2").setBoost(2).build(),
+                CategoryQueryContext.builder().setCategory("type1").setBoost(4).build());
         assertSuggestions("foo", typeBoostSuggest, "suggestion9", "suggestion5", "suggestion6", "suggestion1", "suggestion2");
 
         // boost on both contexts
@@ -320,18 +320,18 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
         // query context order should never matter
         if (randomBoolean()) {
             multiContextBoostSuggest.categoryContexts("type",
-                    new CategoryQueryContext("type2", 2),
-                    new CategoryQueryContext("type1", 4));
+                    CategoryQueryContext.builder().setCategory("type2").setBoost(2).build(),
+                    CategoryQueryContext.builder().setCategory("type1").setBoost(4).build());
             multiContextBoostSuggest.categoryContexts("cat",
-                    new CategoryQueryContext("cat0", 3),
-                    new CategoryQueryContext("cat1"));
+                    CategoryQueryContext.builder().setCategory("cat0").setBoost(3).build(),
+                    CategoryQueryContext.builder().setCategory("cat1").build());
         } else {
             multiContextBoostSuggest.categoryContexts("cat",
-                    new CategoryQueryContext("cat0", 3),
-                    new CategoryQueryContext("cat1"));
+                    CategoryQueryContext.builder().setCategory("cat0").setBoost(3).build(),
+                    CategoryQueryContext.builder().setCategory("cat1").build());
             multiContextBoostSuggest.categoryContexts("type",
-                    new CategoryQueryContext("type2", 2),
-                    new CategoryQueryContext("type1", 4));
+                    CategoryQueryContext.builder().setCategory("type2").setBoost(2).build(),
+                    CategoryQueryContext.builder().setCategory("type1").setBoost(4).build());
         }
         assertSuggestions("foo", multiContextBoostSuggest, "suggestion9", "suggestion6", "suggestion5", "suggestion2", "suggestion1");
     }
@@ -431,7 +431,7 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
         createIndexAndMapping(mapping);
         int numDocs = 10;
         List<IndexRequestBuilder> indexRequestBuilders = new ArrayList<>();
-        String[] geoHashes = new String[] {"ezs42e44yx96", "u4pruydqqvj8"};
+        GeoPoint[] geoPoints = new GeoPoint[] {new GeoPoint("ezs42e44yx96"), new GeoPoint("u4pruydqqvj8")};
         for (int i = 0; i < numDocs; i++) {
             XContentBuilder source = jsonBuilder()
                     .startObject()
@@ -439,7 +439,7 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
                     .field("input", "suggestion" + i)
                     .field("weight", i + 1)
                     .startObject("contexts")
-                    .field("geo", (i % 2 == 0) ? geoHashes[0] : geoHashes[1])
+                    .field("geo", (i % 2 == 0) ? geoPoints[0].getGeohash() : geoPoints[1].getGeohash())
                     .endObject()
                     .endObject().endObject();
             indexRequestBuilders.add(client().prepareIndex(INDEX, TYPE, "" + i)
@@ -451,7 +451,7 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
         assertSuggestions("foo", prefix, "suggestion9", "suggestion8", "suggestion7", "suggestion6", "suggestion5");
 
         CompletionSuggestionBuilder geoFilteringPrefix = SuggestBuilders.completionSuggestion("foo").field(FIELD).prefix("sugg")
-                .geoContexts("geo", new GeoQueryContext(geoHashes[0]));
+                .geoContexts("geo", GeoQueryContext.builder().setGeoPoint(new GeoPoint(geoPoints[0])).build());
 
         assertSuggestions("foo", geoFilteringPrefix, "suggestion8", "suggestion6", "suggestion4", "suggestion2", "suggestion0");
     }
@@ -463,7 +463,7 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
         createIndexAndMapping(mapping);
         int numDocs = 10;
         List<IndexRequestBuilder> indexRequestBuilders = new ArrayList<>();
-        String[] geoHashes = new String[] {"ezs42e44yx96", "u4pruydqqvj8"};
+        GeoPoint[] geoPoints = new GeoPoint[] {new GeoPoint("ezs42e44yx96"), new GeoPoint("u4pruydqqvj8")};
         for (int i = 0; i < numDocs; i++) {
             XContentBuilder source = jsonBuilder()
                     .startObject()
@@ -471,7 +471,7 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
                     .field("input", "suggestion" + i)
                     .field("weight", i + 1)
                     .startObject("contexts")
-                    .field("geo", (i % 2 == 0) ? geoHashes[0] : geoHashes[1])
+                    .field("geo", (i % 2 == 0) ? geoPoints[0].getGeohash() : geoPoints[1].getGeohash())
                     .endObject()
                     .endObject().endObject();
             indexRequestBuilders.add(client().prepareIndex(INDEX, TYPE, "" + i)
@@ -482,8 +482,10 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
         CompletionSuggestionBuilder prefix = SuggestBuilders.completionSuggestion("foo").field(FIELD).prefix("sugg");
         assertSuggestions("foo", prefix, "suggestion9", "suggestion8", "suggestion7", "suggestion6", "suggestion5");
 
+        GeoQueryContext context1 = GeoQueryContext.builder().setGeoPoint(geoPoints[0]).setBoost(2).build();
+        GeoQueryContext context2 = GeoQueryContext.builder().setGeoPoint(geoPoints[1]).build();
         CompletionSuggestionBuilder geoBoostingPrefix = SuggestBuilders.completionSuggestion("foo").field(FIELD).prefix("sugg")
-                .geoContexts("geo", new GeoQueryContext(geoHashes[0], 2), new GeoQueryContext(geoHashes[1]));
+                .geoContexts("geo", context1, context2);
 
         assertSuggestions("foo", geoBoostingPrefix, "suggestion8", "suggestion6", "suggestion4", "suggestion9", "suggestion7");
     }
@@ -514,7 +516,7 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
         indexRandom(true, indexRequestBuilders);
         ensureYellow(INDEX);
         CompletionSuggestionBuilder prefix = SuggestBuilders.completionSuggestion("foo").field(FIELD).prefix("sugg")
-                .geoContexts("geo", new GeoQueryContext(new GeoPoint(52.2263, 4.543)));
+                .geoContexts("geo", GeoQueryContext.builder().setGeoPoint(new GeoPoint(52.2263, 4.543)).build());
         assertSuggestions("foo", prefix, "suggestion9", "suggestion8", "suggestion7", "suggestion6", "suggestion5");
     }
 
@@ -555,7 +557,7 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
         assertSuggestions("foo", prefix, "suggestion9", "suggestion8", "suggestion7", "suggestion6", "suggestion5");
 
         CompletionSuggestionBuilder geoNeighbourPrefix = SuggestBuilders.completionSuggestion("foo").field(FIELD).prefix("sugg")
-                .geoContexts("geo", new GeoQueryContext(geohash));
+                .geoContexts("geo", GeoQueryContext.builder().setGeoPoint(GeoPoint.fromGeohash(geohash)).build());
 
         assertSuggestions("foo", geoNeighbourPrefix, "suggestion9", "suggestion8", "suggestion7", "suggestion6", "suggestion5");
     }