From ba24e81ce55a4bf7cdd10d96105bec398c1e7bc0 Mon Sep 17 00:00:00 2001
From: Greg Marzouka <greg.marzouka@gmail.com>
Date: Mon, 6 Oct 2014 15:04:43 -0400
Subject: [PATCH] Add Collate support to phrase suggester

Closes #824
---
 .../Suggest/PhraseSuggestCollateDescriptor.cs | 80 +++++++++++++++++++
 .../DSL/Suggest/PhraseSuggestDescriptor.cs    | 14 ++++
 src/Nest/Nest.csproj                          |  1 +
 .../Search/suggest/PhraseSuggestTests.cs      | 58 ++++++++++++++
 4 files changed, 153 insertions(+)
 create mode 100644 src/Nest/DSL/Suggest/PhraseSuggestCollateDescriptor.cs

diff --git a/src/Nest/DSL/Suggest/PhraseSuggestCollateDescriptor.cs b/src/Nest/DSL/Suggest/PhraseSuggestCollateDescriptor.cs
new file mode 100644
index 00000000000..7284efc421d
--- /dev/null
+++ b/src/Nest/DSL/Suggest/PhraseSuggestCollateDescriptor.cs
@@ -0,0 +1,80 @@
+using Newtonsoft.Json;
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+
+namespace Nest
+{
+	[JsonObject(MemberSerialization = MemberSerialization.OptIn)]
+	public interface IPhraseSuggestCollate
+	{
+		[JsonProperty(PropertyName = "query")]
+		IQueryContainer Query { get; set; }
+
+		[JsonProperty(PropertyName = "filter")]
+		IFilterContainer Filter { get; set; }
+
+		[JsonProperty(PropertyName = "params")]
+		IDictionary<string, object> Params { get; set; }
+
+		[JsonProperty(PropertyName = "preference")]
+		string Preference { get; set; }
+	}
+
+	public class PhraseSuggestCollate : IPhraseSuggestCollate
+	{
+		public IQueryContainer Query { get; set; }
+
+		public IFilterContainer Filter { get; set; }
+
+		public IDictionary<string, object> Params { get; set; }
+
+		public string Preference { get; set; }
+	}
+
+	public class PhraseSuggestCollateDescriptor<T> : IPhraseSuggestCollate
+		where T : class
+	{
+		internal IPhraseSuggestCollate Collate = new PhraseSuggestCollate();
+
+		IQueryContainer IPhraseSuggestCollate.Query { get; set; }
+
+		IFilterContainer IPhraseSuggestCollate.Filter { get; set; }
+
+		IDictionary<string, object> IPhraseSuggestCollate.Params { get; set; }
+
+		string IPhraseSuggestCollate.Preference { get; set; }
+
+		public PhraseSuggestCollateDescriptor<T> Query(Func<QueryDescriptor<T>, QueryContainer> query)
+		{
+			this.Collate.Query = query(new QueryDescriptor<T>());
+			return this;
+		}
+
+		public PhraseSuggestCollateDescriptor<T> Filter(Func<FilterDescriptor<T>, FilterContainer> filter)
+		{
+			this.Collate.Filter = filter(new FilterDescriptor<T>());
+			return this;
+		}
+
+		public PhraseSuggestCollateDescriptor<T> Params(IDictionary<string, object> paramsDictionary)
+		{
+			this.Collate.Params = paramsDictionary;
+			return this;
+		}
+
+		public PhraseSuggestCollateDescriptor<T> Params(Func<FluentDictionary<string, object>, FluentDictionary<string, object>> paramsDictionary)
+		{
+			this.Collate.Params = paramsDictionary(new FluentDictionary<string, object>());
+			return this;
+		}
+
+		public PhraseSuggestCollateDescriptor<T> Preference(string preference)
+		{
+			this.Collate.Preference = preference;
+			return this;
+		}
+	}
+
+}
diff --git a/src/Nest/DSL/Suggest/PhraseSuggestDescriptor.cs b/src/Nest/DSL/Suggest/PhraseSuggestDescriptor.cs
index ce19ac2c231..1622785f66b 100644
--- a/src/Nest/DSL/Suggest/PhraseSuggestDescriptor.cs
+++ b/src/Nest/DSL/Suggest/PhraseSuggestDescriptor.cs
@@ -28,6 +28,9 @@ public interface IPhraseSuggester : ISuggester
 
 		[JsonProperty(PropertyName = "direct_generator")]
 		IEnumerable<IDirectGenerator> DirectGenerator { get; set; }
+
+		[JsonProperty("collate")]
+		IPhraseSuggestCollate Collate { get; set; }
 	}
 
 	public class PhraseSuggester : Suggester, IPhraseSuggester
@@ -38,6 +41,8 @@ public class PhraseSuggester : Suggester, IPhraseSuggester
 		public decimal? MaxErrors { get; set; }
 		public char? Separator { get; set; }
 		public IEnumerable<IDirectGenerator> DirectGenerator { get; set; }
