Пример #1
0
def fit(table_name, X):
    milvus = Milvus()
    milvus.connect(host = SERVER_HOST_DEFAULT, port = SERVER_PORT_DEFAULT) 
    start = time.time()
    status, ids = milvus.add_vectors(table_name, X)
    end = time.time()
    logger(status, round(end - start, 2))
Пример #2
0
    def _add():
        milvus = Milvus()
        status = milvus.connect(**server_config)

        vectors = _generate_vectors(128, 10000)
        print('\n\tPID: {}, insert {} vectors'.format(os.getpid(), 10000))
        status, _ = milvus.add_vectors(_table_name, vectors)

        milvus.disconnect()
Пример #3
0
    def _add():
        milvus = Milvus()
        status = milvus.connect()

        if not status.OK:
            print(f'PID: {os.getpid()}, connect failed')

        status, _ = milvus.add_vectors(_table_name, vectors)

        milvus.disconnect()
Пример #4
0
    def test_search_during_createIndex(self, args):
        loops = 100000
        table = gen_unique_str()
        query_vecs = [vectors[0], vectors[1]]
        uri = "tcp://%s:%s" % (args["ip"], args["port"])
        id_0 = 0
        id_1 = 0
        milvus_instance = Milvus()
        milvus_instance.connect(uri=uri)
        milvus_instance.create_table({
            'table_name': table,
            'dimension': dim,
            'index_file_size': index_file_size,
            'metric_type': MetricType.L2
        })
        for i in range(10):
            status, ids = milvus_instance.add_vectors(table, vectors)
            # logging.getLogger().info(ids)
            if i == 0:
                id_0 = ids[0]
                id_1 = ids[1]

        def create_index(milvus_instance):
            logging.getLogger().info("In create index")
            status = milvus_instance.create_index(table, index_params)
            logging.getLogger().info(status)
            status, result = milvus_instance.describe_index(table)
            logging.getLogger().info(result)

        def add_vectors(milvus_instance):
            logging.getLogger().info("In add vectors")
            status, ids = milvus_instance.add_vectors(table, vectors)
            logging.getLogger().info(status)

        def search(milvus_instance):
            logging.getLogger().info("In search vectors")
            for i in range(loops):
                status, result = milvus_instance.search_vectors(
                    table, top_k, nprobe, query_vecs)
                logging.getLogger().info(status)
                assert result[0][0].id == id_0
                assert result[1][0].id == id_1

        milvus_instance = Milvus()
        milvus_instance.connect(uri=uri)
        p_search = Process(target=search, args=(milvus_instance, ))
        p_search.start()
        milvus_instance = Milvus()
        milvus_instance.connect(uri=uri)
        p_create = Process(target=add_vectors, args=(milvus_instance, ))
        p_create.start()
        p_create.join()
Пример #5
0
 def test_search_multi_table_IP(search, args):
     '''
     target: test search multi tables of IP
     method: add vectors into 10 tables, and search
     expected: search status ok, the length of result
     '''
     num = 10
     top_k = 10
     nprobe = 1
     tables = []
     idx = []
     for i in range(num):
         table = gen_unique_str("test_add_multitable_%d" % i)
         uri = "tcp://%s:%s" % (args["ip"], args["port"])
         param = {
             'table_name': table,
             'dimension': dim,
             'index_file_size': 10,
             'metric_type': MetricType.L2
         }
         # create table
         milvus = Milvus()
         milvus.connect(uri=uri)
         milvus.create_table(param)
         status, ids = milvus.add_vectors(table, vectors)
         assert status.OK()
         assert len(ids) == len(vectors)
         tables.append(table)
         idx.append(ids[0])
         idx.append(ids[10])
         idx.append(ids[20])
     time.sleep(6)
     query_vecs = [vectors[0], vectors[10], vectors[20]]
     # start query from random table
     for i in range(num):
         table = tables[i]
         status, result = milvus.search_vectors(table, top_k, nprobe,
                                                query_vecs)
         assert status.OK()
         assert len(result) == len(query_vecs)
         for j in range(len(query_vecs)):
             assert len(result[j]) == top_k
         for j in range(len(query_vecs)):
             assert check_result(result[j], idx[3 * i + j])
