Added diversified sampler aggregation (#5031)

Co-authored-by: Martijn Laarman <Mpdreamz@gmail.com>

Author: mandiscombination

src/Nest/Aggregations/AggregateDictionary.cs

@@ -133,6 +133,8 @@ public FiltersAggregate Filters(string key)
 
 		public SingleBucketAggregate Sampler(string key) => TryGet<SingleBucketAggregate>(key);
 
+		public SingleBucketAggregate DiversifiedSampler(string key) => TryGet<SingleBucketAggregate>(key);
+
 		public GeoCentroidAggregate GeoCentroid(string key) => TryGet<GeoCentroidAggregate>(key);
 
 		public SignificantTermsAggregate<TKey> SignificantTerms<TKey>(string key)

src/Nest/Aggregations/AggregationContainer.cs

@@ -135,6 +135,9 @@ public interface IAggregationContainer
 		[DataMember(Name = "derivative")]
 		IDerivativeAggregation Derivative { get; set; }
 
+		[DataMember(Name = "diversified_sampler")]
+		IDiversifiedSamplerAggregation DiversifiedSampler { get; set; }
+
 		[DataMember(Name = "extended_stats")]
 		IExtendedStatsAggregation ExtendedStats { get; set; }
 
@@ -324,6 +327,8 @@ public class AggregationContainer : IAggregationContainer
 
 		public IDerivativeAggregation Derivative { get; set; }
 
+		public IDiversifiedSamplerAggregation DiversifiedSampler { get; set; }
+
 		public IExtendedStatsAggregation ExtendedStats { get; set; }
 
 		public IExtendedStatsBucketAggregation ExtendedStatsBucket { get; set; }
@@ -485,6 +490,8 @@ public class AggregationContainerDescriptor<T> : DescriptorBase<AggregationConta
 
 		IDerivativeAggregation IAggregationContainer.Derivative { get; set; }
 
+		IDiversifiedSamplerAggregation IAggregationContainer.DiversifiedSampler { get; set; }
+
 		IExtendedStatsAggregation IAggregationContainer.ExtendedStats { get; set; }
 
 		IExtendedStatsBucketAggregation IAggregationContainer.ExtendedStatsBucket { get; set; }
@@ -860,6 +867,11 @@ Func<SamplerAggregationDescriptor<T>, ISamplerAggregation> selector
 		) =>
 			_SetInnerAggregation(name, selector, (a, d) => a.Sampler = d);
 
+		public AggregationContainerDescriptor<T> DiversifiedSampler(string name,
+			Func<DiversifiedSamplerAggregationDescriptor<T>, IDiversifiedSamplerAggregation> selector
+		) =>
+			_SetInnerAggregation(name, selector, (a, d) => a.DiversifiedSampler = d);
+
 		public AggregationContainerDescriptor<T> GeoCentroid(string name,
 			Func<GeoCentroidAggregationDescriptor<T>, IGeoCentroidAggregation> selector
 		) =>

src/Nest/Aggregations/Bucket/DiversifiedSampler/DiversifiedSamplerAggregation.cs

