milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add batch import

pull/3/head
jinhai 2019-04-02 19:33:22 +08:00
parent bd10891c15
commit a7636a2636
5 changed files with 62 additions and 59 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
pyengine/engine/controller/tests/test_vector_engine.py
View File

@ -11,6 +11,7 @@ logger = logging.getLogger(__name__)
class TestVectorEngine:
def setup_class(self):
self.__vectors = [[1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8]]
self.__vector = [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8]
self.__limit = 1
@ -44,29 +45,29 @@ class TestVectorEngine:
assert group_list == [{'group_name': 'test_group', 'file_number': 0}]
# Add Vector for not exist group
code, vector_id = VectorEngine.AddVector('not_exist_group', self.__vector)
code, vector_id = VectorEngine.AddVector('not_exist_group', self.__vectors)
assert code == VectorEngine.GROUP_NOT_EXIST
assert vector_id == 'invalid'
# Add vector for exist group
code, vector_id = VectorEngine.AddVector('test_group', self.__vector)
code, vector_id = VectorEngine.AddVector('test_group', self.__vectors)
assert code == VectorEngine.SUCCESS_CODE
assert vector_id == 'test_group.0'
assert vector_id == ['test_group.0']
# Add vector for exist group
code, vector_id = VectorEngine.AddVector('test_group', self.__vector)
code, vector_id = VectorEngine.AddVector('test_group', self.__vectors)
assert code == VectorEngine.SUCCESS_CODE
assert vector_id == 'test_group.1'
assert vector_id == ['test_group.1']
# Add vector for exist group
code, vector_id = VectorEngine.AddVector('test_group', self.__vector)
code, vector_id = VectorEngine.AddVector('test_group', self.__vectors)
assert code == VectorEngine.SUCCESS_CODE
assert vector_id == 'test_group.2'
assert vector_id == ['test_group.2']
# Add vector for exist group
code, vector_id = VectorEngine.AddVector('test_group', self.__vector)
code, vector_id = VectorEngine.AddVector('test_group', self.__vectors)
assert code == VectorEngine.SUCCESS_CODE
assert vector_id == 'test_group.3'
assert vector_id == ['test_group.3']
# Check search vector interface
code, vector_id = VectorEngine.SearchVector('test_group', self.__vector, self.__limit)

6
pyengine/engine/controller/tests/test_views.py
View File

@ -48,17 +48,17 @@ class TestViews:
assert resp.status_code == 200
assert self.loads(resp)['code'] == 0
vector = {"vector": [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8]}
vector = {"vector": [[1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8]]}
resp = test_client.post('/vector/add/6', data=json.dumps(vector), headers = TestViews.HEADERS)
assert resp.status_code == 200
assert self.loads(resp)['code'] == 0
vector = {"vector": [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8]}
vector = {"vector": [[1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8]]}
resp = test_client.post('/vector/add/6', data=json.dumps(vector), headers = TestViews.HEADERS)
assert resp.status_code == 200
assert self.loads(resp)['code'] == 0
vector = {"vector": [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8]}
vector = {"vector": [[1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8]]}
resp = test_client.post('/vector/add/6', data=json.dumps(vector), headers = TestViews.HEADERS)
assert resp.status_code == 200
assert self.loads(resp)['code'] == 0

91
pyengine/engine/controller/vector_engine.py
View File