Пример #6
0
 connect.delete_table(table)
 # sys.exit()
 # time.sleep(2)
 print(connect.get_table_row_count(table))
 param = {
     'table_name': table,
     'dimension': dim,
     'metric_type': MetricType.L2,
     'index_file_size': 10
 }
 status = connect.create_table(param)
 print(status)
 print(connect.get_table_row_count(table))
 # add vectors
 for i in range(10):
     status, ids = connect.add_vectors(table, vectors)
     print(status)
     print(ids[0])
 # print(ids[0])
 index_params = {"index_type": IndexType.IVFLAT, "nlist": 16384}
 status = connect.create_index(table, index_params)
 print(status)
 # sys.exit()
 query_vec = [vectors[0]]
 # print(numpy.inner(numpy.array(query_vec[0]), numpy.array(query_vec[0])))
 top_k = 12
 nprobe = 1
 for i in range(2):
     result = connect.search_vectors(table, top_k, nprobe, query_vec)
     print(result)
 sys.exit()
Пример #7
0
from milvus import Milvus, Prepare, IndexType, Status
import random

milvus = Milvus()

# Connect Milvus server, please change HOST and PORT to correct one
milvus.connect(host='localhost', port='33001')

# Table name is defined
table_name = 'table_' + str(random.randint(0, 100))

# Create table: table name, vector dimension and index type
milvus.create_table(
    Prepare.table_schema(table_name, dimension=256,
                         index_type=IndexType.IDMAP))

# Add 20 256-dim-vectors into table
vectors = Prepare.records([[random.random() for _ in range(256)]
                           for _ in range(20)])
milvus.add_vectors(table_name=table_name, records=vectors)

# Get table row count
_, result = milvus.get_table_row_count(table_name=table_name)
print('Table {}, row counts: {}'.format(table_name, result))
Пример #8
0
class MilvusClient(object):
    def __init__(self, table_name=None, ip=None, port=None):
        self._milvus = Milvus()
        self._table_name = table_name
        try:
            if not ip:
                self._milvus.connect(
                    host = SERVER_HOST_DEFAULT,
                    port = SERVER_PORT_DEFAULT)
            else:
                self._milvus.connect(
                    host = ip,
                    port = port)
        except Exception as e:
            raise e

    def __str__(self):
        return 'Milvus table %s' % self._table_name

    def check_status(self, status):
        if not status.OK():
            logger.error(status.message)
            raise Exception("Status not ok")

    def create_table(self, table_name, dimension, index_file_size, metric_type):
        if not self._table_name:
            self._table_name = table_name
        if metric_type == "l2":
            metric_type = MetricType.L2
        elif metric_type == "ip":
            metric_type = MetricType.IP
        else:
            logger.error("Not supported metric_type: %s" % metric_type)
        create_param = {'table_name': table_name,
                 'dimension': dimension,
                 'index_file_size': index_file_size, 
                 "metric_type": metric_type}
        status = self._milvus.create_table(create_param)
        self.check_status(status)

    @time_wrapper
    def insert(self, X, ids=None):
        status, result = self._milvus.add_vectors(self._table_name, X, ids)
        self.check_status(status)
        return status, result

    @time_wrapper
    def create_index(self, index_type, nlist):
        if index_type == "flat":
            index_type = IndexType.FLAT
        elif index_type == "ivf_flat":
            index_type = IndexType.IVFLAT
        elif index_type == "ivf_sq8":
            index_type = IndexType.IVF_SQ8
        elif index_type == "nsg":
            index_type = IndexType.NSG
        elif index_type == "ivf_sq8h":
            index_type = IndexType.IVF_SQ8H
        elif index_type == "ivf_pq":
            index_type = IndexType.IVF_PQ
        index_params = {
            "index_type": index_type,
            "nlist": nlist,
        }
        logger.info("Building index start, table_name: %s, index_params: %s" % (self._table_name, json.dumps(index_params)))
        status = self._milvus.create_index(self._table_name, index=index_params)
        self.check_status(status)

    def describe_index(self):
        return self._milvus.describe_index(self._table_name)

    def drop_index(self):
        logger.info("Drop index: %s" % self._table_name)
        return self._milvus.drop_index(self._table_name)

    @time_wrapper
    def query(self, X, top_k, nprobe):
        status, result = self._milvus.search_vectors(self._table_name, top_k, nprobe, X)
        self.check_status(status)
        return status, result

    def count(self):
        return self._milvus.get_table_row_count(self._table_name)[1]

    def delete(self, timeout=60):
        logger.info("Start delete table: %s" % self._table_name)
        self._milvus.delete_table(self._table_name)
        i = 0
        while i < timeout:
            if self.count():
                time.sleep(1)
                i = i + 1
                continue
            else:
                break
        if i < timeout:
            logger.error("Delete table timeout")

    def describe(self):
        return self._milvus.describe_table(self._table_name)

    def show_tables(self):
        return self._milvus.show_tables()

    def exists_table(self):
        status, res = self._milvus.has_table(self._table_name)
        self.check_status(status)
        return res

    @time_wrapper
    def preload_table(self):
        return self._milvus.preload_table(self._table_name, timeout=3000)
