乐趣区

关于数据库:进一寸有进一寸的欢喜谈谈如何优化-Milvus-数据库的向量查询功能

✏️ 编者按

每年暑期,Milvus 社区都会携手中科院软件所,在「开源之夏」流动中为高校学生们筹备丰盛的工程项目,并安顿导师答疑解惑。张煜旻同学在「开源之夏」流动中体现优良,他置信进一寸有进一寸的欢喜,尝试在奉献开源的过程中超越自我。

他的我的项目为 Milvus 数据库的向量查问操作提供精度管制,能让开发者自定义返回精度,在缩小内存耗费的同时,进步了返回后果的可读性。

想要理解更多优质开源我的项目和我的项目教训分享?请戳:有哪些值得参加的开源我的项目?

我的项目简介

项目名称:反对指定搜寻时返回的间隔精度

学生简介:张煜旻,中国科学院大学电子信息软件工程业余硕士在读

我的项目导师:Zilliz 软件工程师张财

导师评语:张煜旻同学优化了 Milvus 数据库的查问性能,使其在搜寻时能够用指定精度去进行查问,使搜寻过程更灵便,用户能够依据本人的需要用不同的精度进行查问,给用户带来了便当。

反对指定搜寻时返回的间隔精度

工作简介

在进行向量查问时,搜寻申请返回 id 和 distance 字段,其中的 distance 字段类型是浮点数。Milvus 数据库所计算的间隔是一个 32 位浮点数,然而 Python SDK 返回并以 64 位浮点显示它,导致某些精度有效。本我的项目的奉献是,反对指定搜寻时返回的间隔精度,解决了在 Python 端显示时局部精度有效的状况,并缩小局部内存开销。

我的项目指标

  • 解决计算结果和显示精度不匹配的问题
  • 反对搜寻时返回指定的间隔精度
  • 补充相干文档

我的项目步骤

  • 后期调研,了解 Milvus 整体框架
  • 明确各模块之间的调用关系
  • 设计解决方案和确认后果

我的项目综述

什么是 Milvus 数据库?

Milvus 是一款开源向量数据库,赋能 AI 利用和向量类似度搜寻。在零碎设计上,Milvus 数据库的前端有不便用户应用的 Python SDK(Client);在 Milvus 数据库的后端,整个零碎分为了接入层(Access Layer)、协调服务(Coordinator Server)、执行节点(Worker Node)和存储服务(Storge)四个层面:

(1)接入层(Access Layer):零碎的门面,蕴含了一组对等的 Proxy 节点。接入层是裸露给用户的对立 endpoint,负责转发申请并收集执行后果。

(2)协调服务(Coordinator Service):零碎的大脑,负责分配任务给执行节点。共有四类协调者角色:root 协调者、data 协调者、query 协调者和 index 协调者。

(3)执行节点(Worker Node):零碎的四肢,执行节点只负责被动执行协调服务发动的读写申请。目前有三类执行节点:data 节点、query 节点和 index 节点。

(4)存储服务(Storage):零碎的骨骼,是所有其余性能实现的根底。Milvus  数据库依赖三类存储:元数据存储、音讯存储(log broker)和对象存储。从语言角度来看,则能够看作三个语言层,别离是 Python 形成的 SDK 层、Go 形成的中间层和 C++ 形成的外围计算层。

Milvus 数据库的架构图

向量查问 Search 时,到底产生了什么?

在 Python SDK 端,当用户发动一个 Search API 调用时,这个调用会被封装成 gRPC 申请并发送给 Milvus 后端,同时 SDK 开始期待。而在后端,Proxy 节点首先承受了从 Python SDK 发送过去的申请,而后会对承受的申请进行解决,最初将其封装成 message,经由 Producer 发送到生产队列中。当音讯被发送到生产队列后,Coordinator 将会对其进行协调,将信息发送到适合的 query node 中进行生产。而当 query node 接管到音讯后,则会对音讯进行进一步的解决,最初将信息传递给由 C++ 形成的计算层。在计算层,则会依据不同的情景,调用不同的计算函数对向量间的间隔进行计算。当计算实现后,后果则会顺次向上传递,直到达到 SDK 端。

