Skip to content

Commit 95e82ed

Browse files
yhmoyhmo
committed
Add GetLoadState refresh progress and segment info parity
Signed-off-by: yhmo <yihua.mo@zilliz.com>
1 parent 89b654f commit 95e82ed

16 files changed

Lines changed: 272 additions & 20 deletions

‎src/impl/MilvusClientImpl.cpp‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1481,10 +1481,11 @@ MilvusClientImpl::GetPersistentSegmentInfo(const std::string& collection_name, S
14811481
return Status::OK();
14821482
};
14831483

1484-
auto post = [&segments_info](const proto::milvus::GetPersistentSegmentInfoResponse& response) {
1484+
auto post = [&collection_name, &segments_info](const proto::milvus::GetPersistentSegmentInfoResponse& response) {
14851485
for (const auto& info : response.infos()) {
14861486
segments_info.emplace_back(info.collectionid(), info.partitionid(), info.segmentid(), info.num_rows(),
1487-
SegmentStateCast(info.state()));
1487+
SegmentStateCast(info.state()), collection_name, SegmentLevelCast(info.level()),
1488+
info.storage_version(), info.is_sorted());
14881489
}
14891490
return Status::OK();
14901491
};
@@ -1501,14 +1502,17 @@ MilvusClientImpl::GetQuerySegmentInfo(const std::string& collection_name, QueryS
15011502
return Status::OK();
15021503
};
15031504

1504-
auto post = [&segments_info](const proto::milvus::GetQuerySegmentInfoResponse& response) {
1505+
auto post = [&collection_name, &segments_info](const proto::milvus::GetQuerySegmentInfoResponse& response) {
15051506
for (const auto& info : response.infos()) {
15061507
std::vector<int64_t> ids;
1508+
ids.reserve(info.nodeids_size());
15071509
for (auto id : info.nodeids()) {
15081510
ids.push_back(id);
15091511
}
15101512
segments_info.emplace_back(info.collectionid(), info.partitionid(), info.segmentid(), info.num_rows(),
1511-
milvus::SegmentStateCast(info.state()), info.index_name(), info.indexid(), ids);
1513+
milvus::SegmentStateCast(info.state()), info.index_name(), info.indexid(), ids,
1514+
collection_name, info.mem_size(), SegmentLevelCast(info.level()),
1515+
info.storage_version(), info.is_sorted());
15121516
}
15131517
return Status::OK();
15141518
};

‎src/impl/MilvusClientV2Impl.cpp‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -585,14 +585,19 @@ MilvusClientV2Impl::getLoadState(const GetLoadStateRequest& request, GetLoadStat
585585
auto state = rpc_response.state();
586586
response.SetState(LoadStateCast(state));
587587

588+
response.SetProgress(0);
589+
response.SetRefreshProgress(0);
588590
if (state == proto::common::LoadState::LoadStateLoading) {
589591
uint32_t progress = 0;
592+
uint32_t refresh_progress = 0;
590593
auto status = connection_.GetLoadingProgress(request.DatabaseName(), request.CollectionName(),
591-
request.PartitionNames(), progress, rpc_timeout_ms);
594+
request.PartitionNames(), progress, refresh_progress,
595+
rpc_timeout_ms);
592596
if (!status.IsOk()) {
593597
return status;
594598
}
595599
response.SetProgress(progress);
600+
response.SetRefreshProgress(refresh_progress);
596601
} else if (state == proto::common::LoadState::LoadStateLoaded) {
597602
response.SetProgress(100);
598603
}
@@ -1977,12 +1982,13 @@ MilvusClientV2Impl::ListPersistentSegments(const ListPersistentSegmentsRequest&
19771982
return Status::OK();
19781983
};
19791984