@ -80,62 +80,65 @@ class VectorEngine(object):
@staticmethod
def AddVector(group_id, vector):
print(group_id, vector)
def AddVector(group_id, vectors):
print(group_id, vectors)
code, _, _ = VectorEngine.GetGroup(group_id)
if code == VectorEngine.FAULT_CODE:
return VectorEngine.GROUP_NOT_EXIST, 'invalid'
file = FileTable.query.filter(FileTable.group_name == group_id).filter(FileTable.type == 'raw').first()
group = GroupTable.query.filter(GroupTable.group_name == group_id).first()
if file:
print('insert into exist file')
# create vector id
vector_id = file.seq_no + 1
# insert into raw file
VectorEngine.InsertVectorIntoRawFile(group_id, file.filename, vector, vector_id)
vector_str_list = []
for vector in vectors:
file = FileTable.query.filter(FileTable.group_name == group_id).filter(FileTable.type == 'raw').first()
group = GroupTable.query.filter(GroupTable.group_name == group_id).first()
# check if the file can be indexed
if file.row_number + 1 >= ROW_LIMIT:
raw_vector_array, raw_vector_id_array = VectorEngine.GetVectorListFromRawFile(group_id)
d = group.dimension
if file:
print('insert into exist file')
# create vector id
vector_id = file.seq_no + 1
# insert into raw file
VectorEngine.InsertVectorIntoRawFile(group_id, file.filename, vector, vector_id)
# create index
index_builder = build_index.FactoryIndex()
index = index_builder().build(d, raw_vector_array, raw_vector_id_array)
# check if the file can be indexed
if file.row_number + 1 >= ROW_LIMIT:
raw_vector_array, raw_vector_id_array = VectorEngine.GetVectorListFromRawFile(group_id)
d = group.dimension
# TODO(jinhai): store index into Cache
index_filename = file.filename + '_index'
serialize.write_index(file_name=index_filename, index=index)
# create index
index_builder = build_index.FactoryIndex()
index = index_builder().build(d, raw_vector_array, raw_vector_id_array)
# TODO(jinhai): store index into Cache
index_filename = file.filename + '_index'
serialize.write_index(file_name=index_filename, index=index)
FileTable.query.filter(FileTable.group_name == group_id).filter(FileTable.type == 'raw').update({'row_number':file.row_number + 1,
'type': 'index',
'filename': index_filename,
'seq_no': file.seq_no + 1})
db.session.commit()
VectorEngine.group_dict = None
else:
# we still can insert into exist raw file, update database
FileTable.query.filter(FileTable.group_name == group_id).filter(FileTable.type == 'raw').update({'row_number':file.row_number + 1,
'seq_no': file.seq_no + 1})
db.session.commit()
print('Update db for raw file insertion')
FileTable.query.filter(FileTable.group_name == group_id).filter(FileTable.type == 'raw').update({'row_number':file.row_number + 1,
'type': 'index',
'filename': index_filename,
'seq_no': file.seq_no + 1})
db.session.commit()
VectorEngine.group_dict = None
else:
# we still can insert into exist raw file, update database
FileTable.query.filter(FileTable.group_name == group_id).filter(FileTable.type == 'raw').update({'row_number':file.row_number + 1,
'seq_no': file.seq_no + 1})
print('add a new raw file')
# first raw file
raw_filename = group_id + '.raw'
# create vector id
vector_id = 0
# create and insert vector into raw file
VectorEngine.InsertVectorIntoRawFile(group_id, raw_filename, vector, vector_id)
# insert a record into database
db.session.add(FileTable(group_id, raw_filename, 'raw', 1))
db.session.commit()
print('Update db for raw file insertion')
pass
else:
print('add a new raw file')
# first raw file
raw_filename = group_id + '.raw'
# create vector id
vector_id = 0
# create and insert vector into raw file
VectorEngine.InsertVectorIntoRawFile(group_id, raw_filename, vector, vector_id)
# insert a record into database
db.session.add(FileTable(group_id, raw_filename, 'raw', 1))
db.session.commit()
vector_str_list.append(group_id + '.' + str(vector_id))
vector_id_str = group_id + '.' + str(vector_id)
return VectorEngine.SUCCESS_CODE, vector_id_str
return VectorEngine.SUCCESS_CODE, vector_str_list
@staticmethod

3
pyengine/engine/controller/views.py
View File

@ -14,10 +14,9 @@ from flask_restful import request
class Vector(Resource):
def __init__(self):
self.__parser = reqparse.RequestParser()
self.__parser.add_argument('vector', type=float, action='append', location=['json'])
self.__parser.add_argument('vector', type=list, action='append', location=['json'])
def post(self, group_id):
print(request.json)
args = self.__parser.parse_args()
vector = args['vector']
code, vector_id = VectorEngine.AddVector(group_id, vector)

2
pyengine/engine/run_test.sh
View File

@ -1 +1 @@
pytest -v --disable-warnings
pytest -vv --disable-warnings