解决方案设计

通过前文简略介绍,咱们对向量查问的过程有了一个大抵的概念。同时,咱们也能够分明地意识到,为了实现查问指标,咱们须要对 Python 形成的 SDK 层、Go 形成的中间层和 C++ 形成的计算层都进行批改,批改计划如下:

1. 在 Python 层中的批改步骤:

为向量查问 Search 申请增加一个 round_decimal 参数,从而确定返回的精度信息。同时,须要对参数进行一些合法性检查和异样解决,从而构建 gRPC 的申请:

round_decimal = param_copy("round_decimal", 3)
if not isinstance(round_decimal, (int, str))
  raise ParamError("round_decimal must be int or str")
try:
  round_decimal = int(round_decimal)
except Exception:
  raise ParamError("round_decimal is not illegal")

if round_decimal < 0 or round_decimal > 6:
  raise ParamError("round_decimal must be greater than zero and less than seven")
if not instance(params, dict):
  raise ParamError("Search params must be a dict")
search_params = {"anns_field": anns_field, "topk": limit, "metric_type": metric_type, "params": params, "round_decimal": round_decimal}

2. 在 Go 层中的批改步骤:

在 task.go 文件中增加 RoundDecimalKey 这个常量,放弃格调对立并不便后续调取:

const (
 InsertTaskName                  = "InsertTask"
 CreateCollectionTaskName        = "CreateCollectionTask"
 DropCollectionTaskName          = "DropCollectionTask"
 SearchTaskName                  = "SearchTask"
 RetrieveTaskName                = "RetrieveTask"
 QueryTaskName                   = "QueryTask"
 AnnsFieldKey                    = "anns_field"
 TopKKey                         = "topk"
 MetricTypeKey                   = "metric_type"
 SearchParamsKey                 = "params"
 RoundDecimalKey                 = "round_decimal"
 HasCollectionTaskName           = "HasCollectionTask"
 DescribeCollectionTaskName      = "DescribeCollectionTask"

接着,批改 PreExecute 函数,获取 round_decimal 的值,构建 queryInfo 变量,并增加异样解决:

searchParams, err := funcutil.GetAttrByKeyFromRepeatedKV(SearchParamsKey, st.query.SearchParams)
        if err != nil {return errors.New(SearchParamsKey + "not found in search_params")
        }
        roundDecimalStr, err := funcutil.GetAttrByKeyFromRepeatedKV(RoundDecimalKey, st.query.SearchParams)
        if err != nil {return errors.New(RoundDecimalKey + "not found in search_params")
        }
        roundDeciaml, err := strconv.Atoi(roundDecimalStr)
        if err != nil {return errors.New(RoundDecimalKey + "" + roundDecimalStr +" is not invalid")
        }

        queryInfo := &planpb.QueryInfo{Topk:         int64(topK),
            MetricType:   metricType,
            SearchParams: searchParams,
            RoundDecimal: int64(roundDeciaml),
        }

同时,批改 query 的 proto 文件,为 QueryInfo 增加 round_decimal 变量:

message QueryInfo {
int64 topk = 1;
string metric_type = 3;
string search_params = 4;
int64 round_decimal = 5;
}

3. 在 C++ 层中的批改步骤:

在 SearchInfo 构造体中增加新的变量 round\_decimal\_,从而承受 Go 层传来的 round_decimal 值:

struct SearchInfo {
   int64_t topk_;
   int64_t round_decimal_;
   FieldOffset field_offset_;
   MetricType metric_type_;
   nlohmann::json search_params_;
};

在 ParseVecNode 和 PlanNodeFromProto 函数中,SearchInfo 构造体须要承受 Go 层中 round_decimal 值:

std::unique_ptr<VectorPlanNode>
Parser::ParseVecNode(const Json& out_body) {Assert(out_body.is_object());
    Assert(out_body.size() == 1);
    auto iter = out_body.begin();
    auto field_name = FieldName(iter.key());

    auto& vec_info = iter.value();
    Assert(vec_info.is_object());
    auto topk = vec_info["topk"];
    AssertInfo(topk > 0, "topk must greater than 0");
    AssertInfo(topk < 16384, "topk is too large");

    auto field_offset = schema.get_offset(field_name);

    auto vec_node = [&]() -> std::unique_ptr<VectorPlanNode> {auto& field_meta = schema.operator[](field_name);
        auto data_type = field_meta.get_data_type();
        if (data_type == DataType::VECTOR_FLOAT) {return std::make_unique<FloatVectorANNS>();
        } else {return std::make_unique<BinaryVectorANNS>();
        }
    }();
    vec_node->search_info_.topk_ = topk;
    vec_node->search_info_.metric_type_ = GetMetricType(vec_info.at("metric_type"));
    vec_node->search_info_.search_params_ = vec_info.at("params");
    vec_node->search_info_.field_offset_ = field_offset;
    vec_node->search_info_.round_decimal_ = vec_info.at("round_decimal");
    vec_node->placeholder_tag_ = vec_info.at("query");
    auto tag = vec_node->placeholder_tag_;
    AssertInfo(!tag2field_.count(tag), "duplicated placeholder tag");
    tag2field_.emplace(tag, field_offset);
    return vec_node;
}
std::unique_ptr<VectorPlanNode>
ProtoParser::PlanNodeFromProto(const planpb::PlanNode& plan_node_proto) {
    // TODO: add more buffs
    Assert(plan_node_proto.has_vector_anns());
    auto& anns_proto = plan_node_proto.vector_anns();
    auto expr_opt = [&]() -> std::optional<ExprPtr> {if (!anns_proto.has_predicates()) {return std::nullopt;} else {return ParseExpr(anns_proto.predicates());
        }
    }();

    auto& query_info_proto = anns_proto.query_info();

    SearchInfo search_info;
    auto field_id = FieldId(anns_proto.field_id());
    auto field_offset = schema.get_offset(field_id);
    search_info.field_offset_ = field_offset;

    search_info.metric_type_ = GetMetricType(query_info_proto.metric_type());
    search_info.topk_ = query_info_proto.topk();
    search_info.round_decimal_ = query_info_proto.round_decimal();
    search_info.search_params_ = json::parse(query_info_proto.search_params());

    auto plan_node = [&]() -> std::unique_ptr<VectorPlanNode> {if (anns_proto.is_binary()) {return std::make_unique<BinaryVectorANNS>();
        } else {return std::make_unique<FloatVectorANNS>();
        }
    }();
    plan_node->placeholder_tag_ = anns_proto.placeholder_tag();
    plan_node->predicate_ = std::move(expr_opt);
    plan_node->search_info_ = std::move(search_info);
    return plan_node;
}

在 SubSearchResult 类增加新的成员变量 round_decimal,同时批改每一处的 SubSearchResult 变量申明:

class SubSearchResult {
public:
   SubSearchResult(int64_t num_queries, int64_t topk, MetricType metric_type)
       : metric_type_(metric_type),
         num_queries_(num_queries),
         topk_(topk),
         labels_(num_queries * topk, -1),
         values_(num_queries * topk, init_value(metric_type)) {}

在 SubSearchResult 类增加一个新的成员函数,以便最初对每一个后果进行四舍五入精度管制:

void
SubSearchResult::round_values() {if (round_decimal_ == -1)
        return;
    const float multiplier = pow(10.0, round_decimal_);
    for (auto it = this->values_.begin(); it != this->values_.end(); it++) {*it = round(*it * multiplier) / multiplier;
    }
}

为 SearchDataset 构造体增加新的变量 round_decimal,同时批改每一处的 SearchDataset 变量申明:

struct SearchDataset {
    MetricType metric_type;
    int64_t num_queries;
    int64_t topk;
    int64_t round_decimal;
    int64_t dim;
    const void* query_data;
};

批改 C++ 层中各个间隔计算函数(FloatSearch、BinarySearchBruteForceFast 等等),使其承受 round_decomal 值:

Status
FloatSearch(const segcore::SegmentGrowingImpl& segment,
            const query::SearchInfo& info,
            const float* query_data,
            int64_t num_queries,
            int64_t ins_barrier,
            const BitsetView& bitset,
            SearchResult& results) {auto& schema = segment.get_schema();
    auto& indexing_record = segment.get_indexing_record();
    auto& record = segment.get_insert_record();
    // step 1: binary search to find the barrier of the snapshot
    // auto del_barrier = get_barrier(deleted_record_, timestamp);

#if 0
    auto bitmap_holder = get_deleted_bitmap(del_barrier, timestamp, ins_barrier);
    Assert(bitmap_holder);
    auto bitmap = bitmap_holder->bitmap_ptr;
#endif

    // step 2.1: get meta
    // step 2.2: get which vector field to search
    auto vecfield_offset = info.field_offset_;
    auto& field = schema[vecfield_offset];

    AssertInfo(field.get_data_type() == DataType::VECTOR_FLOAT, "[FloatSearch]Field data type isn't VECTOR_FLOAT");
    auto dim = field.get_dim();
    auto topk = info.topk_;
    auto total_count = topk * num_queries;
    auto metric_type = info.metric_type_;
    auto round_decimal = info.round_decimal_;
    // step 3: small indexing search
    // std::vector<int64_t> final_uids(total_count, -1);
    // std::vector<float> final_dis(total_count, std::numeric_limits<float>::max());
    SubSearchResult final_qr(num_queries, topk, metric_type, round_decimal);
    dataset::SearchDataset search_dataset{metric_type, num_queries, topk, round_decimal, dim, query_data};
    auto vec_ptr = record.get_field_data<FloatVector>(vecfield_offset);

    int current_chunk_id = 0;
SubSearchResult
BinarySearchBruteForceFast(MetricType metric_type,
                           int64_t dim,
                           const uint8_t* binary_chunk,
                           int64_t size_per_chunk,
                           int64_t topk,
                           int64_t num_queries,
                           int64_t round_decimal,
                           const uint8_t* query_data,
                           const faiss::BitsetView& bitset) {SubSearchResult sub_result(num_queries, topk, metric_type, round_decimal);
    float* result_distances = sub_result.get_values();
    idx_t* result_labels = sub_result.get_labels();

    int64_t code_size = dim / 8;
    const idx_t block_size = size_per_chunk;

    raw_search(metric_type, binary_chunk, size_per_chunk, code_size, num_queries, query_data, topk, result_distances,
               result_labels, bitset);
    sub_result.round_values();
    return sub_result;
}

后果确认

1. 对 Milvus 数据库进行从新编译:

2. 启动环境容器:

3. 启动 Milvus 数据库:

4. 构建向量查问申请:

5. 确认后果,默认保留 3 位小数,0 舍去:

总结和感想

加入这次的冬季开源流动,对 我来说是十分贵重的经验。在这次流动中,我第一次尝试浏览开源我的项目代码,第一次尝试接触多语言形成的我的项目,第一次接触到 Make、gRPc、pytest 等等。在编写代码和测试代码阶段,我也遇到来许多意想不到的问题,例如,「奇奇怪怪」的依赖问题、因为 Conda 环境导致的编译失败问题、测试无奈通过等等。面对这些问题,我 慢慢学会急躁细心地查看报错日志,积极思考、查看代码并进行测试,一步一步放大谬误范畴,定位错误代码并尝试各种解决方案。

通过这次的流动,我汲取 了 许 多 教训和教训,同时也非常感激张财导师,感激他在我开发过程中 急躁地帮我答疑解惑、领导方向!同时,心愿大家能多多关注 Milvus 社区,置信 肯定可能有所播种!

最初,欢送大家多多与我交换(📮 deepmin@mail.deepexplore.top),我次要的钻研方向是自然语言解决,平时喜爱看科幻小说、动画和折腾服务器集体网站,每日晃荡 Stack Overflow 和 GitHub。我置信进一寸有进一寸的欢喜,心愿能和你一起共同进步。

退出移动版