1980-
auto post = [&response](const proto::milvus::GetPersistentSegmentInfoResponse& rpc_response) {
1985+
auto post = [&request, &response](const proto::milvus::GetPersistentSegmentInfoResponse& rpc_response) {
19811986
SegmentsInfo segments_info;
19821987
segments_info.reserve(rpc_response.infos_size());
19831988
for (const auto& info : rpc_response.infos()) {
19841989
segments_info.emplace_back(info.collectionid(), info.partitionid(), info.segmentid(), info.num_rows(),
1985-
SegmentStateCast(info.state()));
1990+
SegmentStateCast(info.state()), request.CollectionName(),
1991+
SegmentLevelCast(info.level()), info.storage_version(), info.is_sorted());
19861992
}
19871993
response.SetResult(std::move(segments_info));
19881994
return Status::OK();
@@ -2001,16 +2007,19 @@ MilvusClientV2Impl::ListQuerySegments(const ListQuerySegmentsRequest& request, L
20012007
return Status::OK();
20022008
};
20032009

2004-
auto post = [&response](const proto::milvus::GetQuerySegmentInfoResponse& rpc_response) {
2010+
auto post = [&request, &response](const proto::milvus::GetQuerySegmentInfoResponse& rpc_response) {
20052011
QuerySegmentsInfo segments_info;
20062012
segments_info.reserve(rpc_response.infos_size());
20072013
for (const auto& info : rpc_response.infos()) {
20082014
std::vector<int64_t> ids;
2015+
ids.reserve(info.nodeids_size());
20092016
for (auto id : info.nodeids()) {
20102017
ids.push_back(id);
20112018
}
20122019
segments_info.emplace_back(info.collectionid(), info.partitionid(), info.segmentid(), info.num_rows(),
2013-
milvus::SegmentStateCast(info.state()), info.index_name(), info.indexid(), ids);
2020+
milvus::SegmentStateCast(info.state()), info.index_name(), info.indexid(), ids,
2021+
request.CollectionName(), info.mem_size(), SegmentLevelCast(info.level()),
2022+
info.storage_version(), info.is_sorted());
20142023
}
20152024
response.SetResult(std::move(segments_info));
20162025
return Status::OK();

‎src/impl/response/collection/GetLoadStateResponse.cpp‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,4 +38,14 @@ GetLoadStateResponse::SetProgress(int64_t progress) {
3838
progress_ = progress;
3939
}
4040

41+
int64_t
42+
GetLoadStateResponse::RefreshProgress() const {
43+
return refresh_progress_;
44+
}
45+
46+
void
47+
GetLoadStateResponse::SetRefreshProgress(int64_t refresh_progress) {
48+
refresh_progress_ = refresh_progress;
49+
}
50+
4151
} // namespace milvus

‎src/impl/types/SegmentInfo.cpp‎

Lines changed: 50 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,15 +16,27 @@
1616

1717
#include "milvus/types/SegmentInfo.h"
1818

19+
#include <utility>
20+
1921
namespace milvus {
2022

2123
SegmentInfo::SegmentInfo(int64_t collection_id, int64_t partition_id, int64_t segment_id, int64_t row_count,
2224
SegmentState state)
25+
: SegmentInfo(collection_id, partition_id, segment_id, row_count, state, "", SegmentLevel::UNKNOWN, 0, false) {
26+
}
27+
28+
SegmentInfo::SegmentInfo(int64_t collection_id, int64_t partition_id, int64_t segment_id, int64_t row_count,
29+
SegmentState state, std::string collection_name, SegmentLevel level, int64_t storage_version,
30+
bool is_sorted)
2331
: collection_id_{collection_id},
2432
partition_id_{partition_id},
2533
segment_id_{segment_id},
2634
row_count_{row_count},
27-
state_(state) {
35+
state_(state),
36+
collection_name_{std::move(collection_name)},
37+
level_{level},
38+
storage_version_{storage_version},
39+
is_sorted_{is_sorted} {
2840
}
2941

3042
int64_t
@@ -52,13 +64,43 @@ SegmentInfo::State() const {
5264
return state_;
5365
}
5466

67+
const std::string&
68+
SegmentInfo::CollectionName() const {
69+
return collection_name_;
70+
}
71+
72+
SegmentLevel
73+
SegmentInfo::Level() const {
74+
return level_;
75+
}
76+
77+
int64_t
78+
SegmentInfo::StorageVersion() const {
79+
return storage_version_;
80+
}
81+
82+
bool
83+
SegmentInfo::IsSorted() const {
84+
return is_sorted_;
85+
}
86+
5587
QuerySegmentInfo::QuerySegmentInfo(int64_t collection_id, int64_t partition_id, int64_t segment_id, int64_t row_count,
5688
SegmentState state, std::string index_name, int64_t index_id,
5789
const std::vector<int64_t>& node_ids)
58-
: SegmentInfo(collection_id, partition_id, segment_id, row_count, state),
90+
: QuerySegmentInfo(collection_id, partition_id, segment_id, row_count, state, std::move(index_name), index_id,
91+
node_ids, "", 0, SegmentLevel::UNKNOWN, 0, false) {
92+
}
93+
94+
QuerySegmentInfo::QuerySegmentInfo(int64_t collection_id, int64_t partition_id, int64_t segment_id, int64_t row_count,
95+
SegmentState state, std::string index_name, int64_t index_id,
96+
const std::vector<int64_t>& node_ids, std::string collection_name, int64_t mem_size,
97+
SegmentLevel level, int64_t storage_version, bool is_sorted)
98+
: SegmentInfo(collection_id, partition_id, segment_id, row_count, state, std::move(collection_name), level,
99+
storage_version, is_sorted),
59100
index_name_{std::move(index_name)},
60101
index_id_{index_id},
61-
node_ids_{node_ids} {
102+
node_ids_{node_ids},
103+
mem_size_{mem_size} {
62104
}
63105

64106
std::string
@@ -84,4 +126,9 @@ QuerySegmentInfo::NodeIDs() const {
84126
return node_ids_;
85127
}
86128

129+
int64_t
130+
QuerySegmentInfo::MemSize() const {
131+
return mem_size_;
132+
}
133+
87134
} // namespace milvus

‎src/impl/utils/CompareUtils.cpp‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -305,14 +305,19 @@ operator==(const proto::schema::FieldData& lhs, const Field& rhs) {
305305
bool
306306
operator==(const SegmentInfo& lhs, const SegmentInfo& rhs) {
307307
return lhs.CollectionID() == rhs.CollectionID() && lhs.PartitionID() == rhs.PartitionID() &&
308-
lhs.RowCount() == rhs.RowCount() && lhs.SegmentID() == rhs.SegmentID() && lhs.State() == rhs.State();
308+
lhs.RowCount() == rhs.RowCount() && lhs.SegmentID() == rhs.SegmentID() && lhs.State() == rhs.State() &&
309+
lhs.CollectionName() == rhs.CollectionName() && lhs.Level() == rhs.Level() &&
310+
lhs.StorageVersion() == rhs.StorageVersion() && lhs.IsSorted() == rhs.IsSorted();
309311
}
310312

311313
bool
312314
operator==(const QuerySegmentInfo& lhs, const QuerySegmentInfo& rhs) {
313315
return lhs.CollectionID() == rhs.CollectionID() && lhs.PartitionID() == rhs.PartitionID() &&
314316
lhs.RowCount() == rhs.RowCount() && lhs.SegmentID() == rhs.SegmentID() && lhs.State() == rhs.State() &&
315-
lhs.IndexName() == rhs.IndexName() && lhs.IndexID() == rhs.IndexID() && lhs.NodeID() == rhs.NodeID();
317+
lhs.CollectionName() == rhs.CollectionName() && lhs.Level() == rhs.Level() &&
318+
lhs.StorageVersion() == rhs.StorageVersion() && lhs.IsSorted() == rhs.IsSorted() &&
319+
lhs.IndexName() == rhs.IndexName() && lhs.IndexID() == rhs.IndexID() && lhs.NodeIDs() == rhs.NodeIDs() &&
320+
lhs.MemSize() == rhs.MemSize();
316321
}
317322

318323
} // namespace milvus

‎src/impl/utils/ConnectionHandler.cpp‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,14 @@ Status
111111
ConnectionHandler::GetLoadingProgress(const std::string& db_name, const std::string& collection_name,
112112
const std::set<std::string> partition_names, uint32_t& progress,
113113
uint64_t rpc_timeout_ms) {
114+
uint32_t refresh_progress = 0;
115+
return GetLoadingProgress(db_name, collection_name, partition_names, progress, refresh_progress, rpc_timeout_ms);
116+
}
117+
118+
Status
119+
ConnectionHandler::GetLoadingProgress(const std::string& db_name, const std::string& collection_name,
120+
const std::set<std::string> partition_names, uint32_t& progress,
121+
uint32_t& refresh_progress, uint64_t rpc_timeout_ms) {
114122
MilvusConnectionPtr connection;
115123
uint64_t timeout = 0;
116124
{
@@ -138,6 +146,7 @@ ConnectionHandler::GetLoadingProgress(const std::string& db_name, const std::str
138146
return status;
139147
}
140148
progress = static_cast<uint32_t>(progress_resp.progress());
149+
refresh_progress = static_cast<uint32_t>(progress_resp.refresh_progress());
141150
return Status::OK();
142151
}
143152

‎src/impl/utils/ConnectionHandler.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,11 @@ class ConnectionHandler {
6767
GetLoadingProgress(const std::string& db_name, const std::string& collection_name,
6868
const std::set<std::string> partition_names, uint32_t& progress, uint64_t rpc_timeout_ms = 0);
6969

70+
Status
71+
GetLoadingProgress(const std::string& db_name, const std::string& collection_name,
72+
const std::set<std::string> partition_names, uint32_t& progress,
73+
uint32_t& refresh_progress, uint64_t rpc_timeout_ms = 0);
74+
7075
/**
7176
* Internal wait for status query done.
7277
*

‎src/impl/utils/TypeUtils.cpp‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -700,6 +700,38 @@ SegmentStateCast(SegmentState state) {
700700
}
701701
}
702702

703+
SegmentLevel
704+
SegmentLevelCast(proto::common::SegmentLevel level) {
705+
switch (level) {
706+
case proto::common::SegmentLevel::Legacy:
707+
return SegmentLevel::LEGACY;
708+
case proto::common::SegmentLevel::L0:
709+
return SegmentLevel::L0;
710+
case proto::common::SegmentLevel::L1:
711+
return SegmentLevel::L1;
712+
case proto::common::SegmentLevel::L2:
713+
return SegmentLevel::L2;
714+
default:
715+
return SegmentLevel::UNKNOWN;
716+
}
717+
}
718+
719+
proto::common::SegmentLevel
720+
SegmentLevelCast(SegmentLevel level) {
721+
switch (level) {
722+
case SegmentLevel::LEGACY:
723+
return proto::common::SegmentLevel::Legacy;
724+
case SegmentLevel::L0:
725+
return proto::common::SegmentLevel::L0;
726+
case SegmentLevel::L1:
727+
return proto::common::SegmentLevel::L1;
728+
case SegmentLevel::L2:
729+
return proto::common::SegmentLevel::L2;
730+
default:
731+
return proto::common::SegmentLevel::Legacy;
732+
}
733+
}
734+
703735
IndexStateCode
704736
IndexStateCast(proto::common::IndexState state) {
705737
switch (state) {

‎src/impl/utils/TypeUtils.h‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,12 @@ SegmentStateCast(proto::common::SegmentState state);
107107
proto::common::SegmentState
108108
SegmentStateCast(SegmentState state);
109109

110+
SegmentLevel
111+
SegmentLevelCast(proto::common::SegmentLevel level);
112+
113+
proto::common::SegmentLevel
114+
SegmentLevelCast(SegmentLevel level);
115+
110116
IndexStateCode
111117
IndexStateCast(proto::common::IndexState state);
112118

‎src/include/milvus/response/collection/GetLoadStateResponse.h‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,9 +57,22 @@ class MILVUS_SDK_API GetLoadStateResponse {
5757
void
5858
SetProgress(int64_t progress);
5959

60+
/**
61+
* @brief Get percent value of refresh load progress.
62+
*/
63+
int64_t
64+
RefreshProgress() const;
65+
66+
/**
67+
* @brief Set percent value of refresh load progress.
68+
*/
69+
void
70+
SetRefreshProgress(int64_t refresh_progress);
71+
6072
private:
6173
LoadState state_{LoadState::LOAD_STATE_NOT_EXIST};
6274
int64_t progress_{0};
75+
int64_t refresh_progress_{0};
6376
};
6477

6578
} // namespace milvus

0 commit comments

Comments
 (0)