Пример #9
0
class MilvusClient(object):
    def __init__(self, collection_name=None, ip=None, port=None, timeout=60):
        self._collection_name = collection_name
        try:
            i = 1
            start_time = time.time()
            if not ip:
                self._milvus = Milvus(host=SERVER_HOST_DEFAULT,
                                      port=SERVER_PORT_DEFAULT)
            else:
                # retry connect for remote server
                while time.time() < start_time + timeout:
                    try:
                        self._milvus = Milvus(host=ip, port=port)
                        if self._milvus.server_status():
                            logger.debug(
                                "Try connect times: %d, %s" %
                                (i, round(time.time() - start_time, 2)))
                            break
                    except Exception as e:
                        logger.debug("Milvus connect failed")
                        i = i + 1

        except Exception as e:
            raise e

    def __str__(self):
        return 'Milvus collection %s' % self._collection_name

    def check_status(self, status):
        if not status.OK():
            logger.error(status.message)
            raise Exception("Status not ok")

    def check_result_ids(self, result):
        for index, item in enumerate(result):
            if item[0].distance >= epsilon:
                logger.error(index)
                logger.error(item[0].distance)
                raise Exception("Distance wrong")

    def create_collection(self, collection_name, dimension, index_file_size,
                          metric_type):
        if not self._collection_name:
            self._collection_name = collection_name
        if metric_type == "l2":
            metric_type = MetricType.L2
        elif metric_type == "ip":
            metric_type = MetricType.IP
        elif metric_type == "jaccard":
            metric_type = MetricType.JACCARD
        elif metric_type == "hamming":
            metric_type = MetricType.HAMMING
        elif metric_type == "sub":
            metric_type = MetricType.SUBSTRUCTURE
        elif metric_type == "super":
            metric_type = MetricType.SUPERSTRUCTURE
        else:
            logger.error("Not supported metric_type: %s" % metric_type)
        create_param = {
            'collection_name': collection_name,
            'dimension': dimension,
            'index_file_size': index_file_size,
            "metric_type": metric_type
        }
        status = self._milvus.create_collection(create_param)
        self.check_status(status)

    @time_wrapper
    def insert(self, X, ids=None):
        status, result = self._milvus.add_vectors(self._collection_name, X,
                                                  ids)
        self.check_status(status)
        return status, result

    @time_wrapper
    def delete_vectors(self, ids):
        status = self._milvus.delete_by_id(self._collection_name, ids)
        self.check_status(status)

    @time_wrapper
    def flush(self):
        status = self._milvus.flush([self._collection_name])
        self.check_status(status)

    @time_wrapper
    def compact(self):
        status = self._milvus.compact(self._collection_name)
        self.check_status(status)

    @time_wrapper
    def create_index(self, index_type, index_param=None):
        index_type = INDEX_MAP[index_type]
        logger.info(
            "Building index start, collection_name: %s, index_type: %s" %
            (self._collection_name, index_type))
        if index_param:
            logger.info(index_param)
        status = self._milvus.create_index(self._collection_name, index_type,
                                           index_param)
        self.check_status(status)

    def describe_index(self):
        status, result = self._milvus.describe_index(self._collection_name)
        self.check_status(status)
        index_type = None
        for k, v in INDEX_MAP.items():
            if result._index_type == v:
                index_type = k
                break
        return {"index_type": index_type, "index_param": result._params}

    def drop_index(self):
        logger.info("Drop index: %s" % self._collection_name)
        return self._milvus.drop_index(self._collection_name)

    @time_wrapper
    def query(self, X, top_k, search_param=None):
        status, result = self._milvus.search_vectors(self._collection_name,
                                                     top_k,
                                                     query_records=X,
                                                     params=search_param)
        self.check_status(status)
        return result

    @time_wrapper
    def query_ids(self, top_k, ids, search_param=None):
        status, result = self._milvus.search_by_ids(self._collection_name,
                                                    ids,
                                                    top_k,
                                                    params=search_param)
        self.check_result_ids(result)
        return result

    def count(self):
        return self._milvus.count_collection(self._collection_name)[1]

    def delete(self, timeout=120):
        timeout = int(timeout)
        logger.info("Start delete collection: %s" % self._collection_name)
        self._milvus.drop_collection(self._collection_name)
        i = 0
        while i < timeout:
            if self.count():
                time.sleep(1)
                i = i + 1
                continue
            else:
                break
        if i >= timeout:
            logger.error("Delete collection timeout")

    def describe(self):
        return self._milvus.describe_collection(self._collection_name)

    def show_collections(self):
        return self._milvus.show_collections()

    def exists_collection(self, collection_name=None):
        if collection_name is None:
            collection_name = self._collection_name
        status, res = self._milvus.has_collection(collection_name)
        # self.check_status(status)
        return res

    @time_wrapper
    def preload_collection(self):
        status = self._milvus.preload_collection(self._collection_name,
                                                 timeout=3000)
        self.check_status(status)
        return status

    def get_server_version(self):
        status, res = self._milvus.server_version()
        return res

    def get_server_mode(self):
        return self.cmd("mode")

    def get_server_commit(self):
        return self.cmd("build_commit_id")

    def get_server_config(self):
        return json.loads(self.cmd("get_config *"))

    def get_mem_info(self):
        result = json.loads(self.cmd("get_system_info"))
        result_human = {
            # unit: Gb
            "memory_used":
            round(int(result["memory_used"]) / (1024 * 1024 * 1024), 2)
        }
        return result_human

    def cmd(self, command):
        status, res = self._milvus._cmd(command)
        logger.info("Server command: %s, result: %s" % (command, res))
        self.check_status(status)
        return res
