Skip to content

Commit 86f5cc3

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

24 files changed

Lines changed: 331 additions & 72 deletions

‎.github/workflows/main.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,7 @@ jobs:
153153
sh scripts/install_deps.sh
154154
- name: Testing With Coverage
155155
run: |
156-
make coverage
156+
JOBS=4 CMAKE_BUILD_EXAMPLES=OFF make coverage
157157
- name: Upload coverage
158158
uses: codecov/codecov-action@v5
159159
with:

‎CMakeLists.txt‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ set_option_category("Build")
5151

5252
define_option(MILVUS_BUILD_TEST "Build with testing" OFF)
5353
define_option(MILVUS_BUILD_COVERAGE "Build with coverage" OFF)
54+
define_option(MILVUS_BUILD_EXAMPLES "Build examples" ON)
5455

5556
if (DEFINED MILVUS_SDK_VERSION)
5657
message("MILVUS_SDK_VERSION = ${MILVUS_SDK_VERSION}")
@@ -133,6 +134,9 @@ add_subdirectory(src)
133134

134135
if (MILVUS_BUILD_TEST)
135136
add_subdirectory(test)
137+
endif ()
138+
139+
if (MILVUS_BUILD_EXAMPLES)
136140
add_subdirectory(examples)
137141
endif ()
138142

‎scripts/build.sh‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ BUILD_SHARED_LIBS=${BUILD_SHARED_LIBS:-ON}
3232
BUILD_FROM_CONAN="ON"
3333

3434

35-
JOBS="$(nproc 2>/dev/null || sysctl -n hw.logicalcpu 2>/dev/null || echo 3)"
35+
JOBS="${JOBS:-$(nproc 2>/dev/null || sysctl -n hw.logicalcpu 2>/dev/null || echo 3)}"
3636
if [ ${JOBS} -lt 3 ] ; then
3737
JOBS=3
3838
fi
@@ -216,11 +216,14 @@ if [[ -n "${MILVUS_SDK_VERSION}" ]]; then
216216
CMAKE_VERSION_ARG="-DMILVUS_SDK_VERSION=${MILVUS_SDK_VERSION}"
217217
fi
218218

219+
CMAKE_BUILD_EXAMPLES=${CMAKE_BUILD_EXAMPLES:-ON}
220+
219221
CMAKE_CMD="cmake \
220222
-DCMAKE_TOOLCHAIN_FILE=${TOOLCHAIN_FILE} \
221223
-DCMAKE_BUILD_TYPE=${BUILD_TYPE} \
222224
-DMILVUS_BUILD_TEST=${BUILD_TEST} \
223225
-DMILVUS_BUILD_COVERAGE=${BUILD_COVERAGE} \
226+
-DMILVUS_BUILD_EXAMPLES=${CMAKE_BUILD_EXAMPLES} \
224227
${CMAKE_VERSION_ARG} \
225228
-DCMAKE_INSTALL_PREFIX=${CMAKE_INSTALL_PREFIX} \
226229
-DBUILD_SHARED_LIBS=${BUILD_SHARED_LIBS} \

‎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: 19 additions & 10 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;
590-
auto status = connection_.GetLoadingProgress(request.DatabaseName(), request.CollectionName(),
591-
request.PartitionNames(), progress, rpc_timeout_ms);
592+
uint32_t refresh_progress = 0;
593+
auto status =
594+
connection_.GetLoadingProgress(request.DatabaseName(), request.CollectionName(),
595+
request.PartitionNames(), progress, refresh_progress, 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
}
@@ -1919,8 +1924,8 @@ MilvusClientV2Impl::FlushAll(const FlushAllRequest& request, FlushAllResponse& r
19191924

19201925
auto wait_for_status = [this, &request, &progress_monitor](const proto::milvus::FlushAllResponse& rpc_response) {
19211926
GetFlushAllStateRequest state_request = GetFlushAllStateRequest()
1922-
.WithDatabaseName(request.DatabaseName())
1923-
.WithFlushAllTs(rpc_response.flush_all_ts());
1927+
.WithDatabaseName(request.DatabaseName())
1928+
.WithFlushAllTs(rpc_response.flush_all_ts());
19241929
return ConnectionHandler::WaitForStatus(
19251930
[this, &state_request](Progress& p) -> Status {
19261931
p.total_ = 1;
@@ -1951,7 +1956,7 @@ MilvusClientV2Impl::GetFlushAllState(const GetFlushAllStateRequest& request, Get
19511956

19521957
Status
19531958
MilvusClientV2Impl::getFlushAllState(const GetFlushAllStateRequest& request, GetFlushAllStateResponse& response,
1954-
uint64_t rpc_timeout_ms) {
1959+
uint64_t rpc_timeout_ms) {
19551960
auto pre = [&request](proto::milvus::GetFlushAllStateRequest& rpc_request) {
19561961
rpc_request.set_db_name(request.DatabaseName());
19571962
rpc_request.set_flush_all_ts(request.FlushAllTs());
@@ -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();
@@ -2433,7 +2442,7 @@ MilvusClientV2Impl::GetCompactionPlans(const GetCompactionPlansRequest& request,
24332442

24342443
Status
24352444
MilvusClientV2Impl::GetReplicateConfiguration(const GetReplicateConfigurationRequest& request,
2436-
GetReplicateConfigurationResponse& response) {
2445+
GetReplicateConfigurationResponse& response) {
24372446
auto post = [&response](const proto::milvus::GetReplicateConfigurationResponse& rpc_response) {
24382447
ReplicateConfiguration configuration;
24392448
ConvertReplicateConfiguration(rpc_response.configuration(), configuration);

‎src/impl/MilvusConnection.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -678,8 +678,8 @@ MilvusConnection::UpdateReplicateConfiguration(const proto::milvus::UpdateReplic
678678

679679
Status
680680
MilvusConnection::GetReplicateInfo(const proto::milvus::GetReplicateInfoRequest& request,
681-
proto::milvus::GetReplicateInfoResponse& response,
682-
const GrpcContextOptions& options) {
681+
proto::milvus::GetReplicateInfoResponse& response,
682+
const GrpcContextOptions& options) {
683683
return grpcCall("GetReplicateInfo", &Stub::GetReplicateInfo, request, response, options);
684684
}
685685

‎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

0 commit comments

Comments
 (0)