enhance: return topks when search in restful v2

Signed-off-by: lixinguo <[email protected]>
milvus-io · Feb 13, 2025 · 55cefd6 · 55cefd6
1 parent 16ea4fc
commit 55cefd6
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 3 deletions.
diff --git a/internal/distributed/proxy/httpserver/constant.go b/internal/distributed/proxy/httpserver/constant.go
@@ -126,6 +126,7 @@ const (
 	HTTPReturnRecalls        = "recalls"
 	HTTPReturnLoadState      = "loadState"
 	HTTPReturnLoadProgress   = "loadProgress"
+	HTTPReturnTopks          = "topks"
 
 	HTTPReturnHas = "has"
 

diff --git a/internal/distributed/proxy/httpserver/handler_v2.go b/internal/distributed/proxy/httpserver/handler_v2.go
@@ -1198,9 +1198,9 @@ func (h *HandlersV2) search(ctx context.Context, c *gin.Context, anyReq any, dbN
 				})
 			} else {
 				if len(searchResp.Results.Recalls) > 0 {
-					HTTPReturnStream(c, http.StatusOK, gin.H{HTTPReturnCode: merr.Code(nil), HTTPReturnData: outputData, HTTPReturnCost: cost, HTTPReturnRecalls: searchResp.Results.Recalls})
+					HTTPReturnStream(c, http.StatusOK, gin.H{HTTPReturnCode: merr.Code(nil), HTTPReturnData: outputData, HTTPReturnCost: cost, HTTPReturnRecalls: searchResp.Results.Recalls, HTTPReturnTopks: searchResp.Results.Topks})
 				} else {
-					HTTPReturnStream(c, http.StatusOK, gin.H{HTTPReturnCode: merr.Code(nil), HTTPReturnData: outputData, HTTPReturnCost: cost})
+					HTTPReturnStream(c, http.StatusOK, gin.H{HTTPReturnCode: merr.Code(nil), HTTPReturnData: outputData, HTTPReturnCost: cost, HTTPReturnTopks: searchResp.Results.Topks})
 				}
 			}
 		}
@@ -1302,7 +1302,7 @@ func (h *HandlersV2) advancedSearch(ctx context.Context, c *gin.Context, anyReq
 					HTTPReturnMessage: merr.ErrInvalidSearchResult.Error() + ", error: " + err.Error(),
 				})
 			} else {
-				HTTPReturnStream(c, http.StatusOK, gin.H{HTTPReturnCode: merr.Code(nil), HTTPReturnData: outputData, HTTPReturnCost: cost})
+				HTTPReturnStream(c, http.StatusOK, gin.H{HTTPReturnCode: merr.Code(nil), HTTPReturnData: outputData, HTTPReturnCost: cost, HTTPReturnTopks: searchResp.Results.Topks})
 			}
 		}
 	}

diff --git a/tests/restful_client_v2/testcases/test_vector_operations.py b/tests/restful_client_v2/testcases/test_vector_operations.py
@@ -2596,6 +2596,9 @@ def test_advanced_search_vector_with_multi_float32_vector_datatype(self, nb, dim
         rsp = self.vector_client.vector_advanced_search(payload)
         assert rsp['code'] == 0
         assert len(rsp['data']) == 10
+        assert len(rsp['topks']) == 1
+        assert rsp['topks'][0] == 10
+
 
 
 @pytest.mark.L0