Пример #10
0
def main():
    milvus = Milvus()

    # Connect to Milvus server
    # You may need to change _HOST and _PORT accordingly
    param = {'host': _HOST, 'port': _PORT}
    status = milvus.connect(**param)

    # Create table demo_table if it dosen't exist.
    table_name = 'demo_table_01'

    if not milvus.has_table(table_name):
        param = {
            'table_name': table_name,
            'dimension': 16,
            'index_file_size': 1024,
            'metric_type': MetricType.L2
        }

        milvus.create_table(param)

    # Show tables in Milvus server
    _, tables = milvus.show_tables()

    # Describe demo_table
    _, table = milvus.describe_table(table_name)

    # 10 vectors with 16 dimension
    vectors = [
        [
            0.66, 0.01, 0.29, 0.64, 0.75, 0.94, 0.26, 0.79, 0.61, 0.11, 0.25,
            0.50, 0.74, 0.37, 0.28, 0.63
        ],
        [
            0.77, 0.65, 0.57, 0.68, 0.29, 0.93, 0.17, 0.15, 0.95, 0.09, 0.78,
            0.37, 0.76, 0.21, 0.42, 0.15
        ],
        [
            0.61, 0.38, 0.32, 0.39, 0.54, 0.93, 0.09, 0.81, 0.52, 0.30, 0.20,
            0.59, 0.15, 0.27, 0.04, 0.37
        ],
        [
            0.33, 0.03, 0.87, 0.47, 0.79, 0.61, 0.46, 0.77, 0.62, 0.70, 0.85,
            0.01, 0.30, 0.41, 0.74, 0.98
        ],
        [
            0.19, 0.80, 0.03, 0.75, 0.22, 0.49, 0.52, 0.91, 0.40, 0.91, 0.79,
            0.08, 0.27, 0.16, 0.07, 0.24
        ],
        [
            0.44, 0.36, 0.16, 0.88, 0.30, 0.79, 0.45, 0.31, 0.45, 0.99, 0.15,
            0.93, 0.37, 0.25, 0.78, 0.84
        ],
        [
            0.33, 0.37, 0.59, 0.66, 0.76, 0.11, 0.19, 0.38, 0.14, 0.37, 0.97,
            0.50, 0.08, 0.69, 0.16, 0.67
        ],
        [
            0.68, 0.97, 0.20, 0.13, 0.30, 0.16, 0.85, 0.21, 0.26, 0.17, 0.81,
            0.96, 0.18, 0.40, 0.13, 0.74
        ],
        [
            0.11, 0.26, 0.44, 0.91, 0.89, 0.79, 0.98, 0.91, 0.09, 0.45, 0.07,
            0.88, 0.71, 0.35, 0.97, 0.41
        ],
        [
            0.17, 0.54, 0.61, 0.58, 0.25, 0.63, 0.65, 0.71, 0.26, 0.80, 0.28,
            0.77, 0.69, 0.02, 0.63, 0.60
        ],
    ]

    # Insert vectors into demo_table
    status, ids = milvus.add_vectors(table_name=table_name, records=vectors)

    # Wait for 6 seconds, since Milvus server persist vector data every 5 seconds by default.
    # You can set data persist interval in Milvus config file.
    time.sleep(6)

    # Get demo_table row count
    status, result = milvus.get_table_row_count(table_name)

    # Use the 3rd vector for similarity search
    query_vectors = [vectors[3]]

    # execute vector similarity search
    param = {
        'table_name': table_name,
        'query_records': query_vectors,
        'top_k': 1,
        'nprobe': 16
    }
    status, results = milvus.search_vectors(**param)

    if results[0][0].distance == 0.0 or results[0][0].id == ids[3]:
        print('Query result is correct')
    else:
        print('Query result isn\'t correct')

    # Delete demo_table
    status = milvus.delete_table(table_name)

    # Disconnect from Milvus
    status = milvus.disconnect()