@@ -0,0 +1,69 @@
+using System;
+using System.Linq.Expressions;
+using System.Runtime.Serialization;
+using Elasticsearch.Net.Utf8Json;
+
+namespace Nest
+{
+	[InterfaceDataContract]
+	[ReadAs(typeof(DiversifiedSamplerAggregation))]
+	public interface IDiversifiedSamplerAggregation : IBucketAggregation
+	{
+		[DataMember(Name ="execution_hint")]
+		DiversifiedSamplerAggregationExecutionHint? ExecutionHint { get; set; }
+
+		[DataMember(Name = "field")]
+		Field Field { get; set; }
+
+		[DataMember(Name ="max_docs_per_value")]
+		int? MaxDocsPerValue { get; set; }
+
+		[DataMember(Name ="script")]
+		IScript Script { get; set; }
+
+		[DataMember(Name ="shard_size")]
+		int? ShardSize { get; set; }
+	}
+
+	public class DiversifiedSamplerAggregation : BucketAggregationBase, IDiversifiedSamplerAggregation
+	{
+		internal DiversifiedSamplerAggregation() { }
+
+		public DiversifiedSamplerAggregation(string name) : base(name) { }
+
+		public DiversifiedSamplerAggregationExecutionHint? ExecutionHint { get; set; }
+		public Field Field { get; set; }
+		public int? MaxDocsPerValue { get; set; }
+		public IScript Script { get; set; }
+		public int? ShardSize { get; set; }
+
+		internal override void WrapInContainer(AggregationContainer c) => c.DiversifiedSampler = this;
+	}
+
+	public class DiversifiedSamplerAggregationDescriptor<T>
+		: BucketAggregationDescriptorBase<DiversifiedSamplerAggregationDescriptor<T>, IDiversifiedSamplerAggregation, T>, IDiversifiedSamplerAggregation
+		where T : class
+	{
+		DiversifiedSamplerAggregationExecutionHint? IDiversifiedSamplerAggregation.ExecutionHint { get; set; }
+		Field IDiversifiedSamplerAggregation.Field { get; set; }
+		int? IDiversifiedSamplerAggregation.MaxDocsPerValue { get; set; }
+		IScript IDiversifiedSamplerAggregation.Script { get; set; }
+		int? IDiversifiedSamplerAggregation.ShardSize { get; set; }
+
+		public DiversifiedSamplerAggregationDescriptor<T> ExecutionHint(DiversifiedSamplerAggregationExecutionHint? executionHint) =>
+			Assign(executionHint, (a, v) => a.ExecutionHint = v);
+
+		public DiversifiedSamplerAggregationDescriptor<T> Field(Field field) => Assign(field, (a, v) => a.Field = v);
+
+		public DiversifiedSamplerAggregationDescriptor<T> Field<TValue>(Expression<Func<T, TValue>> field) => Assign(field, (a, v) => a.Field = v);
+
+		public DiversifiedSamplerAggregationDescriptor<T> MaxDocsPerValue(int? maxDocs) => Assign(maxDocs, (a, v) => a.MaxDocsPerValue = v);
+
+		public DiversifiedSamplerAggregationDescriptor<T> Script(string script) => Assign((InlineScript)script, (a, v) => a.Script = v);
+
+		public DiversifiedSamplerAggregationDescriptor<T> Script(Func<ScriptDescriptor, IScript> scriptSelector) =>
+			Assign(scriptSelector, (a, v) => a.Script = v?.Invoke(new ScriptDescriptor()));
+
+		public DiversifiedSamplerAggregationDescriptor<T> ShardSize(int? shardSize) => Assign(shardSize, (a, v) => a.ShardSize = v);
+	}
+}

src/Nest/Aggregations/Bucket/DiversifiedSampler/DiversifiedSamplerAggregationExecutionHint.cs

@@ -0,0 +1,18 @@
+using System.Runtime.Serialization;
+using Elasticsearch.Net;
+
+namespace Nest
+{
+	[StringEnum]
+	public enum DiversifiedSamplerAggregationExecutionHint
+	{
+		[EnumMember(Value = "map")]
+		Map,
+
+		[EnumMember(Value = "global_ordinals")]
+		GlobalOrdinals,
+
+		[EnumMember(Value = "bytes_hash")]
+		BytesHash
+	}
+}

src/Nest/Aggregations/Visitor/AggregationVisitor.cs

@@ -130,6 +130,8 @@ public interface IAggregationVisitor
 
 		void Visit(ISamplerAggregation aggregation);
 
+		void Visit(IDiversifiedSamplerAggregation aggregation);
+
 		void Visit(IGeoCentroidAggregation aggregation);
 
 		void Visit(ICompositeAggregation aggregation);
@@ -213,6 +215,8 @@ public virtual void Visit(IBucketScriptAggregation aggregation) { }
 
 		public virtual void Visit(ISamplerAggregation aggregation) { }
 
+		public virtual void Visit(IDiversifiedSamplerAggregation aggregation) { }
+
 		public virtual void Visit(IBucketSelectorAggregation aggregation) { }
 
 		public virtual void Visit(IBucketSortAggregation aggregation) { }

tests/Examples/Aggregations/Bucket/DiversifiedSamplerAggregationPage.cs

