Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
BaiXuePrincess
milvus
提交
508131e4
milvus
项目概览
BaiXuePrincess
/
milvus
与 Fork 源项目一致
从无法访问的项目Fork
通知
7
Star
4
Fork
0
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
Issue
0
列表
看板
标记
里程碑
合并请求
0
Wiki
0
Wiki
分析
仓库
DevOps
项目成员
Pages
milvus
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Issue
0
Issue
0
列表
看板
标记
里程碑
合并请求
0
合并请求
0
Pages
分析
分析
仓库分析
DevOps
Wiki
0
Wiki
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
创建新Issue
提交
Issue看板
体验新版 GitCode,发现更多精彩内容 >>
提交
508131e4
编写于
6月 10, 2019
作者:
Y
Yang Xuan
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
feat(python): complete sdk 0.0.1
Former-commit-id: 713b3a629fcf86291d1e35c1d1b038b6fc599cb8
上级
964d3f37
变更
6
隐藏空白更改
内联
并排
Showing
6 changed file
with
353 addition
and
727 deletion
+353
-727
python/sdk/client/Abstract.py
python/sdk/client/Abstract.py
+89
-235
python/sdk/client/Client.py
python/sdk/client/Client.py
+106
-269
python/sdk/client/Exceptions.py
python/sdk/client/Exceptions.py
+0
-4
python/sdk/client/Status.py
python/sdk/client/Status.py
+8
-6
python/sdk/examples/connection_exp.py
python/sdk/examples/connection_exp.py
+69
-56
python/sdk/tests/TestClient.py
python/sdk/tests/TestClient.py
+81
-157
未找到文件。
python/sdk/client/Abstract.py
浏览文件 @
508131e4
from
enum
import
IntEnum
from
enum
import
IntEnum
from
.Exceptions
import
ConnectParamMissingError
class
AbstactIndexType
(
object
):
class
IndexType
(
IntEnum
):
RAW
=
'1'
INVALIDE
=
0
IVFFLAT
=
'2'
IDMAP
=
1
IVFLAT
=
2
class
AbstractColumnType
(
object
):
INVALID
=
1
INT8
=
2
INT16
=
3
INT32
=
4
INT64
=
5
FLOAT32
=
6
FLOAT64
=
7
DATE
=
8
VECTOR
=
9
class
Column
(
object
):
"""
Table column description
:type type: ColumnType
:param type: (Required) type of the column
:type name: str
:param name: (Required) name of the column
"""
def
__init__
(
self
,
name
=
None
,
type
=
AbstractColumnType
.
INVALID
):
self
.
type
=
type
self
.
name
=
name
class
VectorColumn
(
Column
):
"""
Table vector column description
:type dimension: int, int64
:param dimension: (Required) vector dimension
:type index_type: string IndexType
:param index_type: (Required) IndexType
:type store_raw_vector: bool
:param store_raw_vector: (Required) Is vector self stored in the table
`Column`:
:type name: str
:param name: (Required) Name of the column
:type type: ColumnType
:param type: (Required) Default type is ColumnType.VECTOR, can't change
"""
def
__init__
(
self
,
name
,
dimension
=
0
,
index_type
=
None
,
store_raw_vector
=
False
,
type
=
None
):
self
.
dimension
=
dimension
self
.
index_type
=
index_type
self
.
store_raw_vector
=
store_raw_vector
super
(
VectorColumn
,
self
).
__init__
(
name
,
type
=
type
)
class
TableSchema
(
object
):
class
TableSchema
(
object
):
...
@@ -74,30 +14,26 @@ class TableSchema(object):
...
@@ -74,30 +14,26 @@ class TableSchema(object):
:type table_name: str
:type table_name: str
:param table_name: (Required) name of table
:param table_name: (Required) name of table
:type vector_columns: list[VectorColumn]
:type index_type: IndexType
:param vector_columns: (Required) a list of VectorColumns,
:param index_type: (Optional) index type, default = 0
Stores different types of vectors
:type attribute_columns: list[Column]
:param attribute_columns: (Optional) Columns description
List of `Columns` whose type isn't VECTOR
`IndexType`: 0-invalid, 1-idmap, 2-ivflat
:type
partition_column_names: list[str]
:type
dimension: int64
:param
partition_column_names: (Optional) Partition column name
:param
dimension: (Required) dimension of vector
`Partition columns` are `attribute columns`, the number of
:type store_raw_vector: bool
partition columns may be less than or equal to attribute columns,
:param store_raw_vector: (Optional) default = False
this param only stores `column name`
"""
"""
def
__init__
(
self
,
table_name
,
vector_columns
,
def
__init__
(
self
,
table_name
,
attribute_columns
,
partition_column_names
,
**
kwargs
):
dimension
=
0
,
index_type
=
IndexType
.
INVALIDE
,
store_raw_vector
=
False
):
self
.
table_name
=
table_name
self
.
table_name
=
table_name
self
.
vector_columns
=
vector_columns
self
.
index_type
=
index_type
self
.
attribute_columns
=
attribute_columns
self
.
dimension
=
dimension
self
.
partition_column_names
=
partition_column_names
self
.
store_raw_vector
=
store_raw_vector
class
Range
(
object
):
class
Range
(
object
):
...
@@ -105,10 +41,10 @@ class Range(object):
...
@@ -105,10 +41,10 @@ class Range(object):
Range information
Range information
:type start: str
:type start: str
:param start:
(Required)
Range start value
:param start: Range start value
:type end: str
:type end: str
:param end:
(Required)
Range end value
:param end: Range end value
"""
"""
def
__init__
(
self
,
start
,
end
):
def
__init__
(
self
,
start
,
end
):
...
@@ -116,97 +52,37 @@ class Range(object):
...
@@ -116,97 +52,37 @@ class Range(object):
self
.
end
=
end
self
.
end
=
end
class
CreateTablePartitionParam
(
object
):
"""
Create table partition parameters
:type table_name: str
:param table_name: (Required) Table name,
VECTOR/FLOAT32/FLOAT64 ColumnType is not allowed for partition
:type partition_name: str
:param partition_name: (Required) partition name, created partition name
:type column_name_to_range: dict{str : Range}
:param column_name_to_range: (Required) Column name to PartitionRange dictionary
"""
# TODO Iterable
def
__init__
(
self
,
table_name
,
partition_name
,
column_name_to_range
):
self
.
table_name
=
table_name
self
.
partition_name
=
partition_name
self
.
column_name_to_range
=
column_name_to_range
class
DeleteTablePartitionParam
(
object
):
"""
Delete table partition parameters
:type table_name: str
:param table_name: (Required) Table name
:type partition_names: iterable, str
:param partition_names: (Required) Partition name array
"""
# TODO Iterable
def
__init__
(
self
,
table_name
,
partition_names
):
self
.
table_name
=
table_name
self
.
partition_names
=
partition_names
class
RowRecord
(
object
):
class
RowRecord
(
object
):
"""
"""
Record inserted
Record inserted
:type
column_name_to_vector: dict{str : list[float]}
:type
vector_data: binary str
:param
column_name_to_vector: (Required) Column name to vector map
:param
vector_data: (Required) a vector
:type column_name_to_attribute: dict{str: str}
:param column_name_to_attribute: (Optional) Other attribute columns
"""
"""
def
__init__
(
self
,
column_name_to_vector
,
column_name_to_attribute
):
def
__init__
(
self
,
vector_data
):
self
.
column_name_to_vector
=
column_name_to_vector
self
.
vector_data
=
vector_data
self
.
column_name_to_attribute
=
column_name_to_attribute
class
QueryRecord
(
object
):
"""
Query record
:type column_name_to_vector: (Required) dict{str : list[float]}
:param column_name_to_vector: Query vectors, column name to vector map
:type selected_columns: list[str]
:param selected_columns: (Optional) Output column array
:type name_to_partition_ranges: dict{str : list[Range]}
:param name_to_partition_ranges: (Optional) Range used to select partitions
"""
def
__init__
(
self
,
column_name_to_vector
,
selected_columns
,
name_to_partition_ranges
):
self
.
column_name_to_vector
=
column_name_to_vector
self
.
selected_columns
=
selected_columns
self
.
name_to_partition_ranges
=
name_to_partition_ranges
class
QueryResult
(
object
):
class
QueryResult
(
object
):
"""
"""
Query result
Query result
:type id: int
:type id: int
64
:param id:
Output result
:param id:
id of the vector
:type score: float
:type score: float
:param score: Vector similarity 0 <= score <= 100
:param score: Vector similarity 0 <= score <= 100
:type column_name_to_attribute: dict{str : str}
:param column_name_to_attribute: Other columns
"""
"""
def
__init__
(
self
,
id
,
score
,
column_name_to_attribute
):
def
__init__
(
self
,
id
,
score
):
self
.
id
=
id
self
.
id
=
id
self
.
score
=
score
self
.
score
=
score
self
.
column_name_to_value
=
column_name_to_attribute
def
__repr__
(
self
):
L
=
[
'%s=%r'
%
(
key
,
value
)
for
key
,
value
in
self
.
__dict__
.
items
()]
return
'%s(%s)'
%
(
self
.
__class__
.
__name__
,
', '
.
join
(
L
))
class
TopKQueryResult
(
object
):
class
TopKQueryResult
(
object
):
...
@@ -220,6 +96,12 @@ class TopKQueryResult(object):
...
@@ -220,6 +96,12 @@ class TopKQueryResult(object):
def
__init__
(
self
,
query_results
):
def
__init__
(
self
,
query_results
):
self
.
query_results
=
query_results
self
.
query_results
=
query_results
def
__repr__
(
self
):
L
=
[
'%s=%r'
%
(
key
,
value
)
for
key
,
value
in
self
.
__dict__
.
items
()]
return
'%s(%s)'
%
(
self
.
__class__
.
__name__
,
', '
.
join
(
L
))
def
_abstract
():
def
_abstract
():
raise
NotImplementedError
(
'You need to override this function'
)
raise
NotImplementedError
(
'You need to override this function'
)
...
@@ -232,114 +114,71 @@ class ConnectIntf(object):
...
@@ -232,114 +114,71 @@ class ConnectIntf(object):
"""
"""
@
staticmethod
def
connect
(
self
,
host
=
None
,
port
=
None
,
uri
=
None
):
def
create
():
"""Create a connection instance and return it
should be implemented
:return connection: Connection
"""
_abstract
()
@
staticmethod
def
destroy
(
connection
):
"""Destroy the connection instance
should be implemented
:type connection: Connection
:param connection: The connection instance to be destroyed
:return bool, return True if destroy is successful
"""
_abstract
()
def
connect
(
self
,
param
=
None
,
uri
=
None
):
"""
"""
Connect method should be called before any operations
Connect method should be called before any operations
Server will be connected after connect return OK
Server will be connected after connect return OK
should be implemented
Should be implemented
:type host: str
:param host: host
:type p
aram: ConnectParam
:type p
ort: str
:param p
aram: ConnectParam
:param p
ort: port
:type uri: str
:type uri: str
:param uri:
uri param
:param uri:
(Optional) uri
:return
:
Status, indicate if connect is successful
:return Status, indicate if connect is successful
"""
"""
if
(
not
param
and
not
uri
)
or
(
param
and
uri
):
raise
ConnectParamMissingError
(
'You need to parse exact one param'
)
_abstract
()
_abstract
()
def
connected
(
self
):
def
connected
(
self
):
"""
"""
connected, connection status
connected, connection status
s
hould be implemented
S
hould be implemented
:return
:
Status, indicate if connect is successful
:return Status, indicate if connect is successful
"""
"""
_abstract
()
_abstract
()
def
disconnect
(
self
):
def
disconnect
(
self
):
"""
"""
Disconnect, server will be disconnected after disconnect return
OK
Disconnect, server will be disconnected after disconnect return
SUCCESS
s
hould be implemented
S
hould be implemented
:return
:
Status, indicate if connect is successful
:return Status, indicate if connect is successful
"""
"""
_abstract
()
_abstract
()
def
create_table
(
self
,
param
):
def
create_table
(
self
,
param
):
"""
"""
Create table
Create table
s
hould be implemented
S
hould be implemented
:type param: TableSchema
:type param: TableSchema
:param param: provide table information to be created
:param param: provide table information to be created
:return
:
Status, indicate if connect is successful
:return Status, indicate if connect is successful
"""
"""
_abstract
()
_abstract
()
def
delete_table
(
self
,
table_name
):
def
delete_table
(
self
,
table_name
):
"""
"""
Delete table
Delete table
s
hould be implemented
S
hould be implemented
:type table_name: str
:type table_name: str
:param table_name: table_name of the deleting table
:param table_name: table_name of the deleting table
:return
:
Status, indicate if connect is successful
:return Status, indicate if connect is successful
"""
"""
_abstract
()
_abstract
()
def
create_table_partition
(
self
,
param
):
def
add_vectors
(
self
,
table_name
,
records
):
"""
Create table partition
should be implemented
:type param: CreateTablePartitionParam
:param param: provide partition information
:return: Status, indicate if table partition is created successfully
"""
_abstract
()
def
delete_table_partition
(
self
,
param
):
"""
Delete table partition
should be implemented
:type param: DeleteTablePartitionParam
:param param: provide partition information to be deleted
:return: Status, indicate if partition is deleted successfully
"""
_abstract
()
def
add_vector
(
self
,
table_name
,
records
):
"""
"""
Add vectors to table
Add vectors to table
s
hould be implemented
S
hould be implemented
:type table_name: str
:type table_name: str
:param table_name: table name been inserted
:param table_name: table name been inserted
...
@@ -347,27 +186,31 @@ class ConnectIntf(object):
...
@@ -347,27 +186,31 @@ class ConnectIntf(object):
:type records: list[RowRecord]
:type records: list[RowRecord]
:param records: list of vectors been inserted
:param records: list of vectors been inserted
:returns
:
:returns
Status : indicate if vectors inserted successfully
Status : indicate if vectors inserted successfully
ids :list of id, after inserted every vector is given a id
ids :list of id, after inserted every vector is given a id
"""
"""
_abstract
()
_abstract
()
def
search_vector
(
self
,
table_name
,
query_record
s
,
top_k
):
def
search_vector
s
(
self
,
table_name
,
query_records
,
query_range
s
,
top_k
):
"""
"""
Query vectors in a table
Query vectors in a table
s
hould be implemented
S
hould be implemented
:type table_name: str
:type table_name: str
:param table_name: table name been queried
:param table_name: table name been queried
:type query_records: list[
Query
Record]
:type query_records: list[
Row
Record]
:param query_records: all vectors going to be queried
:param query_records: all vectors going to be queried
:type query_ranges: list[Range]
:param query_ranges: Optional ranges for conditional search.
If not specified, search whole table
:type top_k: int
:type top_k: int
:param top_k: how many similar vectors will be searched
:param top_k: how many similar vectors will be searched
:returns
:
:returns
Status: indicate if query is successful
Status: indicate if query is successful
query_results: list[TopKQueryResult]
query_results: list[TopKQueryResult]
"""
"""
...
@@ -376,23 +219,37 @@ class ConnectIntf(object):
...
@@ -376,23 +219,37 @@ class ConnectIntf(object):
def
describe_table
(
self
,
table_name
):
def
describe_table
(
self
,
table_name
):
"""
"""
Show table information
Show table information
s
hould be implemented
S
hould be implemented
:type table_name: str
:type table_name: str
:param table_name: which table to be shown
:param table_name: which table to be shown
:returns
:
:returns
Status: indicate if query is successful
Status: indicate if query is successful
table_schema: TableSchema, given when operation is successful
table_schema: TableSchema, given when operation is successful
"""
"""
_abstract
()
_abstract
()
def
get_table_row_count
(
self
,
table_name
):
"""
Get table row count
Should be implemented
:type table_name, str
:param table_name, target table name.
:returns
Status: indicate if operation is successful
count: int, table row count
"""
_abstract
()
def
show_tables
(
self
):
def
show_tables
(
self
):
"""
"""
Show all tables in database
Show all tables in database
should be implemented
should be implemented
:return
:
:return
Status: indicate if this operation is successful
Status: indicate if this operation is successful
tables: list[str], list of table names
tables: list[str], list of table names
"""
"""
...
@@ -403,31 +260,28 @@ class ConnectIntf(object):
...
@@ -403,31 +260,28 @@ class ConnectIntf(object):
Provide client version
Provide client version
should be implemented
should be implemented
:return:
C
lient version
:return:
str, c
lient version
"""
"""
_abstract
()
_abstract
()
pass
def
server_version
(
self
):
def
server_version
(
self
):
"""
"""
Provide server version
Provide server version
should be implemented
should be implemented
:return:
S
erver version
:return:
str, s
erver version
"""
"""
_abstract
()
def
server_status
(
self
,
cmd
):
def
server_status
(
self
,
cmd
):
"""
"""
Provide server status
Provide server status
should be implemented
should be implemented
# TODO What is cmd
:type cmd, str
:type cmd
:param cmd
:return:
S
erver status
:return:
str, s
erver status
"""
"""
_abstract
()
_abstract
()
pass
...
...
python/sdk/client/Client.py
浏览文件 @
508131e4
...
@@ -9,21 +9,20 @@ from thrift.Thrift import TException, TApplicationException, TType
...
@@ -9,21 +9,20 @@ from thrift.Thrift import TException, TApplicationException, TType
from
megasearch.thrift
import
MegasearchService
from
megasearch.thrift
import
MegasearchService
from
megasearch.thrift
import
ttypes
from
megasearch.thrift
import
ttypes
from
client.Abstract
import
(
from
client.Abstract
import
(
ConnectIntf
,
TableSchema
,
ConnectIntf
,
AbstactIndexType
,
AbstractColumnType
,
TableSchema
,
Column
,
IndexType
,
VectorColumn
,
Range
,
Range
,
CreateTablePartitionParam
,
RowRecord
,
DeleteTablePartitionParam
,
QueryResult
,
RowRecord
,
QueryRecord
,
TopKQueryResult
QueryResult
,
TopKQueryResult
)
)
from
client.Status
import
Status
from
client.Status
import
Status
from
client.Exceptions
import
(
from
client.Exceptions
import
(
RepeatingConnectError
,
ConnectParamMissingError
,
RepeatingConnectError
,
DisconnectNotConnectedClientError
,
DisconnectNotConnectedClientError
,
ParamError
,
NotConnectError
NotConnectError
)
)
LOGGER
=
logging
.
getLogger
(
__name__
)
LOGGER
=
logging
.
getLogger
(
__name__
)
...
@@ -32,125 +31,35 @@ __VERSION__ = '0.0.1'
...
@@ -32,125 +31,35 @@ __VERSION__ = '0.0.1'
__NAME__
=
'Thrift_Client'
__NAME__
=
'Thrift_Client'
class
IndexType
(
AbstactIndexType
):
# TODO thrift in IndexType
RAW
=
'1'
IVFFLAT
=
'2'
class
ColumnType
(
AbstractColumnType
):
FLOAT32
=
6
FLOAT64
=
7
DATE
=
8
INVALID
=
TType
.
STOP
INT8
=
TType
.
I08
INT16
=
TType
.
I16
INT32
=
TType
.
I32
INT64
=
TType
.
I64
VECTOR
=
TType
.
LIST
# TODO Required and Optional
# TODO Examples
# TODO ORM
class
Prepare
(
object
):
class
Prepare
(
object
):
@
classmethod
@
classmethod
def
column
(
cls
,
name
,
type
):
def
table_schema
(
cls
,
"""
table_name
,
*
,
Table column param
dimension
,
:param type: (Required) ColumnType, type of the column
index_type
,
:param name: (Required) str, name of the column
store_raw_vector
):
:return Column
"""
temp_column
=
Column
(
name
=
name
,
type
=
type
)
return
ttypes
.
Column
(
name
=
temp_column
.
name
,
type
=
temp_column
.
type
)
@
classmethod
def
vector_column
(
cls
,
name
,
dimension
,
index_type
=
IndexType
.
RAW
,
store_raw_vector
=
False
):
"""
"""
Table vector column description
:param dimension: (Required) int64, vector dimension
:param index_type: (Required) IndexType
:param store_raw_vector: (Required) Bool
`Column`:
:param name: (Required) Name of the column
:param type: (Required) Default type is ColumnType.VECTOR, can't change
:return VectorColumn
"""
temp
=
VectorColumn
(
name
=
name
,
dimension
=
dimension
,
index_type
=
index_type
,
store_raw_vector
=
store_raw_vector
)
base
=
ttypes
.
Column
(
name
=
temp
.
name
,
type
=
ColumnType
.
VECTOR
)
return
ttypes
.
VectorColumn
(
base
=
base
,
dimension
=
temp
.
dimension
,
store_raw_vector
=
temp
.
store_raw_vector
,
index_type
=
temp
.
index_type
)
# Without IndexType
# temp = VectorColumn(name=name, dimension=dimension,
# store_raw_vector=store_raw_vector)
# return ttypes.VectorColumn(base=base, dimension=temp.dimension,
# store_raw_vector=temp.store_raw_vector)
@
classmethod
def
table_schema
(
cls
,
table_name
,
vector_columns
,
attribute_columns
,
partition_column_names
):
"""
:param table_name: (Required) Name of the table
:param vector_columns: (Required) List of VectorColumns
`VectorColumn`:
:param table_name: str, (Required) name of table
- dimension: int, default = 0
:param index_type: IndexType, (Required) index type, default = IndexType.INVALID
Dimension of the vector, different vector_columns'
:param dimension: int64, (Optional) dimension of the table
dimension may vary
:param store_raw_vector: bool, (Optional) default = False
- index_type: (optional) IndexType, default=IndexType.RAW
Vector's index type
- store_raw_vector : (optional) bool, default=False
- name: str
Name of the column
- type: ColumnType, default=ColumnType.VECTOR, can't change
:param attribute_columns: (Optional) List of Columns. Attribute columns are Columns,
whose types aren't ColumnType.VECTOR
`Column`:
- name: str
- type: ColumnType, default=ColumnType.INVALID
:param partition_column_names: (Optional) List of str.
Partition columns name
indicates which attribute columns is used for partition, can
have lots of partition columns as long as:
-> No. partition_column_names <= No. attribute_columns
-> partition_column_names IN attribute_column_names
:return: TableSchema
:return: TableSchema
"""
"""
temp
=
TableSchema
(
table_name
,
vector_columns
,
temp
=
TableSchema
(
table_name
,
dimension
,
index_type
,
store_raw_vector
)
attribute_columns
,
partition_column_names
)
return
ttypes
.
TableSchema
(
table_name
=
temp
.
table_name
,
return
ttypes
.
TableSchema
(
table_name
=
temp
.
table_name
,
vector_column_array
=
temp
.
vector_columns
,
dimension
=
dimension
,
attribute_column_array
=
temp
.
attribute_columns
,
index_type
=
index_type
,
partition_column_name_array
=
temp
.
partition_column_names
)
store_raw_vector
=
store_raw_vector
)
@
classmethod
@
classmethod
def
range
(
cls
,
start
,
end
):
def
range
(
cls
,
start
,
end
):
"""
"""
:param start:
(Required) Partition range start value
:param start:
str, (Required) range start
:param end:
(Required) Partition range end value
:param end:
str (Required) range end
:return Range
:return Range
"""
"""
...
@@ -158,142 +67,66 @@ class Prepare(object):
...
@@ -158,142 +67,66 @@ class Prepare(object):
return
ttypes
.
Range
(
start_value
=
temp
.
start
,
end_value
=
temp
.
end
)
return
ttypes
.
Range
(
start_value
=
temp
.
start
,
end_value
=
temp
.
end
)
@
classmethod
@
classmethod
def
create_table_partition_param
(
cls
,
def
row_record
(
cls
,
vector_data
):
table_name
,
partition_name
,
column_name_to_range
):
"""
Create table partition parameters
:param table_name: (Required) str, Table name,
VECTOR/FLOAT32/FLOAT64 ColumnType is not allowed for partition
:param partition_name: (Required) str partition name, created partition name
:param column_name_to_range: (Required) dict, column name to partition range dictionary
:return CreateTablePartitionParam
"""
temp
=
CreateTablePartitionParam
(
table_name
=
table_name
,
partition_name
=
partition_name
,
column_name_to_range
=
column_name_to_range
)
return
ttypes
.
CreateTablePartitionParam
(
table_name
=
temp
.
table_name
,
partition_name
=
temp
.
partition_name
,
range_map
=
temp
.
column_name_to_range
)
@
classmethod
def
delete_table_partition_param
(
cls
,
table_name
,
partition_names
):
"""
Delete table partition parameters
:param table_name: (Required) Table name
:param partition_names: (Required) List of partition names
:return DeleteTablePartitionParam
"""
temp
=
DeleteTablePartitionParam
(
table_name
=
table_name
,
partition_names
=
partition_names
)
return
ttypes
.
DeleteTablePartitionParam
(
table_name
=
table_name
,
partition_name_array
=
partition_names
)
@
classmethod
def
row_record
(
cls
,
column_name_to_vector
,
column_name_to_attribute
):
"""
:param column_name_to_vector: (Required) dict{str : list[float]}
Column name to vector map
:param column_name_to_attribute: (Optional) dict{str: str}
Other attribute columns
"""
temp
=
RowRecord
(
column_name_to_vector
=
column_name_to_vector
,
column_name_to_attribute
=
column_name_to_attribute
)
return
ttypes
.
RowRecord
(
vector_map
=
temp
.
column_name_to_vector
,
attribute_map
=
temp
.
column_name_to_attribute
)
@
classmethod
def
query_record
(
cls
,
column_name_to_vector
,
selected_columns
,
name_to_partition_ranges
):
"""
"""
:param column_name_to_vector: (Required) dict{str : list[float]}
Record inserted
Query vectors, column name to vector map
:param selected_columns: (Optional) list[str_column_name]
:param vector_data: float binary str, (Required) a binary str
List of Output columns
:param name_to_partition_ranges: (Optional) dict{str : list[Range]}
Partition Range used to search
`Range`:
:param start: Partition range start value
:param end: Partition range end value
:return QueryRecord
"""
"""
temp
=
QueryRecord
(
column_name_to_vector
=
column_name_to_vector
,
temp
=
RowRecord
(
vector_data
)
selected_columns
=
selected_columns
,
return
ttypes
.
RowRecord
(
vector_data
=
temp
.
vector_data
)
name_to_partition_ranges
=
name_to_partition_ranges
)
return
ttypes
.
QueryRecord
(
vector_map
=
temp
.
column_name_to_vector
,
selected_column_array
=
temp
.
selected_columns
,
partition_filter_column_map
=
name_to_partition_ranges
)
class
MegaSearch
(
ConnectIntf
):
class
MegaSearch
(
ConnectIntf
):
def
__init__
(
self
):
def
__init__
(
self
):
self
.
transport
=
None
self
.
client
=
None
self
.
status
=
None
self
.
status
=
None
self
.
_transport
=
None
self
.
_client
=
None
def
__repr__
(
self
):
def
__repr__
(
self
):
return
'{}'
.
format
(
self
.
status
)
return
'{}'
.
format
(
self
.
status
)
@
staticmethod
def
create
():
# TODO in python, maybe this method is useless
return
MegaSearch
()
@
staticmethod
def
destroy
(
connection
):
"""Destroy the connection instance"""
# TODO in python, maybe this method is useless
pass
def
connect
(
self
,
host
=
'localhost'
,
port
=
'9090'
,
uri
=
None
):
def
connect
(
self
,
host
=
'localhost'
,
port
=
'9090'
,
uri
=
None
):
# TODO URI
# TODO URI
if
self
.
status
and
self
.
status
==
Status
(
message
=
'Connected'
)
:
if
self
.
status
and
self
.
status
==
Status
.
SUCCESS
:
raise
RepeatingConnectError
(
"You have already connected!"
)
raise
RepeatingConnectError
(
"You have already connected!"
)
transport
=
TSocket
.
TSocket
(
host
=
host
,
port
=
port
)
transport
=
TSocket
.
TSocket
(
host
=
host
,
port
=
port
)
self
.
transport
=
TTransport
.
TBufferedTransport
(
transport
)
self
.
_
transport
=
TTransport
.
TBufferedTransport
(
transport
)
protocol
=
TBinaryProtocol
.
TBinaryProtocol
(
transport
)
protocol
=
TBinaryProtocol
.
TBinaryProtocol
(
transport
)
self
.
client
=
MegasearchService
.
Client
(
protocol
)
self
.
_
client
=
MegasearchService
.
Client
(
protocol
)
try
:
try
:
transport
.
open
()
transport
.
open
()
self
.
status
=
Status
(
Status
.
OK
,
'Connected'
)
self
.
status
=
Status
(
Status
.
SUCCESS
,
'Connected'
)
LOGGER
.
info
(
'Connected!'
)
LOGGER
.
info
(
'Connected!'
)
except
(
TTransport
.
TTransportException
,
TException
)
as
e
:
except
(
TTransport
.
TTransportException
,
TException
)
as
e
:
self
.
status
=
Status
(
Status
.
INVALI
D
,
message
=
str
(
e
))
self
.
status
=
Status
(
Status
.
CONNECT_FAILE
D
,
message
=
str
(
e
))
LOGGER
.
error
(
'logger.error: {}'
.
format
(
self
.
status
))
LOGGER
.
error
(
'logger.error: {}'
.
format
(
self
.
status
))
finally
:
finally
:
return
self
.
status
return
self
.
status
@
property
@
property
def
connected
(
self
):
def
connected
(
self
):
return
self
.
status
==
Status
()
return
self
.
status
==
Status
.
SUCCESS
def
disconnect
(
self
):
def
disconnect
(
self
):
if
not
self
.
transport
:
if
not
self
.
_
transport
:
raise
DisconnectNotConnectedClientError
(
'Error'
)
raise
DisconnectNotConnectedClientError
(
'Error'
)
try
:
try
:
self
.
transport
.
close
()
self
.
_
transport
.
close
()
LOGGER
.
info
(
'Client Disconnected!'
)
LOGGER
.
info
(
'Client Disconnected!'
)
self
.
status
=
None
self
.
status
=
None
except
TException
as
e
:
except
TException
as
e
:
return
Status
(
Status
.
INVALI
D
,
str
(
e
))
return
Status
(
Status
.
PERMISSION_DENIE
D
,
str
(
e
))
return
Status
(
Status
.
OK
,
'Disconnected'
)
return
Status
(
Status
.
SUCCESS
,
'Disconnected'
)
def
create_table
(
self
,
param
):
def
create_table
(
self
,
param
):
"""Create table
"""Create table
...
@@ -304,15 +137,14 @@ class MegaSearch(ConnectIntf):
...
@@ -304,15 +137,14 @@ class MegaSearch(ConnectIntf):
:return: Status, indicate if operation is successful
:return: Status, indicate if operation is successful
"""
"""
if
not
self
.
client
:
if
not
self
.
_
client
:
raise
NotConnectError
(
'Please Connect to the server first!'
)
raise
NotConnectError
(
'Please Connect to the server first!'
)
try
:
try
:
LOGGER
.
error
(
param
)
self
.
_client
.
CreateTable
(
param
)
self
.
client
.
CreateTable
(
param
)
except
(
TApplicationException
,
)
as
e
:
except
(
TApplicationException
,
)
as
e
:
LOGGER
.
error
(
'Unable to create table'
)
LOGGER
.
error
(
'Unable to create table'
)
return
Status
(
Status
.
INVALI
D
,
str
(
e
))
return
Status
(
Status
.
PERMISSION_DENIE
D
,
str
(
e
))
return
Status
(
message
=
'Table {} created!'
.
format
(
param
.
table_name
))
return
Status
(
message
=
'Table {} created!'
.
format
(
param
.
table_name
))
def
delete_table
(
self
,
table_name
):
def
delete_table
(
self
,
table_name
):
...
@@ -323,48 +155,13 @@ class MegaSearch(ConnectIntf):
...
@@ -323,48 +155,13 @@ class MegaSearch(ConnectIntf):
:return: Status, indicate if operation is successful
:return: Status, indicate if operation is successful
"""
"""
try
:
try
:
self
.
client
.
DeleteTable
(
table_name
)
self
.
_
client
.
DeleteTable
(
table_name
)
except
(
TApplicationException
,
TException
)
as
e
:
except
(
TApplicationException
,
TException
)
as
e
:
LOGGER
.
error
(
'Unable to delete table {}'
.
format
(
table_name
))
LOGGER
.
error
(
'Unable to delete table {}'
.
format
(
table_name
))
return
Status
(
Status
.
INVALI
D
,
str
(
e
))
return
Status
(
Status
.
PERMISSION_DENIE
D
,
str
(
e
))
return
Status
(
message
=
'Table {} deleted!'
.
format
(
table_name
))
return
Status
(
message
=
'Table {} deleted!'
.
format
(
table_name
))
def
create_table_partition
(
self
,
param
):
def
add_vectors
(
self
,
table_name
,
records
):
"""
Create table partition
:type param: CreateTablePartitionParam, provide partition information
`Please use Prepare.create_table_partition_param generate param`
:return: Status, indicate if table partition is created successfully
"""
try
:
self
.
client
.
CreateTablePartition
(
param
)
except
(
TApplicationException
,
TException
)
as
e
:
LOGGER
.
error
(
'{}'
.
format
(
e
))
return
Status
(
Status
.
INVALID
,
str
(
e
))
return
Status
(
message
=
'Table partition created successfully!'
)
def
delete_table_partition
(
self
,
param
):
"""
Delete table partition
:type param: DeleteTablePartitionParam
:param param: provide partition information to be deleted
`Please use Prepare.delete_table_partition_param generate param`
:return: Status, indicate if partition is deleted successfully
"""
try
:
self
.
client
.
DeleteTablePartition
(
param
)
except
(
TApplicationException
,
TException
)
as
e
:
LOGGER
.
error
(
'{}'
.
format
(
e
))
return
Status
(
Status
.
INVALID
,
str
(
e
))
return
Status
(
message
=
'Table partition deleted successfully!'
)
def
add_vector
(
self
,
table_name
,
records
):
"""
"""
Add vectors to table
Add vectors to table
...
@@ -378,13 +175,13 @@ class MegaSearch(ConnectIntf):
...
@@ -378,13 +175,13 @@ class MegaSearch(ConnectIntf):
ids :list of id, after inserted every vector is given a id
ids :list of id, after inserted every vector is given a id
"""
"""
try
:
try
:
ids
=
self
.
client
.
AddVector
(
table_name
=
table_name
,
record_array
=
records
)
ids
=
self
.
_
client
.
AddVector
(
table_name
=
table_name
,
record_array
=
records
)
except
(
TApplicationException
,
TException
)
as
e
:
except
(
TApplicationException
,
TException
)
as
e
:
LOGGER
.
error
(
'{}'
.
format
(
e
))
LOGGER
.
error
(
'{}'
.
format
(
e
))
return
Status
(
Status
.
INVALI
D
,
str
(
e
)),
None
return
Status
(
Status
.
PERMISSION_DENIE
D
,
str
(
e
)),
None
return
Status
(
message
=
'Vector added successfully!'
),
ids
return
Status
(
message
=
'Vector
s
added successfully!'
),
ids
def
search_vector
(
self
,
table_name
,
query_records
,
top_k
):
def
search_vector
s
(
self
,
table_name
,
top_k
,
query_records
,
query_ranges
=
None
):
"""
"""
Query vectors in a table
Query vectors in a table
...
@@ -394,20 +191,29 @@ class MegaSearch(ConnectIntf):
...
@@ -394,20 +191,29 @@ class MegaSearch(ConnectIntf):
`Please use Prepare.query_record generate QueryRecord`
`Please use Prepare.query_record generate QueryRecord`
:param top_k: int, how many similar vectors will be searched
:param top_k: int, how many similar vectors will be searched
:param query_ranges, (Optional) list[Range], search range
:returns:
:returns:
Status: indicate if query is successful
Status: indicate if query is successful
query_result
s: list[TopKQueryResult], return when operation is successful
re
s: list[TopKQueryResult], return when operation is successful
"""
"""
# TODO topk_query_results
res
=
[]
try
:
try
:
topk_query_results
=
self
.
client
.
SearchVector
(
top_k_query_results
=
self
.
_client
.
SearchVector
(
table_name
=
table_name
,
query_record_array
=
query_records
,
topk
=
top_k
)
table_name
=
table_name
,
query_record_array
=
query_records
,
query_range_array
=
query_ranges
,
topk
=
top_k
)
if
top_k_query_results
:
for
top_k
in
top_k_query_results
:
res
.
append
(
TopKQueryResult
([
QueryResult
(
qr
.
id
,
qr
.
score
)
for
qr
in
top_k
.
query_result_arrays
]))
except
(
TApplicationException
,
TException
)
as
e
:
except
(
TApplicationException
,
TException
)
as
e
:
LOGGER
.
error
(
'{}'
.
format
(
e
))
LOGGER
.
error
(
'{}'
.
format
(
e
))
return
Status
(
Status
.
INVALI
D
,
str
(
e
)),
None
return
Status
(
Status
.
PERMISSION_DENIE
D
,
str
(
e
)),
None
return
Status
(
message
=
'Success!'
),
topk_query_result
s
return
Status
(
message
=
'Success!'
),
re
s
def
describe_table
(
self
,
table_name
):
def
describe_table
(
self
,
table_name
):
"""
"""
...
@@ -420,12 +226,14 @@ class MegaSearch(ConnectIntf):
...
@@ -420,12 +226,14 @@ class MegaSearch(ConnectIntf):
table_schema: TableSchema, return when operation is successful
table_schema: TableSchema, return when operation is successful
"""
"""
try
:
try
:
thrift_table_schema
=
self
.
client
.
DescribeTable
(
table_name
)
temp
=
self
.
_client
.
DescribeTable
(
table_name
)
# res = TableSchema(table_name=temp.table_name, dimension=temp.dimension,
# index_type=temp.index_type, store_raw_vector=temp.store_raw_vector)
except
(
TApplicationException
,
TException
)
as
e
:
except
(
TApplicationException
,
TException
)
as
e
:
LOGGER
.
error
(
'{}'
.
format
(
e
))
LOGGER
.
error
(
'{}'
.
format
(
e
))
return
Status
(
Status
.
INVALID
,
str
(
e
)),
None
return
Status
(
Status
.
PERMISSION_DENIED
,
str
(
e
)),
None
# TODO Table Schema
return
Status
(
message
=
'Success!'
),
temp
return
Status
(
message
=
'Success!'
),
thrift_table_schema
def
show_tables
(
self
):
def
show_tables
(
self
):
"""
"""
...
@@ -437,12 +245,36 @@ class MegaSearch(ConnectIntf):
...
@@ -437,12 +245,36 @@ class MegaSearch(ConnectIntf):
is successful
is successful
"""
"""
try
:
try
:
tables
=
self
.
client
.
ShowTables
()
res
=
self
.
_client
.
ShowTables
()
tables
=
[]
if
res
:
tables
,
_
=
res
except
(
TApplicationException
,
TException
)
as
e
:
except
(
TApplicationException
,
TException
)
as
e
:
LOGGER
.
error
(
'{}'
.
format
(
e
))
LOGGER
.
error
(
'{}'
.
format
(
e
))
return
Status
(
Status
.
INVALI
D
,
str
(
e
)),
None
return
Status
(
Status
.
PERMISSION_DENIE
D
,
str
(
e
)),
None
return
Status
(
message
=
'Success!'
),
tables
return
Status
(
message
=
'Success!'
),
tables
def
get_table_row_count
(
self
,
table_name
):
"""
Get table row count
:type table_name, str
:param table_name, target table name.
:returns:
Status: indicate if operation is successful
res: int, table row count
"""
try
:
count
,
_
=
self
.
_client
.
GetTableRowCount
(
table_name
)
except
(
TApplicationException
,
TException
)
as
e
:
LOGGER
.
error
(
'{}'
.
format
(
e
))
return
Status
(
Status
.
PERMISSION_DENIED
,
str
(
e
)),
None
return
Status
(
message
=
'Success'
),
count
def
client_version
(
self
):
def
client_version
(
self
):
"""
"""
Provide client version
Provide client version
...
@@ -457,8 +289,10 @@ class MegaSearch(ConnectIntf):
...
@@ -457,8 +289,10 @@ class MegaSearch(ConnectIntf):
:return: Server version
:return: Server version
"""
"""
# TODO How to get server version
if
not
self
.
connected
:
pass
raise
NotConnectError
(
'You have to connect first'
)
return
self
.
_client
.
Ping
(
'version'
)
def
server_status
(
self
,
cmd
=
None
):
def
server_status
(
self
,
cmd
=
None
):
"""
"""
...
@@ -466,4 +300,7 @@ class MegaSearch(ConnectIntf):
...
@@ -466,4 +300,7 @@ class MegaSearch(ConnectIntf):
:return: Server status
:return: Server status
"""
"""
return
self
.
client
.
Ping
(
cmd
)
if
not
self
.
connected
:
raise
NotConnectError
(
'You have to connect first'
)
return
self
.
_client
.
Ping
(
cmd
)
python/sdk/client/Exceptions.py
浏览文件 @
508131e4
...
@@ -2,10 +2,6 @@ class ParamError(ValueError):
...
@@ -2,10 +2,6 @@ class ParamError(ValueError):
pass
pass
class
ConnectParamMissingError
(
ParamError
):
pass
class
ConnectError
(
ValueError
):
class
ConnectError
(
ValueError
):
pass
pass
...
...
python/sdk/client/Status.py
浏览文件 @
508131e4
...
@@ -3,13 +3,15 @@ class Status(object):
...
@@ -3,13 +3,15 @@ class Status(object):
:attribute code : int (optional) default as ok
:attribute code : int (optional) default as ok
:attribute message : str (optional) current status message
:attribute message : str (optional) current status message
"""
"""
OK
=
0
SUCCESS
=
0
INVALID
=
1
CONNECT_FAILED
=
1
UNKNOWN_ERROR
=
2
PERMISSION_DENIED
=
2
NOT_SUPPORTED
=
3
TABLE_NOT_EXISTS
=
3
NOT_CONNECTED
=
4
ILLEGAL_ARGUMENT
=
4
ILLEGAL_RANGE
=
5
ILLEGAL_DIMENSION
=
6
def
__init__
(
self
,
code
=
OK
,
message
=
None
):
def
__init__
(
self
,
code
=
SUCCESS
,
message
=
None
):
self
.
code
=
code
self
.
code
=
code
self
.
message
=
message
self
.
message
=
message
...
...
python/sdk/examples/connection_exp.py
浏览文件 @
508131e4
from
client.Client
import
MegaSearch
,
Prepare
,
IndexType
,
ColumnType
from
client.Client
import
MegaSearch
,
Prepare
,
IndexType
from
client.Status
import
Status
from
client.Status
import
Status
import
time
import
time
import
random
import
struct
from
pprint
import
pprint
from
megasearch.thrift
import
MegasearchService
,
ttypes
from
megasearch.thrift
import
MegasearchService
,
ttypes
def
main
():
def
main
():
# Get client version
mega
=
MegaSearch
()
mega
=
MegaSearch
()
print
(
mega
.
client_version
(
))
print
(
'# Client version: {}'
.
format
(
mega
.
client_version
()
))
# Connect
# Connect
param
=
{
'host'
:
'192.168.1.129'
,
'port'
:
'33001'
}
param
=
{
'host'
:
'192.168.1.129'
,
'port'
:
'33001'
}
cnn_status
=
mega
.
connect
(
**
param
)
cnn_status
=
mega
.
connect
(
**
param
)
print
(
'Connect Status: {}'
.
format
(
cnn_status
))
print
(
'
#
Connect Status: {}'
.
format
(
cnn_status
))
# Check if connected
is_connected
=
mega
.
connected
is_connected
=
mega
.
connected
print
(
'Connect status: {}'
.
format
(
is_connected
))
print
(
'# Is connected: {}'
.
format
(
is_connected
))
# Create table with 1 vector column, 1 attribute column and 1 partition column
# 1. prepare table_schema
# table_schema = Prepare.table_schema(
# table_name='fake_table_name' + time.strftime('%H%M%S'),
#
# vector_columns=[Prepare.vector_column(
# name='fake_vector_name' + time.strftime('%H%M%S'),
# store_raw_vector=False,
# dimension=256)],
#
# attribute_columns=[],
#
# partition_column_names=[]
# )
# get server version
print
(
mega
.
server_status
(
'version'
))
print
(
mega
.
client
.
Ping
(
'version'
))
# show tables and their description
statu
,
tables
=
mega
.
show_tables
()
print
(
tables
)
for
table
in
tables
:
s
,
t
=
mega
.
describe_table
(
table
)
print
(
'table: {}'
.
format
(
t
))
# Create table
# Get server version
# 1. create table schema
print
(
'# Server version: {}'
.
format
(
mega
.
server_version
()))
table_schema_full
=
MegasearchService
.
TableSchema
(
table_name
=
'fake'
+
time
.
strftime
(
'%H%M%S'
),
vector_column_array
=
[
MegasearchService
.
VectorColumn
(
base
=
MegasearchService
.
Column
(
name
=
'111'
,
type
=
ttypes
.
TType
.
LIST
),
index_type
=
"aaa"
,
dimension
=
256
,
store_raw_vector
=
False
,
)],
attribute_column_array
=
[],
partition_column_name_array
=
[]
# Show tables and their description
)
status
,
tables
=
mega
.
show_tables
()
print
(
'# Show tables: {}'
.
format
(
tables
))
# 2. Create Table
# Create table
create_status
=
mega
.
client
.
CreateTable
(
param
=
table_schema_full
)
# 01.Prepare data
print
(
'Create table status: {}'
.
format
(
create_status
))
param
=
{
'table_name'
:
'test'
+
str
(
random
.
randint
(
0
,
999
)),
# add_vector
'dimension'
:
256
,
'index_type'
:
IndexType
.
IDMAP
,
'store_raw_vector'
:
False
}
# 02.Create table
res_status
=
mega
.
create_table
(
Prepare
.
table_schema
(
**
param
))
print
(
'# Create table status: {}'
.
format
(
res_status
))
# Describe table
table_name
=
'test01'
res_status
,
table
=
mega
.
describe_table
(
table_name
)
print
(
'# Describe table status: {}'
.
format
(
res_status
))
print
(
'# Describe table:{}'
.
format
(
table
))
# Add vectors to table 'test01'
# 01. Prepare data
dim
=
256
# list of binary vectors
vectors
=
[
Prepare
.
row_record
(
struct
.
pack
(
str
(
dim
)
+
'd'
,
*
[
random
.
random
()
for
_
in
range
(
dim
)]))
for
_
in
range
(
20
)]
# 02. Add vectors
status
,
ids
=
mega
.
add_vectors
(
table_name
=
table_name
,
records
=
vectors
)
print
(
'# Add vector status: {}'
.
format
(
status
))
pprint
(
ids
)
# Search vectors
q_records
=
[
Prepare
.
row_record
(
struct
.
pack
(
str
(
dim
)
+
'd'
,
*
[
random
.
random
()
for
_
in
range
(
dim
)]))
for
_
in
range
(
5
)]
param
=
{
'table_name'
:
'test01'
,
'query_records'
:
q_records
,
'top_k'
:
10
,
# 'query_ranges': None # Optional
}
sta
,
results
=
mega
.
search_vectors
(
**
param
)
print
(
'# Search vectors status: {}'
.
format
(
sta
))
pprint
(
results
)
# Get table row count
sta
,
result
=
mega
.
get_table_row_count
(
table_name
)
print
(
'# Status: {}'
.
format
(
sta
))
print
(
'# Count: {}'
.
format
(
result
))
# Delete table 'test01'
res_status
=
mega
.
delete_table
(
table_name
)
print
(
'# Delete table status: {}'
.
format
(
res_status
))
# Disconnect
# Disconnect
discnn_status
=
mega
.
disconnect
()
discnn_status
=
mega
.
disconnect
()
print
(
'
Disconnect Status
{}'
.
format
(
discnn_status
))
print
(
'
# Disconnect Status:
{}'
.
format
(
discnn_status
))
if
__name__
==
'__main__'
:
if
__name__
==
'__main__'
:
main
()
main
()
\ No newline at end of file
python/sdk/tests/TestClient.py
浏览文件 @
508131e4
...
@@ -3,17 +3,20 @@ import pytest
...
@@ -3,17 +3,20 @@ import pytest
import
mock
import
mock
import
faker
import
faker
import
random
import
random
import
struct
from
faker.providers
import
BaseProvider
from
faker.providers
import
BaseProvider
from
client.Client
import
MegaSearch
,
Prepare
,
IndexType
,
ColumnType
from
client.Client
import
MegaSearch
,
Prepare
from
client.Abstract
import
IndexType
,
TableSchema
from
client.Status
import
Status
from
client.Status
import
Status
from
client.Exceptions
import
(
from
client.Exceptions
import
(
RepeatingConnectError
,
RepeatingConnectError
,
DisconnectNotConnectedClientError
DisconnectNotConnectedClientError
)
)
from
megasearch.thrift
import
ttypes
,
MegasearchService
from
thrift.transport.TSocket
import
TSocket
from
thrift.transport.TSocket
import
TSocket
from
megasearch.thrift
import
ttypes
,
MegasearchService
from
thrift.transport
import
TTransport
from
thrift.transport.TTransport
import
TTransportException
from
thrift.transport.TTransport
import
TTransportException
LOGGER
=
logging
.
getLogger
(
__name__
)
LOGGER
=
logging
.
getLogger
(
__name__
)
...
@@ -35,63 +38,37 @@ fake = faker.Faker()
...
@@ -35,63 +38,37 @@ fake = faker.Faker()
fake
.
add_provider
(
FakerProvider
)
fake
.
add_provider
(
FakerProvider
)
def
vector_column_factory
():
return
{
'name'
:
fake
.
name
(),
'dimension'
:
fake
.
dim
(),
'store_raw_vector'
:
True
}
def
column_factory
():
return
{
'name'
:
fake
.
table_name
(),
'type'
:
ColumnType
.
INT32
}
def
range_factory
():
def
range_factory
():
return
{
param
=
{
'start'
:
str
(
random
.
randint
(
1
,
10
)),
'start'
:
str
(
random
.
randint
(
1
,
10
)),
'end'
:
str
(
random
.
randint
(
11
,
20
)),
'end'
:
str
(
random
.
randint
(
11
,
20
)),
}
}
return
Prepare
.
range
(
**
param
)
def
ranges_factory
():
return
[
range_factory
()
for
_
in
range
(
5
)]
def
table_schema_factory
():
def
table_schema_factory
():
vec_params
=
[
vector_column_factory
()
for
i
in
range
(
10
)]
column_params
=
[
column_factory
()
for
i
in
range
(
5
)]
param
=
{
param
=
{
'table_name'
:
fake
.
table_name
(),
'table_name'
:
fake
.
table_name
(),
'
vector_columns'
:
[
Prepare
.
vector_column
(
**
pa
)
for
pa
in
vec_params
]
,
'
dimension'
:
random
.
randint
(
0
,
999
)
,
'
attribute_columns'
:
[
Prepare
.
column
(
**
pa
)
for
pa
in
column_params
]
,
'
index_type'
:
IndexType
.
IDMAP
,
'
partition_column_names'
:
[
str
(
x
)
for
x
in
range
(
2
)]
'
store_raw_vector'
:
False
}
}
return
Prepare
.
table_schema
(
**
param
)
return
Prepare
.
table_schema
(
**
param
)
def
create_table_partition_param_factory
():
def
row_record_factory
(
dimension
):
param
=
{
vec
=
[
random
.
random
()
+
random
.
randint
(
0
,
9
)
for
_
in
range
(
dimension
)]
'table_name'
:
fake
.
table_name
(),
bin_vec
=
struct
.
pack
(
str
(
dimension
)
+
"d"
,
*
vec
)
'partition_name'
:
fake
.
table_name
(),
'column_name_to_range'
:
{
fake
.
name
():
range_factory
()
for
_
in
range
(
3
)}
}
return
Prepare
.
create_table_partition_param
(
**
param
)
return
Prepare
.
row_record
(
vector_data
=
bin_vec
)
def
delete_table_partition_param_factory
():
param
=
{
'table_name'
:
fake
.
table_name
(),
'partition_names'
:
[
fake
.
name
()
for
i
in
range
(
5
)]
}
return
Prepare
.
delete_table_partition_param
(
**
param
)
def
row_record_factory
():
def
row_records_factory
(
dimension
):
param
=
{
return
[
row_record_factory
(
dimension
)
for
_
in
range
(
20
)]
'column_name_to_vector'
:
{
fake
.
name
():
[
random
.
random
()
for
i
in
range
(
256
)]},
'column_name_to_attribute'
:
{
fake
.
name
():
fake
.
name
()}
}
return
Prepare
.
row_record
(
**
param
)
class
TestConnection
:
class
TestConnection
:
...
@@ -103,9 +80,8 @@ class TestConnection:
...
@@ -103,9 +80,8 @@ class TestConnection:
cnn
=
MegaSearch
()
cnn
=
MegaSearch
()
cnn
.
connect
(
**
self
.
param
)
cnn
.
connect
(
**
self
.
param
)
assert
cnn
.
status
==
Status
.
OK
assert
cnn
.
status
==
Status
.
SUCCESS
assert
cnn
.
connected
assert
cnn
.
connected
assert
isinstance
(
cnn
.
client
,
MegasearchService
.
Client
)
with
pytest
.
raises
(
RepeatingConnectError
):
with
pytest
.
raises
(
RepeatingConnectError
):
cnn
.
connect
(
**
self
.
param
)
cnn
.
connect
(
**
self
.
param
)
...
@@ -114,12 +90,23 @@ class TestConnection:
...
@@ -114,12 +90,23 @@ class TestConnection:
def
test_false_connect
(
self
):
def
test_false_connect
(
self
):
cnn
=
MegaSearch
()
cnn
=
MegaSearch
()
cnn
.
connect
(
self
.
param
)
cnn
.
connect
(
**
self
.
param
)
assert
cnn
.
status
!=
Status
.
OK
assert
cnn
.
status
!=
Status
.
SUCCESS
@
mock
.
patch
.
object
(
TTransport
.
TBufferedTransport
,
'close'
)
@
mock
.
patch
.
object
(
TSocket
,
'open'
)
def
test_disconnected
(
self
,
close
,
open
):
close
.
return_value
=
None
open
.
return_value
=
None
cnn
=
MegaSearch
()
cnn
.
connect
(
**
self
.
param
)
assert
cnn
.
disconnect
()
==
Status
.
SUCCESS
def
test_disconnected_error
(
self
):
def
test_disconnected_error
(
self
):
cnn
=
MegaSearch
()
cnn
=
MegaSearch
()
cnn
.
connect_status
=
Status
(
Status
.
INVALI
D
)
cnn
.
connect_status
=
Status
(
Status
.
PERMISSION_DENIE
D
)
with
pytest
.
raises
(
DisconnectNotConnectedClientError
):
with
pytest
.
raises
(
DisconnectNotConnectedClientError
):
cnn
.
disconnect
()
cnn
.
disconnect
()
...
@@ -142,26 +129,26 @@ class TestTable:
...
@@ -142,26 +129,26 @@ class TestTable:
param
=
table_schema_factory
()
param
=
table_schema_factory
()
res
=
client
.
create_table
(
param
)
res
=
client
.
create_table
(
param
)
assert
res
==
Status
.
OK
assert
res
==
Status
.
SUCCESS
def
test_false_create_table
(
self
,
client
):
def
test_false_create_table
(
self
,
client
):
param
=
table_schema_factory
()
param
=
table_schema_factory
()
with
pytest
.
raises
(
TTransportException
):
with
pytest
.
raises
(
TTransportException
):
res
=
client
.
create_table
(
param
)
res
=
client
.
create_table
(
param
)
LOGGER
.
error
(
'{}'
.
format
(
res
))
LOGGER
.
error
(
'{}'
.
format
(
res
))
assert
res
!=
Status
.
OK
assert
res
!=
Status
.
SUCCESS
@
mock
.
patch
.
object
(
MegasearchService
.
Client
,
'DeleteTable'
)
@
mock
.
patch
.
object
(
MegasearchService
.
Client
,
'DeleteTable'
)
def
test_delete_table
(
self
,
DeleteTable
,
client
):
def
test_delete_table
(
self
,
DeleteTable
,
client
):
DeleteTable
.
return_value
=
None
DeleteTable
.
return_value
=
None
table_name
=
'fake_table_name'
table_name
=
'fake_table_name'
res
=
client
.
delete_table
(
table_name
)
res
=
client
.
delete_table
(
table_name
)
assert
res
==
Status
.
OK
assert
res
==
Status
.
SUCCESS
def
test_false_delete_table
(
self
,
client
):
def
test_false_delete_table
(
self
,
client
):
table_name
=
'fake_table_name'
table_name
=
'fake_table_name'
res
=
client
.
delete_table
(
table_name
)
res
=
client
.
delete_table
(
table_name
)
assert
res
!=
Status
.
OK
assert
res
!=
Status
.
SUCCESS
class
TestVector
:
class
TestVector
:
...
@@ -176,70 +163,46 @@ class TestVector:
...
@@ -176,70 +163,46 @@ class TestVector:
cnn
.
connect
(
**
param
)
cnn
.
connect
(
**
param
)
return
cnn
return
cnn
@
mock
.
patch
.
object
(
MegasearchService
.
Client
,
'CreateTablePartition'
)
def
test_create_table_partition
(
self
,
CreateTablePartition
,
client
):
CreateTablePartition
.
return_value
=
None
param
=
create_table_partition_param_factory
()
res
=
client
.
create_table_partition
(
param
)
assert
res
==
Status
.
OK
def
test_false_table_partition
(
self
,
client
):
param
=
create_table_partition_param_factory
()
res
=
client
.
create_table_partition
(
param
)
assert
res
!=
Status
.
OK
@
mock
.
patch
.
object
(
MegasearchService
.
Client
,
'DeleteTablePartition'
)
def
test_delete_table_partition
(
self
,
DeleteTablePartition
,
client
):
DeleteTablePartition
.
return_value
=
None
param
=
delete_table_partition_param_factory
()
res
=
client
.
delete_table_partition
(
param
)
assert
res
==
Status
.
OK
def
test_false_delete_table_partition
(
self
,
client
):
param
=
delete_table_partition_param_factory
()
res
=
client
.
delete_table_partition
(
param
)
assert
res
!=
Status
.
OK
@
mock
.
patch
.
object
(
MegasearchService
.
Client
,
'AddVector'
)
@
mock
.
patch
.
object
(
MegasearchService
.
Client
,
'AddVector'
)
def
test_add_vector
(
self
,
AddVector
,
client
):
def
test_add_vector
(
self
,
AddVector
,
client
):
AddVector
.
return_value
=
None
AddVector
.
return_value
=
None
param
=
{
param
=
{
'table_name'
:
fake
.
table_name
(),
'table_name'
:
fake
.
table_name
(),
'records'
:
[
row_record_factory
()
for
_
in
range
(
1000
)]
'records'
:
row_records_factory
(
256
)
}
}
res
,
ids
=
client
.
add_vector
(
**
param
)
res
,
ids
=
client
.
add_vector
s
(
**
param
)
assert
res
==
Status
.
OK
assert
res
==
Status
.
SUCCESS
def
test_false_add_vector
(
self
,
client
):
def
test_false_add_vector
(
self
,
client
):
param
=
{
param
=
{
'table_name'
:
fake
.
table_name
(),
'table_name'
:
fake
.
table_name
(),
'records'
:
[
row_record_factory
()
for
_
in
range
(
1000
)]
'records'
:
row_records_factory
(
256
)
}
}
res
,
ids
=
client
.
add_vector
(
**
param
)
res
,
ids
=
client
.
add_vector
s
(
**
param
)
assert
res
!=
Status
.
OK
assert
res
!=
Status
.
SUCCESS
@
mock
.
patch
.
object
(
MegasearchService
.
Client
,
'SearchVector'
)
@
mock
.
patch
.
object
(
MegasearchService
.
Client
,
'SearchVector'
)
def
test_search_vector
(
self
,
SearchVector
,
client
):
def
test_search_vector
(
self
,
SearchVector
,
client
):
SearchVector
.
return_value
=
None
SearchVector
.
return_value
=
None
,
None
param
=
{
param
=
{
'table_name'
:
fake
.
table_name
(),
'table_name'
:
fake
.
table_name
(),
'query_records'
:
[
row_record_factory
()
for
_
in
range
(
1000
)],
'query_records'
:
row_records_factory
(
256
),
'top_k'
:
random
.
randint
(
0
,
10
)
'query_ranges'
:
ranges_factory
(),
'top_k'
:
random
.
randint
(
0
,
10
)
}
}
res
,
results
=
client
.
search_vector
(
**
param
)
res
,
results
=
client
.
search_vector
s
(
**
param
)
assert
res
==
Status
.
OK
assert
res
==
Status
.
SUCCESS
def
test_false_vector
(
self
,
client
):
def
test_false_vector
(
self
,
client
):
param
=
{
param
=
{
'table_name'
:
fake
.
table_name
(),
'table_name'
:
fake
.
table_name
(),
'query_records'
:
[
row_record_factory
()
for
_
in
range
(
1000
)],
'query_records'
:
row_records_factory
(
256
),
'top_k'
:
random
.
randint
(
0
,
10
)
'query_ranges'
:
ranges_factory
(),
'top_k'
:
random
.
randint
(
0
,
10
)
}
}
res
,
results
=
client
.
search_vector
(
**
param
)
res
,
results
=
client
.
search_vector
s
(
**
param
)
assert
res
!=
Status
.
OK
assert
res
!=
Status
.
SUCCESS
@
mock
.
patch
.
object
(
MegasearchService
.
Client
,
'DescribeTable'
)
@
mock
.
patch
.
object
(
MegasearchService
.
Client
,
'DescribeTable'
)
def
test_describe_table
(
self
,
DescribeTable
,
client
):
def
test_describe_table
(
self
,
DescribeTable
,
client
):
...
@@ -247,27 +210,38 @@ class TestVector:
...
@@ -247,27 +210,38 @@ class TestVector:
table_name
=
fake
.
table_name
()
table_name
=
fake
.
table_name
()
res
,
table_schema
=
client
.
describe_table
(
table_name
)
res
,
table_schema
=
client
.
describe_table
(
table_name
)
assert
res
==
Status
.
OK
assert
res
==
Status
.
SUCCESS
assert
isinstance
(
table_schema
,
ttypes
.
TableSchema
)
assert
isinstance
(
table_schema
,
TableSchema
)
def
test_false_decribe_table
(
self
,
client
):
def
test_false_decribe_table
(
self
,
client
):
table_name
=
fake
.
table_name
()
table_name
=
fake
.
table_name
()
res
,
table_schema
=
client
.
describe_table
(
table_name
)
res
,
table_schema
=
client
.
describe_table
(
table_name
)
assert
res
!=
Status
.
OK
assert
res
!=
Status
.
SUCCESS
assert
not
table_schema
assert
not
table_schema
@
mock
.
patch
.
object
(
MegasearchService
.
Client
,
'ShowTables'
)
@
mock
.
patch
.
object
(
MegasearchService
.
Client
,
'ShowTables'
)
def
test_show_tables
(
self
,
ShowTables
,
client
):
def
test_show_tables
(
self
,
ShowTables
,
client
):
ShowTables
.
return_value
=
[
fake
.
table_name
()
for
_
in
range
(
10
)]
ShowTables
.
return_value
=
[
fake
.
table_name
()
for
_
in
range
(
10
)]
,
None
res
,
tables
=
client
.
show_tables
()
res
,
tables
=
client
.
show_tables
()
assert
res
==
Status
.
OK
assert
res
==
Status
.
SUCCESS
assert
isinstance
(
tables
,
list
)
assert
isinstance
(
tables
,
list
)
def
test_false_show_tables
(
self
,
client
):
def
test_false_show_tables
(
self
,
client
):
res
,
tables
=
client
.
show_tables
()
res
,
tables
=
client
.
show_tables
()
assert
res
!=
Status
.
OK
assert
res
!=
Status
.
SUCCESS
assert
not
tables
assert
not
tables
@
mock
.
patch
.
object
(
MegasearchService
.
Client
,
'GetTableRowCount'
)
def
test_get_table_row_count
(
self
,
GetTableRowCount
,
client
):
GetTableRowCount
.
return_value
=
22
,
None
res
,
count
=
client
.
get_table_row_count
(
'fake_table'
)
assert
res
==
Status
.
SUCCESS
def
test_false_get_table_row_count
(
self
,
client
):
res
,
count
=
client
.
get_table_row_count
(
'fake_table'
)
assert
res
!=
Status
.
SUCCESS
assert
not
count
def
test_client_version
(
self
,
client
):
def
test_client_version
(
self
,
client
):
res
=
client
.
client_version
()
res
=
client
.
client_version
()
assert
res
==
'0.0.1'
assert
res
==
'0.0.1'
...
@@ -275,34 +249,13 @@ class TestVector:
...
@@ -275,34 +249,13 @@ class TestVector:
class
TestPrepare
:
class
TestPrepare
:
def
test_column
(
self
):
param
=
{
'name'
:
'test01'
,
'type'
:
ColumnType
.
DATE
}
res
=
Prepare
.
column
(
**
param
)
LOGGER
.
error
(
'{}'
.
format
(
res
))
assert
res
.
name
==
'test01'
assert
res
.
type
==
ColumnType
.
DATE
assert
isinstance
(
res
,
ttypes
.
Column
)
def
test_vector_column
(
self
):
param
=
vector_column_factory
()
res
=
Prepare
.
vector_column
(
**
param
)
LOGGER
.
error
(
'{}'
.
format
(
res
))
assert
isinstance
(
res
,
ttypes
.
VectorColumn
)
def
test_table_schema
(
self
):
def
test_table_schema
(
self
):
vec_params
=
[
vector_column_factory
()
for
i
in
range
(
10
)]
column_params
=
[
column_factory
()
for
i
in
range
(
5
)]
param
=
{
param
=
{
'table_name'
:
'test03'
,
'table_name'
:
fake
.
table_name
()
,
'
vector_columns'
:
[
Prepare
.
vector_column
(
**
pa
)
for
pa
in
vec_params
]
,
'
dimension'
:
random
.
randint
(
0
,
999
)
,
'
attribute_columns'
:
[
Prepare
.
column
(
**
pa
)
for
pa
in
column_params
]
,
'
index_type'
:
IndexType
.
IDMAP
,
'
partition_column_names'
:
[
str
(
x
)
for
x
in
range
(
2
)]
'
store_raw_vector'
:
False
}
}
res
=
Prepare
.
table_schema
(
**
param
)
res
=
Prepare
.
table_schema
(
**
param
)
assert
isinstance
(
res
,
ttypes
.
TableSchema
)
assert
isinstance
(
res
,
ttypes
.
TableSchema
)
...
@@ -319,39 +272,10 @@ class TestPrepare:
...
@@ -319,39 +272,10 @@ class TestPrepare:
assert
res
.
start_value
==
'200'
assert
res
.
start_value
==
'200'
assert
res
.
end_value
==
'1000'
assert
res
.
end_value
==
'1000'
def
test_create_table_partition_param
(
self
):
param
=
{
'table_name'
:
fake
.
table_name
(),
'partition_name'
:
fake
.
table_name
(),
'column_name_to_range'
:
{
fake
.
name
():
range_factory
()
for
_
in
range
(
3
)}
}
res
=
Prepare
.
create_table_partition_param
(
**
param
)
LOGGER
.
error
(
'{}'
.
format
(
res
))
assert
isinstance
(
res
,
ttypes
.
CreateTablePartitionParam
)
def
test_delete_table_partition_param
(
self
):
param
=
{
'table_name'
:
fake
.
table_name
(),
'partition_names'
:
[
fake
.
name
()
for
i
in
range
(
5
)]
}
res
=
Prepare
.
delete_table_partition_param
(
**
param
)
assert
isinstance
(
res
,
ttypes
.
DeleteTablePartitionParam
)
def
test_row_record
(
self
):
def
test_row_record
(
self
):
param
=
{
vec
=
[
random
.
random
()
+
random
.
randint
(
0
,
9
)
for
_
in
range
(
256
)]
'column_name_to_vector'
:
{
fake
.
name
():
[
random
.
random
()
for
i
in
range
(
256
)]},
bin_vec
=
struct
.
pack
(
str
(
256
)
+
"d"
,
*
vec
)
'column_name_to_attribute'
:
{
fake
.
name
():
fake
.
name
()}
res
=
Prepare
.
row_record
(
bin_vec
)
}
res
=
Prepare
.
row_record
(
**
param
)
assert
isinstance
(
res
,
ttypes
.
RowRecord
)
assert
isinstance
(
res
,
ttypes
.
RowRecord
)
assert
isinstance
(
bin_vec
,
bytes
)
def
test_query_record
(
self
):
param
=
{
'column_name_to_vector'
:
{
fake
.
name
():
[
random
.
random
()
for
i
in
range
(
256
)]},
'selected_columns'
:
[
fake
.
name
()
for
_
in
range
(
10
)],
'name_to_partition_ranges'
:
{
fake
.
name
():
[
range_factory
()
for
_
in
range
(
5
)]}
}
res
=
Prepare
.
query_record
(
**
param
)
assert
isinstance
(
res
,
ttypes
.
QueryRecord
)
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录