Пример #11
0
def main():
    milvus = Milvus()

    # Print client version
    print('# Client version: {}'.format(milvus.client_version()))

    # Connect milvus server
    # Please change HOST and PORT to the correct one
    param = {'host': _HOST, 'port': _PORT}
    cnn_status = milvus.connect(**param)
    print('# Connect Status: {}'.format(cnn_status))

    # Check if connected
    # is_connected = milvus.connected
    print('# Is connected: {}'.format(milvus.connected))

    # Print milvus server version
    print('# Server version: {}'.format(milvus.server_version()))

    # Describe table
    table_name = 'table01'
    res_status, table = milvus.describe_table(table_name)
    print('# Describe table status: {}'.format(res_status))
    print('# Describe table:{}'.format(table))

    # Create table
    # Check if `table01` exists, if not, create a table `table01`
    dimension = 256
    if not table:
        param = {
            'table_name': table_name,
            'dimension': dimension,
            'index_type': IndexType.IDMAP,
            'store_raw_vector': False
        }

        res_status = milvus.create_table(Prepare.table_schema(**param))
        print('# Create table status: {}'.format(res_status))

    # Show tables and their description
    status, tables = milvus.show_tables()
    pprint(tables)

    # Add vectors
    # Prepare vector with 256 dimension
    vectors = Prepare.records([[random.random() for _ in range(dimension)]
                               for _ in range(20)])

    # Insert vectors into table 'table01'
    status, ids = milvus.add_vectors(table_name=table_name, records=vectors)
    print('# Add vector status: {}'.format(status))
    pprint(ids)

    # Search vectors
    # When adding vectors for the first time, server will take at least 5s to
    # persist vector data, so you have to wait for 6s after adding vectors for
    # the first time.
    print('# Waiting for 6s...')
    time.sleep(6)

    q_records = Prepare.records([[random.random() for _ in range(dimension)]
                                 for _ in range(2)])

    param = {
        'table_name': table_name,
        'query_records': q_records,
        'top_k': 10,
    }
    status, results = milvus.search_vectors(**param)
    print('# Search vectors status: {}'.format(status))
    pprint(results)

    # Get table row count
    status, result = milvus.get_table_row_count(table_name)
    print('# Status: {}'.format(status))
    print('# Count: {}'.format(result))

    # Disconnect
    status = milvus.disconnect()
    print('# Disconnect Status: {}'.format(status))
