From ee7fd9879bff8421318330635f9d01cb33d61f62 Mon Sep 17 00:00:00 2001
From: Justin Robb <reptarsrage@gmail.com>
Date: Wed, 8 May 2024 16:38:42 -0700
Subject: [PATCH] support for grouping search (#83)

---
 Milvus.Client.Tests/SearchQueryTests.cs  | 24 ++++++++++++++++++++++++
 Milvus.Client/Constants.cs               |  5 +++++
 Milvus.Client/MilvusCollection.Entity.cs |  9 +++++++++
 Milvus.Client/SearchParameters.cs        |  8 ++++++++
 4 files changed, 46 insertions(+)

diff --git a/Milvus.Client.Tests/SearchQueryTests.cs b/Milvus.Client.Tests/SearchQueryTests.cs
index 7ad50d6..985f25a 100644
--- a/Milvus.Client.Tests/SearchQueryTests.cs
+++ b/Milvus.Client.Tests/SearchQueryTests.cs
@@ -551,6 +551,30 @@ public async Task Search_with_unsupported_vector_type_throws()
         Assert.Equal("vectors", exception.ParamName);
     }
 
+    [Fact]
+    public async Task Search_with_group_by_field()
+    {
+        var parameters = new SearchParameters();
+        parameters.GroupByField = "id";
+        var results = await Collection.SearchAsync(
+            "float_vector",
+            new ReadOnlyMemory<float>[] { new[] { 0.1f, 0.2f } },
+            SimilarityMetricType.L2,
+            parameters: parameters,
+            limit: 2);
+
+        Assert.Equal(CollectionName, results.CollectionName);
+        Assert.Empty(results.FieldsData);
+        Assert.Collection(results.Ids.LongIds!,
+            id => Assert.Equal(1, id),
+            id => Assert.Equal(2, id));
+        Assert.Null(results.Ids.StringIds); // TODO: Need to test string IDs
+        Assert.Equal(1, results.NumQueries);
+        Assert.Equal(2, results.Scores.Count);
+        Assert.Equal(2, results.Limit);
+        Assert.Collection(results.Limits, l => Assert.Equal(2, l));
+    }
+
     public class QueryCollectionFixture : IAsyncLifetime
     {
         public QueryCollectionFixture(MilvusFixture milvusFixture)
diff --git a/Milvus.Client/Constants.cs b/Milvus.Client/Constants.cs
index c59455f..7babf2b 100644
--- a/Milvus.Client/Constants.cs
+++ b/Milvus.Client/Constants.cs
@@ -100,6 +100,11 @@ internal static class Constants
     /// </summary>
     internal const string IgnoreGrowing = "ignore_growing";
 
+    /// <summary>
+    /// Key name.
+    /// </summary>
+    internal const string GroupByField = "group_by_field";
+
     /// <summary>
     /// Default database name.
     /// </summary>
diff --git a/Milvus.Client/MilvusCollection.Entity.cs b/Milvus.Client/MilvusCollection.Entity.cs
index 5e897a0..e4e0c88 100644
--- a/Milvus.Client/MilvusCollection.Entity.cs
+++ b/Milvus.Client/MilvusCollection.Entity.cs
@@ -242,6 +242,15 @@ public async Task<SearchResults> SearchAsync<T>(
                     Key = Constants.IgnoreGrowing, Value = parameters.IgnoreGrowing.Value.ToString()
                 });
             }
+
+            if (parameters.GroupByField is not null)
+            {
+                request.SearchParams.Add(new Grpc.KeyValuePair
+                {
+                    Key = Constants.GroupByField,
+                    Value = parameters.GroupByField
+                });
+            }
         }
 
         // Note that we send both the consistency level and the guarantee timestamp, although the latter is derived
diff --git a/Milvus.Client/SearchParameters.cs b/Milvus.Client/SearchParameters.cs
index c2b054b..878c25f 100644
--- a/Milvus.Client/SearchParameters.cs
+++ b/Milvus.Client/SearchParameters.cs
@@ -72,4 +72,12 @@ public class SearchParameters
     /// searches involve growing segments.
     /// </summary>
     public bool? IgnoreGrowing { get; private set; }
+
+    /// <summary>
+    /// Group search results by the specified field.
+    /// </summary>
+    /// <remarks>
+    /// See <see href="https://milvus.io/docs/single-vector-search.md#Grouping-search" /> for more information.
+    /// </remarks>
+    public string? GroupByField { get; set; }
 }