@@ -10,15 +10,35 @@ namespace Examples.Aggregations.Bucket
 {
 	public class DiversifiedSamplerAggregationPage : ExampleBase
 	{
-		[U(Skip = "Example not implemented")]
+		[U]
 		[Description("aggregations/bucket/diversified-sampler-aggregation.asciidoc:30")]
 		public void Line30()
 		{
 			// tag::3344c3478f1e8bbbef683757638a34f4[]
-			var response0 = new SearchResponse<object>();
+			var searchResponse = client.Search<object>(s => s
+				.Index("stackoverflow")
+				.Query(q => q
+					.QueryString(qs => qs
+						.Query("tags:elasticsearch")
+					)
+				)
+				.Aggregations(a => a
+					.DiversifiedSampler("my_unbiased_sample", s => s
+						.ShardSize(200)
+						.Field("author")
+						.Aggregations(agg => agg
+							.SignificantTerms("keywords", k => k
+								.Field("tags")
+								.Exclude(new string[] { "elasticsearch" })
+							)
+						)
+					)
+				)
+				.Size(0)
+			);
 			// end::3344c3478f1e8bbbef683757638a34f4[]
 
-			response0.MatchesExample(@"POST /stackoverflow/_search?size=0
+			searchResponse.MatchesExample(@"POST /stackoverflow/_search?size=0
 			{
 			    ""query"": {
 			        ""query_string"": {
@@ -41,18 +61,46 @@ public void Line30()
 			            }
 			        }
 			    }
-			}");
+			}", (e, b) =>
+			{
+				e.Uri.Query = e.Uri.Query.Replace("size=0", string.Empty);
+				b["size"] = 0;
+			});
 		}
 
-		[U(Skip = "Example not implemented")]
+		[U]
 		[Description("aggregations/bucket/diversified-sampler-aggregation.asciidoc:95")]
 		public void Line95()
 		{
 			// tag::07afce825c09de17a3d73a02b17a0a97[]
-			var response0 = new SearchResponse<object>();
+			var searchResponse = client.Search<object>(s => s
+				.Index("stackoverflow")
+				.Query(q => q
+					.QueryString(qs => qs
+						.Query("tags:kibana")
+					)
+				)
+				.Aggregations(a => a
+					.DiversifiedSampler("my_unbiased_sample", s => s
+						.ShardSize(200)
+						.MaxDocsPerValue(3)
+						.Script(sc => sc
+							.Source("doc['tags'].hashCode()")
+							.Lang("painless")
+						)
+						.Aggregations(agg => agg
+							.SignificantTerms("keywords", k => k
+								.Field("tags")
+								.Exclude(new string[] { "kibana" })
+							)
+						)
+					)
+				)
+				.Size(0)
+			);
 			// end::07afce825c09de17a3d73a02b17a0a97[]
 
-			response0.MatchesExample(@"POST /stackoverflow/_search?size=0
+			searchResponse.MatchesExample(@"POST /stackoverflow/_search?size=0
 			{
 			    ""query"": {
 			        ""query_string"": {
@@ -79,7 +127,11 @@ public void Line95()
 			            }
 			        }
 			    }
-			}");
+			}", (e, b) =>
+			{
+				e.Uri.Query = e.Uri.Query.Replace("size=0", string.Empty);
+				b["size"] = 0;
+			});
 		}
 	}
 }

tests/Tests/Aggregations/Bucket/DiversifiedSampler/DiversifiedSamplerAggregationUsageTests.cs

@@ -0,0 +1,69 @@
+// Licensed to Elasticsearch B.V under one or more agreements.
+// Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
+// See the LICENSE file in the project root for more information
+
+using System;
+using Elastic.Elasticsearch.Xunit.XunitPlumbing;
+using Nest;
+using Tests.Core.ManagedElasticsearch.Clusters;
+using Tests.Domain;
+using Tests.Framework.EndpointTests.TestState;
+
+namespace Tests.Aggregations.Bucket.DiversifiedSampler
+{
+	[SkipVersion("<7.9.0", "introduced in 7.9.0")]
+	public class DiversifiedSamplerAggregationUsageTests : AggregationUsageTestBase
+	{
+		public DiversifiedSamplerAggregationUsageTests(ReadOnlyCluster i, EndpointUsage usage) : base(i, usage) { }
+
+		protected override object AggregationJson => new
+		{
+			diversified_sample = new
+			{
+				diversified_sampler = new
+				{
+					execution_hint = "global_ordinals",
+					field = "type",
+					max_docs_per_value = 10,
+					shard_size = 200
+				},
+				aggs = new
+				{
+					significant_names = new
+					{
+						significant_terms = new
+						{
+							field = "name"
+						}
+					}
+				}
+			}
+		};
+
+		protected override Func<AggregationContainerDescriptor<Project>, IAggregationContainer> FluentAggs => a => a
+			.DiversifiedSampler("diversified_sample", sm => sm
+				.ExecutionHint(DiversifiedSamplerAggregationExecutionHint.GlobalOrdinals)
+				.Field(doc => doc.Type)
+				.MaxDocsPerValue(10)
+				.ShardSize(200)
+				.Aggregations(aa => aa
+					.SignificantTerms("significant_names", st => st
+						.Field(p => p.Name)
+					)
+				)
+			);
+
+		protected override AggregationDictionary InitializerAggs =>
+			new DiversifiedSamplerAggregation("diversified_sample")
+			{
+				ExecutionHint = DiversifiedSamplerAggregationExecutionHint.GlobalOrdinals,
+				Field = new Field("type"),
+				MaxDocsPerValue = 10,
+				ShardSize = 200,
+				Aggregations = new SignificantTermsAggregation("significant_names")
+				{
+					Field = "name"
+				}
+			};
+	}
+}