Пример #12
0
class MilvusClient(object):
    def __init__(self, table_name=None, ip=None, port=None, timeout=60):
        self._milvus = Milvus()
        self._table_name = table_name
        try:
            i = 1
            start_time = time.time()
            if not ip:
                self._milvus.connect(host=SERVER_HOST_DEFAULT,
                                     port=SERVER_PORT_DEFAULT)
            else:
                # retry connect for remote server
                while time.time() < start_time + timeout:
                    try:
                        self._milvus.connect(host=ip, port=port)
                        if self._milvus.connected() is True:
                            logger.debug(
                                "Try connect times: %d, %s" %
                                (i, round(time.time() - start_time, 2)))
                            break
                    except Exception as e:
                        logger.debug("Milvus connect failed")
                        i = i + 1

        except Exception as e:
            raise e

    def __str__(self):
        return 'Milvus table %s' % self._table_name

    def check_status(self, status):
        if not status.OK():
            logger.error(status.message)
            # raise Exception("Status not ok")

    def create_table(self, table_name, dimension, index_file_size,
                     metric_type):
        if not self._table_name:
            self._table_name = table_name
        if metric_type == "l2":
            metric_type = MetricType.L2
        elif metric_type == "ip":
            metric_type = MetricType.IP
        elif metric_type == "jaccard":
            metric_type = MetricType.JACCARD
        elif metric_type == "hamming":
            metric_type = MetricType.HAMMING
        else:
            logger.error("Not supported metric_type: %s" % metric_type)
        create_param = {
            'table_name': table_name,
            'dimension': dimension,
            'index_file_size': index_file_size,
            "metric_type": metric_type
        }
        status = self._milvus.create_table(create_param)
        self.check_status(status)

    @time_wrapper
    def insert(self, X, ids=None):
        status, result = self._milvus.add_vectors(self._table_name, X, ids)
        self.check_status(status)
        return status, result

    @time_wrapper
    def create_index(self, index_type, nlist):
        index_params = {
            "index_type": INDEX_MAP[index_type],
            "nlist": nlist,
        }
        logger.info("Building index start, table_name: %s, index_params: %s" %
                    (self._table_name, json.dumps(index_params)))
        status = self._milvus.create_index(self._table_name,
                                           index=index_params)
        self.check_status(status)

    def describe_index(self):
        status, result = self._milvus.describe_index(self._table_name)
        index_type = None
        for k, v in INDEX_MAP.items():
            if result._index_type == v:
                index_type = k
                break
        nlist = result._nlist
        res = {"index_type": index_type, "nlist": nlist}
        return res

    def drop_index(self):
        logger.info("Drop index: %s" % self._table_name)
        return self._milvus.drop_index(self._table_name)

    @time_wrapper
    def query(self, X, top_k, nprobe):
        status, result = self._milvus.search_vectors(self._table_name, top_k,
                                                     nprobe, X)
        self.check_status(status)
        return result

    def count(self):
        return self._milvus.get_table_row_count(self._table_name)[1]

    def delete(self, timeout=60):
        logger.info("Start delete table: %s" % self._table_name)
        self._milvus.delete_table(self._table_name)
        i = 0
        while i < timeout:
            if self.count():
                time.sleep(1)
                i = i + 1
                continue
            else:
                break
        if i >= timeout:
            logger.error("Delete table timeout")

    def describe(self):
        return self._milvus.describe_table(self._table_name)

    def show_tables(self):
        return self._milvus.show_tables()

    def exists_table(self, table_name=None):
        if table_name is None:
            table_name = self._table_name
        status, res = self._milvus.has_table(table_name)
        self.check_status(status)
        return res

    @time_wrapper
    def preload_table(self):
        return self._milvus.preload_table(self._table_name, timeout=3000)

    def get_server_version(self):
        status, res = self._milvus.server_version()
        return res

    def get_server_mode(self):
        return self.cmd("mode")

    def get_server_commit(self):
        return self.cmd("build_commit_id")

    def get_server_config(self):
        return json.loads(self.cmd("get_config *"))

    def get_mem_info(self):
        result = json.loads(self.cmd("get_system_info"))
        result_human = {
            # unit: Gb
            "memory_used":
            round(int(result["memory_used"]) / (1024 * 1024 * 1024), 2)
        }
        return result_human

    def cmd(self, command):
        status, res = self._milvus._cmd(command)
        logger.info("Server command: %s, result: %s" % (command, res))
        self.check_status(status)
        return res
