ExecutionEngineImpl.h

/*******************************************************************************
 * Copyright 上海赜睿信息科技有限公司(Zilliz) - All Rights Reserved
 * Unauthorized copying of this file, via any medium is strictly prohibited.
 * Proprietary and confidential.
 ******************************************************************************/
#pragma once

#include "ExecutionEngine.h"
#include "wrapper/knowhere/vec_index.h"

#include <memory>
#include <string>


namespace zilliz {
namespace milvus {
namespace engine {


class ExecutionEngineImpl : public ExecutionEngine {
public:

    ExecutionEngineImpl(uint16_t dimension,
                        const std::string &location,
                        EngineType index_type,
                        MetricType metric_type,
                        int32_t nlist);

    ExecutionEngineImpl(VecIndexPtr index,
                        const std::string &location,
                        EngineType index_type,
                        MetricType metric_type,
                        int32_t nlist);

    Status AddWithIds(long n, const float *xdata, const long *xids) override;

    size_t Count() const override;

    size_t Size() const override;

    size_t Dimension() const override;

    size_t PhysicalSize() const override;

    Status Serialize() override;

    Status Load(bool to_cache) override;

    Status CopyToGpu(uint64_t device_id) override;

    Status CopyToCpu() override;

    Status Merge(const std::string &location) override;

    Status Search(long n,
                  const float *data,
                  long k,
                  long nprobe,
                  float *distances,
                  long *labels) const override;

    ExecutionEnginePtr BuildIndex(const std::string &) override;

    Status Cache() override;

    Status GpuCache(uint64_t gpu_id) override;

    Status Init() override;

    EngineType IndexEngineType() const override { return index_type_; }

    MetricType IndexMetricType() const override { return metric_type_; }

private:
    VecIndexPtr CreatetVecIndex(EngineType type);

    VecIndexPtr Load(const std::string &location);

protected:
    VecIndexPtr index_ = nullptr;
    EngineType index_type_;
    MetricType metric_type_;

    int64_t dim_;
    std::string location_;

    int32_t nlist_ = 0;
    int64_t gpu_num_ = 0;
};


} // namespace engine
} // namespace milvus
} // namespace zilliz