Skip to content

Commit 8694a1c

Browse files
PresburgerPresburger
committed
expose knowhere config as a golang package
Signed-off-by: yusheng.ma <yusheng.ma@zilliz.com>
1 parent e908223 commit 8694a1c

File tree

4 files changed

+158
-14
lines changed

4 files changed

+158
-14
lines changed

‎go.mod

+3
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
module github.com/zilliztech/knowhere
2+
3+
go 1.22.5

‎include/knowhere/config.h

+14-14
Original file line numberDiff line numberDiff line change
@@ -78,13 +78,13 @@ enum PARAM_TYPE {
7878

7979
template <>
8080
struct Entry<CFG_STRING> {
81-
explicit Entry<CFG_STRING>(CFG_STRING* v) {
81+
explicit Entry(CFG_STRING* v) {
8282
val = v;
8383
type = 0x0;
8484
default_val = std::nullopt;
8585
desc = std::nullopt;
8686
}
87-
Entry<CFG_STRING>() {
87+
Entry() {
8888
val = nullptr;
8989
type = 0x0;
9090
default_val = std::nullopt;
@@ -99,14 +99,14 @@ struct Entry<CFG_STRING> {
9999

100100
template <>
101101
struct Entry<CFG_FLOAT> {
102-
explicit Entry<CFG_FLOAT>(CFG_FLOAT* v) {
102+
explicit Entry(CFG_FLOAT* v) {
103103
val = v;
104104
default_val = std::nullopt;
105105
type = 0x0;
106106
range = std::nullopt;
107107
desc = std::nullopt;
108108
}
109-
Entry<CFG_FLOAT>() {
109+
Entry() {
110110
val = nullptr;
111111
default_val = std::nullopt;
112112
type = 0x0;
@@ -124,14 +124,14 @@ struct Entry<CFG_FLOAT> {
124124

125125
template <>
126126
struct Entry<CFG_INT> {
127-
explicit Entry<CFG_INT>(CFG_INT* v) {
127+
explicit Entry(CFG_INT* v) {
128128
val = v;
129129
default_val = std::nullopt;
130130
type = 0x0;
131131
range = std::nullopt;
132132
desc = std::nullopt;
133133
}
134-
Entry<CFG_INT>() {
134+
Entry() {
135135
val = nullptr;
136136
default_val = std::nullopt;
137137
type = 0x0;
@@ -149,14 +149,14 @@ struct Entry<CFG_INT> {
149149

150150
template <>
151151
struct Entry<CFG_LIST> {
152-
explicit Entry<CFG_LIST>(CFG_LIST* v) {
152+
explicit Entry(CFG_LIST* v) {
153153
val = v;
154154
default_val = std::nullopt;
155155
type = 0x0;
156156
desc = std::nullopt;
157157
}
158158

159-
Entry<CFG_LIST>() {
159+
Entry() {
160160
val = nullptr;
161161
default_val = std::nullopt;
162162
type = 0x0;
@@ -172,14 +172,14 @@ struct Entry<CFG_LIST> {
172172

173173
template <>
174174
struct Entry<CFG_BOOL> {
175-
explicit Entry<CFG_BOOL>(CFG_BOOL* v) {
175+
explicit Entry(CFG_BOOL* v) {
176176
val = v;
177177
default_val = std::nullopt;
178178
type = 0x0;
179179
desc = std::nullopt;
180180
}
181181

182-
Entry<CFG_BOOL>() {
182+
Entry() {
183183
val = nullptr;
184184
default_val = std::nullopt;
185185
type = 0x0;
@@ -195,14 +195,14 @@ struct Entry<CFG_BOOL> {
195195

196196
template <>
197197
struct Entry<CFG_MATERIALIZED_VIEW_SEARCH_INFO_TYPE> {
198-
explicit Entry<CFG_MATERIALIZED_VIEW_SEARCH_INFO_TYPE>(CFG_MATERIALIZED_VIEW_SEARCH_INFO_TYPE* v) {
198+
explicit Entry(CFG_MATERIALIZED_VIEW_SEARCH_INFO_TYPE* v) {
199199
val = v;
200200
default_val = std::nullopt;
201201
type = 0x0;
202202
desc = std::nullopt;
203203
}
204204

205-
Entry<CFG_MATERIALIZED_VIEW_SEARCH_INFO_TYPE>() {
205+
Entry() {
206206
val = nullptr;
207207
default_val = std::nullopt;
208208
type = 0x0;
@@ -218,14 +218,14 @@ struct Entry<CFG_MATERIALIZED_VIEW_SEARCH_INFO_TYPE> {
218218

219219
template <>
220220
struct Entry<CFG_BYTES> {
221-
explicit Entry<CFG_BYTES>(CFG_BYTES* v) {
221+
explicit Entry(CFG_BYTES* v) {
222222
val = v;
223223
default_val = std::nullopt;
224224
type = 0x0;
225225
desc = std::nullopt;
226226
}
227227

228-
Entry<CFG_BYTES>() {
228+
Entry() {
229229
val = nullptr;
230230
default_val = std::nullopt;
231231
type = 0x0;

‎knowhere.go

+83
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
package knowhere
2+
3+
/*
4+
#cgo LDFLAGS: -L${SRCDIR}/build/Release -lknowhere
5+
#include <stdlib.h>
6+
int CheckConfig(int index_type, char const* str, int n, int param_type);
7+
*/
8+
import "C"
9+
10+
import (
11+
"encoding/json"
12+
"errors"
13+
"fmt"
14+
"unsafe"
15+
)
16+
17+
func CheckConfig(name string, params map[string]string, ptype string) error {
18+
19+
mJson, err := json.Marshal(params)
20+
if err != nil {
21+
fmt.Println(err.Error())
22+
return err
23+
}
24+
jsonStr := string(mJson)
25+
26+
var indexType int
27+
switch name {
28+
case "FLAT":
29+
indexType = 0
30+
case "DISKANN":
31+
indexType = 1
32+
case "HNSW":
33+
indexType = 2
34+
case "IVF_FLAT":
35+
indexType = 3
36+
case "IVF_PQ":
37+
indexType = 4
38+
case "GPU_CAGRA":
39+
indexType = 5
40+
case "GPU_IVF_PQ":
41+
indexType = 6
42+
case "GPU_IVF_FLAT":
43+
indexType = 7
44+
case "GPU_BRUTE_FORCE":
45+
indexType = 8
46+
default:
47+
return errors.New("index is not supported.")
48+
}
49+
50+
var paramType int
51+
switch ptype {
52+
case "TRAIN":
53+
paramType = 1 << 0
54+
case "SEARCH":
55+
paramType = 1 << 1
56+
case "RANGE_SEARCH":
57+
paramType = 1 << 2
58+
case "FEDER":
59+
paramType = 1 << 3
60+
case "DESERIALIZE":
61+
paramType = 1 << 4
62+
case "DESERIALIZE_FROM_FILE":
63+
paramType = 1 << 5
64+
case "ITERATOR":
65+
paramType = 1 << 6
66+
case "CLUSTER":
67+
paramType = 1 << 7
68+
default:
69+
return errors.New("param type is not supported.")
70+
}
71+
72+
c_str := C.CString(jsonStr)
73+
ret, err := C.CheckConfig(C.int(indexType), c_str, C.int(len(jsonStr)), C.int(paramType))
74+
C.free(unsafe.Pointer(c_str))
75+
if err != nil {
76+
return err
77+
}
78+
if ret == 0 {
79+
return nil
80+
}
81+
82+
return errors.New("Config check failed.")
83+
}

‎src/common/config.cc

+58
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,17 @@
1111

1212
#include "knowhere/config.h"
1313

14+
#include "index/diskann/diskann_config.h"
15+
#include "index/flat/flat_config.h"
16+
#include "index/gpu_raft/gpu_raft_brute_force_config.h"
17+
#include "index/gpu_raft/gpu_raft_cagra_config.h"
18+
#include "index/gpu_raft/gpu_raft_ivf_flat_config.h"
19+
#include "index/gpu_raft/gpu_raft_ivf_pq_config.h"
20+
#include "index/hnsw/hnsw_config.h"
21+
#include "index/ivf/ivf_config.h"
22+
#include "index/sparse/sparse_inverted_index_config.h"
1423
#include "knowhere/log.h"
24+
1525
namespace knowhere {
1626

1727
static const std::unordered_set<std::string> ext_legal_json_keys = {"metric_type",
@@ -108,4 +118,52 @@ Config::FormatAndCheck(const Config& cfg, Json& json, std::string* const err_msg
108118
}
109119
return Status::success;
110120
}
121+
111122
} // namespace knowhere
123+
124+
extern "C" __attribute__((visibility("default"))) int
125+
CheckConfig(int index_type, char const* str, int n, int param_type);
126+
127+
int
128+
CheckConfig(int index_type, const char* str, int n, int param_type) {
129+
knowhere::Json json = knowhere::Json::parse(str, str + n);
130+
std::unique_ptr<knowhere::Config> cfg;
131+
132+
switch (index_type) {
133+
case 0:
134+
cfg = std::make_unique<knowhere::FlatConfig>();
135+
break;
136+
case 1:
137+
cfg = std::make_unique<knowhere::DiskANNConfig>();
138+
break;
139+
case 2:
140+
cfg = std::make_unique<knowhere::HnswConfig>();
141+
break;
142+
case 3:
143+
cfg = std::make_unique<knowhere::IvfFlatConfig>();
144+
break;
145+
case 4:
146+
cfg = std::make_unique<knowhere::IvfPqConfig>();
147+
break;
148+
case 5:
149+
cfg = std::make_unique<knowhere::GpuRaftCagraConfig>();
150+
break;
151+
case 6:
152+
cfg = std::make_unique<knowhere::GpuRaftIvfPqConfig>();
153+
break;
154+
case 7:
155+
cfg = std::make_unique<knowhere::GpuRaftIvfFlatConfig>();
156+
break;
157+
case 8:
158+
cfg = std::make_unique<knowhere::GpuRaftBruteForceConfig>();
159+
break;
160+
default:
161+
return int(knowhere::Status::invalid_args);
162+
}
163+
164+
auto res = knowhere::Config::FormatAndCheck(*cfg, json, nullptr);
165+
if (res != knowhere::Status::success) {
166+
return int(res);
167+
}
168+
return int(knowhere::Config::Load(*cfg, json, knowhere::PARAM_TYPE(param_type), nullptr));
169+
}

0 commit comments

Comments
 (0)