From c6f2cd8e5a10122ebc256df78ad17051472341ba Mon Sep 17 00:00:00 2001 From: congqixia Date: Mon, 10 Jul 2023 10:12:23 +0800 Subject: [PATCH] Support probing diskann (#164) Signed-off-by: Congqi Xia --- states/probe.go | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/states/probe.go b/states/probe.go index 83d6b79..ed8c0a0 100644 --- a/states/probe.go +++ b/states/probe.go @@ -468,6 +468,26 @@ func getMockSearchRequest(ctx context.Context, cli clientv3.KV, basePath string, DmlChannels: []string{}, } return r, nil + case "DISKANN": + metricType := common.GetKVPair(index.GetIndexInfo().GetIndexParams(), "metric_type") + if metricType == "" { + metricType = common.GetKVPair(index.GetIndexInfo().GetTypeParams(), "metric_type") + if metricType == "" { + fmt.Println("no metric_type in IndexParams or TypeParams") + return nil, fmt.Errorf("no metric_type in IndexParams or TypeParams, bad meta") + } + fmt.Println("metric_type is in TypeParams instead of IndexParams") + } + + topK := int64(10) + spStr := genSearchDISKANNParamBytes(20) + req.SerializedExprPlan = getSearchPlan(vectorField.DataType == models.DataTypeBinaryVector, pkField.FieldID, vectorField.FieldID, topK, metricType, string(spStr)) + r := &querypbv2.SearchRequest{ + Req: req, + FromShardLeader: false, + DmlChannels: []string{}, + } + return r, nil default: return nil, fmt.Errorf("probing index type %s not supported yet", indexType) @@ -506,6 +526,13 @@ func genSearchHNSWParamBytes(ef int64) []byte { return bs } +func genSearchDISKANNParamBytes(searchList int) []byte { + m := make(map[string]any) + m["search_list"] = searchList + bs, _ := json.Marshal(m) + return bs +} + func genFloatVector(dim int64) FloatVector { result := make([]float32, 0, dim)