Пример #13
0
class MilvusClient(object):
    def __init__(self, table_name=None, host=None, port=None):
        self._milvus = Milvus()
        self._table_name = table_name
        try:
            if not host:
                self._milvus.connect(host=SERVER_HOST_DEFAULT,
                                     port=SERVER_PORT_DEFAULT)
            else:
                self._milvus.connect(host=host, port=port)
        except Exception as e:
            raise e

    def __str__(self):
        return 'Milvus table %s' % self._table_name

    def check_status(self, status):
        if not status.OK():
            logger.error(status.message)
            raise Exception("Status not ok")

    def create_table(self, table_name, dimension, index_file_size,
                     metric_type):
        if not self._table_name:
            self._table_name = table_name
        if metric_type == "l2":
            metric_type = MetricType.L2
        elif metric_type == "ip":
            metric_type = MetricType.IP
        else:
            logger.error("Not supported metric_type: %s" % metric_type)
        self._metric_type = metric_type
        create_param = {
            'table_name': table_name,
            'dimension': dimension,
            'index_file_size': index_file_size,
            "metric_type": metric_type
        }
        status = self._milvus.create_table(create_param)
        self.check_status(status)

    @time_wrapper
    def insert(self, X, ids):
        if self._metric_type == MetricType.IP:
            logger.info("Set normalize for metric_type: Inner Product")
            X = sklearn.preprocessing.normalize(X, axis=1, norm='l2')
        X = X.astype(numpy.float32)
        status, result = self._milvus.add_vectors(self._table_name,
                                                  X.tolist(),
                                                  ids=ids)
        self.check_status(status)
        return status, result

    @time_wrapper
    def create_index(self, index_type, nlist):
        if index_type == "flat":
            index_type = IndexType.FLAT
        elif index_type == "ivf_flat":
            index_type = IndexType.IVFLAT
        elif index_type == "ivf_sq8":
            index_type = IndexType.IVF_SQ8
        elif index_type == "ivf_sq8h":
            index_type = IndexType.IVF_SQ8H
        elif index_type == "nsg":
            index_type = IndexType.NSG
        elif index_type == "ivf_pq":
            index_type = IndexType.IVF_PQ
        index_params = {
            "index_type": index_type,
            "nlist": nlist,
        }
        logger.info("Building index start, table_name: %s, index_params: %s" %
                    (self._table_name, json.dumps(index_params)))
        status = self._milvus.create_index(self._table_name,
                                           index=index_params,
                                           timeout=6 * 3600)
        self.check_status(status)

    def describe_index(self):
        return self._milvus.describe_index(self._table_name)

    def drop_index(self):
        logger.info("Drop index: %s" % self._table_name)
        return self._milvus.drop_index(self._table_name)

    @time_wrapper
    def query(self, X, top_k, nprobe):
        if self._metric_type == MetricType.IP:
            logger.info("Set normalize for metric_type: Inner Product")
            X = sklearn.preprocessing.normalize(X, axis=1, norm='l2')
        X = X.astype(numpy.float32)
        status, results = self._milvus.search_vectors(self._table_name, top_k,
                                                      nprobe, X.tolist())
        self.check_status(status)
        ids = []
        for result in results:
            tmp_ids = []
            for item in result:
                tmp_ids.append(item.id)
            ids.append(tmp_ids)
        return ids

    def count(self):
        return self._milvus.get_table_row_count(self._table_name)[1]

    def delete(self, table_name):
        logger.info("Start delete table: %s" % table_name)
        return self._milvus.delete_table(table_name)

    def describe(self):
        return self._milvus.describe_table(self._table_name)

    def exists_table(self, table_name):
        return self._milvus.has_table(table_name)

    def get_server_version(self):
        status, res = self._milvus.server_version()
        self.check_status(status)
        return res

    @time_wrapper
    def preload_table(self):
        return self._milvus.preload_table(self._table_name)