ConfAdapter.h 2.6 KB
Newer Older
X
xiaojun.lin 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#pragma once

#include "VecIndex.h"
S
starlord 已提交
21
#include "knowhere/common/Config.h"
X
xiaojun.lin 已提交
22

S
starlord 已提交
23
#include <memory>
X
xiaojun.lin 已提交
24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43

namespace milvus {
namespace engine {

// TODO(linxj): remove later, replace with real metaconf
constexpr int64_t TEMPMETA_DEFAULT_VALUE = -1;
struct TempMetaConf {
    int64_t size = TEMPMETA_DEFAULT_VALUE;
    int64_t nlist = TEMPMETA_DEFAULT_VALUE;
    int64_t dim = TEMPMETA_DEFAULT_VALUE;
    int64_t gpu_id = TEMPMETA_DEFAULT_VALUE;
    int64_t k = TEMPMETA_DEFAULT_VALUE;
    int64_t nprobe = TEMPMETA_DEFAULT_VALUE;
    int64_t search_length = TEMPMETA_DEFAULT_VALUE;
    knowhere::METRICTYPE metric_type = knowhere::DEFAULT_TYPE;
};

class ConfAdapter {
 public:
    virtual knowhere::Config
S
starlord 已提交
44
    Match(const TempMetaConf& metaconf);
X
xiaojun.lin 已提交
45 46

    virtual knowhere::Config
S
starlord 已提交
47
    MatchSearch(const TempMetaConf& metaconf, const IndexType& type);
X
xiaojun.lin 已提交
48

S
starlord 已提交
49 50
    //    virtual void
    //    Dump(){}
51

X
xiaojun.lin 已提交
52
 protected:
S
starlord 已提交
53 54
    static void
    MatchBase(knowhere::Config conf);
X
xiaojun.lin 已提交
55 56 57 58 59 60 61
};

using ConfAdapterPtr = std::shared_ptr<ConfAdapter>;

class IVFConfAdapter : public ConfAdapter {
 public:
    knowhere::Config
S
starlord 已提交
62
    Match(const TempMetaConf& metaconf) override;
X
xiaojun.lin 已提交
63 64

    knowhere::Config
S
starlord 已提交
65
    MatchSearch(const TempMetaConf& metaconf, const IndexType& type) override;
X
xiaojun.lin 已提交
66 67

 protected:
S
starlord 已提交
68 69
    static int64_t
    MatchNlist(const int64_t& size, const int64_t& nlist);
X
xiaojun.lin 已提交
70 71 72 73 74
};

class IVFSQConfAdapter : public IVFConfAdapter {
 public:
    knowhere::Config
S
starlord 已提交
75
    Match(const TempMetaConf& metaconf) override;
X
xiaojun.lin 已提交
76 77 78 79 80
};

class IVFPQConfAdapter : public IVFConfAdapter {
 public:
    knowhere::Config
S
starlord 已提交
81
    Match(const TempMetaConf& metaconf) override;
X
xiaojun.lin 已提交
82 83 84 85 86
};

class NSGConfAdapter : public IVFConfAdapter {
 public:
    knowhere::Config
S
starlord 已提交
87
    Match(const TempMetaConf& metaconf) override;
X
xiaojun.lin 已提交
88 89

    knowhere::Config
S
starlord 已提交
90
    MatchSearch(const TempMetaConf& metaconf, const IndexType& type) final;
X
xiaojun.lin 已提交
91 92
};

S
starlord 已提交
93 94
}  // namespace engine
}  // namespace milvus