+
+		public IPhraseSuggestCollate Collate { get; set; }
 	}
 
 	public class PhraseSuggestDescriptor<T> : BaseSuggestDescriptor<T>, IPhraseSuggester where T : class
@@ -56,6 +61,8 @@ public class PhraseSuggestDescriptor<T> : BaseSuggestDescriptor<T>, IPhraseSugge
 
 		IEnumerable<IDirectGenerator> IPhraseSuggester.DirectGenerator { get; set; }
 
+		IPhraseSuggestCollate IPhraseSuggester.Collate { get; set; }
+
 		public PhraseSuggestDescriptor<T> Text(string text)
 		{
 			Self.Text = text;
@@ -121,5 +128,12 @@ public PhraseSuggestDescriptor<T> DirectGenerator(params Func<DirectGeneratorDes
 			Self.DirectGenerator = generators.Select(g => g(new DirectGeneratorDescriptor<T>())).ToList();
 			return this;
 		}
+
+		public PhraseSuggestDescriptor<T> Collate(Func<PhraseSuggestCollateDescriptor<T>, PhraseSuggestCollateDescriptor<T>> collateDescriptor)
+		{
+			var selector = collateDescriptor(new PhraseSuggestCollateDescriptor<T>());
+			Self.Collate = selector.Collate;
+			return this;
+		}
 	}
 }
diff --git a/src/Nest/Nest.csproj b/src/Nest/Nest.csproj
index d1d68414a66..01b56207251 100644
--- a/src/Nest/Nest.csproj
+++ b/src/Nest/Nest.csproj
@@ -238,6 +238,7 @@
     <Compile Include="DSL\Paths\RepositorySnapshotOptionalPathDescriptor.cs" />
     <Compile Include="DSL\SnapshotStatusDescriptor.cs" />
     <Compile Include="DSL\RecoveryStatusDescriptor.cs" />
+    <Compile Include="DSL\Suggest\PhraseSuggestCollateDescriptor.cs" />
     <Compile Include="DSL\TemplateExistsDescriptor.cs" />
     <Compile Include="Domain\Responses\PingResponse.cs" />
     <Compile Include="Domain\Responses\NodesShutdownResponse.cs" />
diff --git a/src/Tests/Nest.Tests.Unit/Search/suggest/PhraseSuggestTests.cs b/src/Tests/Nest.Tests.Unit/Search/suggest/PhraseSuggestTests.cs
index 68ef3c08abf..c1f72ff9b8b 100644
--- a/src/Tests/Nest.Tests.Unit/Search/suggest/PhraseSuggestTests.cs
+++ b/src/Tests/Nest.Tests.Unit/Search/suggest/PhraseSuggestTests.cs
@@ -1,6 +1,7 @@
 using Elasticsearch.Net;
 using NUnit.Framework;
 using Nest.Tests.MockData.Domain;
+using System.Collections.Generic;
 
 namespace Nest.Tests.Unit.Search.Suggest
 {
@@ -84,5 +85,62 @@ public void PhraseSuggestOnSearchTest()
 			var json = search.ConnectionStatus.Request.Utf8String();
 			Assert.True(json.JsonEquals(expected), json);
 		}
+
+		[Test]
+		public void PhraseSuggestCollateTest()
+		{
+			var search = this._client.Search<ElasticsearchProject>(s => s
+				.SuggestPhrase("myphrasesuggest", ts => ts
+					.Text("n")
+					.Analyzer("body")
+					.OnField("bigram")
+					.Size(1)
+					.GramSize(2)
+					.MaxErrors(0.5m)
+					.Collate(c => c
+						.Query(q => q
+							.Match(m => m
+								.OnField("{{fieldname}}")
+								.Query("{{suggestion}}")
+							)
+						)
+						.Params(ps => ps
+							.Add("{{fieldname}}", "name")
+						)
+						.Preference("_primary")
+					)
+				)
+			);
+
+			var expected = @"{
+				suggest: {
+					myphrasesuggest: {
+						text: ""n"",
+						phrase: {
+                              ""field"": ""bigram"",
+                              ""analyzer"": ""body"",
+                              ""size"": 1,
+							  ""gram_size"": 2,
+                              ""max_errors"": 0.5,
+							  ""collate"": {
+									""query"": {
+										""match"": {
+											""{{fieldname}}"": {
+												""query"": ""{{suggestion}}""
+											}
+										}
+									},
+									""params"": {
+										""{{fieldname}}"": ""name""
+									},
+									""preference"": ""_primary""
+							  }
+						}
+					}
+				}
+			}";
+			var json = search.ConnectionStatus.Request.Utf8String();
+			Assert.True(json.JsonEquals(expected), json);
+		}
 	} 
 }