client/python/test_client.py at master · zcdb/client · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
# -*- coding: utf-8 -*-

import os
import sys
import numpy as np
sys.path.append(os.path.join(os.getcwd(), 'vearchpb'))
import client
import data_model_pb2 as dpb
import router_grpc_pb2 as rpb


TotalNum = 10
Address = "localhost:9002"
DBName = 'ts_db'
SpaceName = 'ts_space'


def println(*args):
    print(*args)


def random_request(fields):
    request = {}
    for field in fields:
        k, v = field.name, field.data_type
        if v == dpb.FieldType.INT:
            request[k] = np.random.randint(TotalNum, dtype=np.int32)
        elif v == dpb.FieldType.LONG:
            request[k] = np.random.randint(TotalNum)
        elif v == dpb.FieldType.FLOAT:
            request[k] = np.random.randn()
        elif v == dpb.FieldType.DOUBLE:
            request[k] = np.random.randn()
        elif v == dpb.FieldType.STRING:
            request[k] = str(np.random.randn())
        elif v == dpb.FieldType.VECTOR:
            dimension = field.vector_meta_info.dimension
            request[k] = [abs(n) for n in np.random.randn(dimension)]
    return request


def test_conn():
    c = client.Client(Address, DBName, SpaceName)
    c.init()
    c.close()


def test_add():
    c = client.Client(Address, DBName, SpaceName)
    c.init()
    for i in range(TotalNum):
        request = random_request(c.table.table_meta_info.field_meta_info)
        request['_id'] = str(i)
        res = c.add(request)
        println(res)
        println(request)
        println('The result of get is {}'.format(c.get([str(i)])))
    c.close()


def test_update():
    c = client.Client(Address, DBName, SpaceName)
    c.init()
    fields = c.table.table_meta_info.field_meta_info
    field_name = ''
    for field in fields:
        if field.data_type == dpb.FieldType.STRING:
            field_name = field.name
            break
    for i in range(TotalNum):
        request = random_request(fields)
        request['_id'] = str(i)
        request[field_name] = 'update' + request[field_name]
        c.update(request)
        println('The result of get is {}'.format(c.get([str(i)])))
    c.close


def test_get():
    c = client.Client(Address, DBName, SpaceName)
    c.init()
    for i in range(TotalNum):
        res = c.get([str(i)])
        println(res)
    c.close()


def test_search():
    c = client.Client(Address, DBName, SpaceName)
    c.init()

    vector_name = ''
    for k, v in c.field_map.items():
        if v == dpb.FieldType.VECTOR:
            vector_name = k
            break

    for i in range(TotalNum):
        request = random_request(c.table.table_meta_info.field_meta_info)
        request['_id'] = str(i)
        vec_field = rpb.VectorQuery(name=vector_name,
                value=np.array(request[vector_name], dtype=np.float32).tobytes(),
                min_score=0.0,
                max_score=100.0)

        res = c.search(10, vec_fields=[vec_field], fields=request.keys())
        println('\n', res)
    c.close()


def test_msearch():
    c = client.Client(Address, DBName, SpaceName)
    c.init()
    vector_name = ''
    for k, v in c.field_map.items():
        if v == dpb.FieldType.VECTOR:
            vector_name = k
            break
    msearch_request = rpb.MSearchRequest(head=c.head)
    for i in range(TotalNum):
        request = random_request(c.table.table_meta_info.field_meta_info)
        vec_field = rpb.VectorQuery(name=vector_name,
                value=np.array(request[vector_name], dtype=np.float32).tobytes(),
                min_score=0.0,
                max_score=1000.0)
        search_request = rpb.SearchRequest(head=c.head,
                req_num=TotalNum,
                topN=10,
                vec_fields=[vec_field],
                fields=['_id'])
        msearch_request.search_requests.append(search_request)
    res = c.msearch(msearch_request)
    println(res)
    c.close()


def test_delete():
    c = client.Client(Address, DBName, SpaceName)
    c.init()
    for i in range(TotalNum):
        res = c.get([str(i)])
        #  println(res)
    c.close()