Skip to content

Commit 54fc6a7

Browse files
yhmoyhmo
authored
Align GetLoadState/LoadCollection/LoadPartitions with PyMilvus (#504)
Signed-off-by: yhmo <yihua.mo@zilliz.com>
1 parent 01b246e commit 54fc6a7

21 files changed

Lines changed: 336 additions & 83 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: 12 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,22 @@ 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+
case "${CMAKE_BUILD_EXAMPLES}" in
221+
ON|OFF)
222+
;;
223+
*)
224+
echo "ERROR! CMAKE_BUILD_EXAMPLES must be ON or OFF"
225+
exit 1
226+
;;
227+
esac
228+
219229
CMAKE_CMD="cmake \
220230
-DCMAKE_TOOLCHAIN_FILE=${TOOLCHAIN_FILE} \
221231
-DCMAKE_BUILD_TYPE=${BUILD_TYPE} \
222232
-DMILVUS_BUILD_TEST=${BUILD_TEST} \
223233
-DMILVUS_BUILD_COVERAGE=${BUILD_COVERAGE} \
234+
-DMILVUS_BUILD_EXAMPLES=${CMAKE_BUILD_EXAMPLES} \
224235
${CMAKE_VERSION_ARG} \
225236
-DCMAKE_INSTALL_PREFIX=${CMAKE_INSTALL_PREFIX} \
226237
-DBUILD_SHARED_LIBS=${BUILD_SHARED_LIBS} \

‎src/impl/MilvusClientImpl.cpp‎

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,9 @@ MilvusClientImpl::LoadCollection(const std::string& collection_name, int replica
174174
progress.total_ = 100;
175175
auto db_name = connection_.CurrentDbName("");
176176
std::set<std::string> partition_names;
177-
return connection_.GetLoadingProgress(db_name, collection_name, partition_names, progress.finished_);
177+
uint32_t refresh_progress = 0;
178+
return connection_.GetLoadingProgress(db_name, collection_name, partition_names, progress.finished_,
179+
refresh_progress);
178180
},
179181
progress_monitor);
180182
};
@@ -411,8 +413,9 @@ MilvusClientImpl::LoadPartitions(const std::string& collection_name, const std::
411413
for (const auto& name : partition_names) {
412414
unique_partition_names.insert(name);
413415
}
416+
uint32_t refresh_progress = 0;
414417
return connection_.GetLoadingProgress(db_name, collection_name, unique_partition_names,
415-
progress.finished_);
418+
progress.finished_, refresh_progress);
416419
},
417420
progress_monitor);
418421
};
@@ -1481,10 +1484,11 @@ MilvusClientImpl::GetPersistentSegmentInfo(const std::string& collection_name, S
14811484
return Status::OK();
14821485
};
14831486

1484-
auto post = [&segments_info](const proto::milvus::GetPersistentSegmentInfoResponse& response) {
1487+
auto post = [&collection_name, &segments_info](const proto::milvus::GetPersistentSegmentInfoResponse& response) {
14851488
for (const auto& info : response.infos()) {
14861489
segments_info.emplace_back(info.collectionid(), info.partitionid(), info.segmentid(), info.num_rows(),
1487-
SegmentStateCast(info.state()));
1490+
SegmentStateCast(info.state()), collection_name, SegmentLevelCast(info.level()),
1491+
info.storage_version(), info.is_sorted());
14881492
}
14891493
return Status::OK();
14901494
};
@@ -1501,14 +1505,17 @@ MilvusClientImpl::GetQuerySegmentInfo(const std::string& collection_name, QueryS
15011505
return Status::OK();
15021506
};
15031507

1504-
auto post = [&segments_info](const proto::milvus::GetQuerySegmentInfoResponse& response) {
1508+
auto post = [&collection_name, &segments_info](const proto::milvus::GetQuerySegmentInfoResponse& response) {
15051509
for (const auto& info : response.infos()) {
15061510
std::vector<int64_t> ids;
1511+
ids.reserve(info.nodeids_size());
15071512
for (auto id : info.nodeids()) {
15081513
ids.push_back(id);
15091514
}
15101515
segments_info.emplace_back(info.collectionid(), info.partitionid(), info.segmentid(), info.num_rows(),
1511-
milvus::SegmentStateCast(info.state()), info.index_name(), info.indexid(), ids);
1516+
milvus::SegmentStateCast(info.state()), info.index_name(), info.indexid(), ids,
1517+
collection_name, info.mem_size(), SegmentLevelCast(info.level()),
1518+
info.storage_version(), info.is_sorted());
15121519
}
15131520
return Status::OK();
15141521
};

‎src/impl/MilvusClientV2Impl.cpp‎

Lines changed: 44 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -312,8 +312,15 @@ MilvusClientV2Impl::LoadCollection(const LoadCollectionRequest& request) {
312312
progress.total_ = 100;
313313
auto db_name = connection_.CurrentDbName(request.DatabaseName());
314314
std::set<std::string> partition_names;
315-
return connection_.GetLoadingProgress(db_name, request.CollectionName(), partition_names,
316-
progress.finished_);
315+
uint32_t loading_progress = 0;
316+
uint32_t refresh_progress = 0;
317+
auto status = connection_.GetLoadingProgress(db_name, request.CollectionName(), partition_names,
318+
loading_progress, refresh_progress);
319+
if (!status.IsOk()) {
320+
return status;
321+
}
322+
progress.finished_ = request.Refresh() ? refresh_progress : loading_progress;
323+
return Status::OK();
317324
},
318325
progress_monitor);
319326
};
@@ -350,8 +357,15 @@ MilvusClientV2Impl::refreshLoad(const RefreshLoadRequest& request, uint64_t rpc_
350357
progress.total_ = 100;
351358
auto db_name = connection_.CurrentDbName(request.DatabaseName());
352359
std::set<std::string> partition_names;
353-
return connection_.GetLoadingProgress(db_name, request.CollectionName(), partition_names,
354-
progress.finished_, rpc_timeout_ms);
360+
uint32_t loading_progress = 0;
361+
uint32_t refresh_progress = 0;
362+
auto status = connection_.GetLoadingProgress(db_name, request.CollectionName(), partition_names,
363+
loading_progress, refresh_progress, rpc_timeout_ms);
364+
if (!status.IsOk()) {
365+
return status;
366+
}
367+
progress.finished_ = refresh_progress;
368+
return Status::OK();
355369
},
356370
progress_monitor);
357371
};
@@ -585,10 +599,13 @@ MilvusClientV2Impl::getLoadState(const GetLoadStateRequest& request, GetLoadStat
585599
auto state = rpc_response.state();
586600
response.SetState(LoadStateCast(state));
587601

602+
response.SetProgress(0);
588603
if (state == proto::common::LoadState::LoadStateLoading) {
589604
uint32_t progress = 0;
590-
auto status = connection_.GetLoadingProgress(request.DatabaseName(), request.CollectionName(),
591-
request.PartitionNames(), progress, rpc_timeout_ms);
605+
uint32_t refresh_progress = 0;
606+
auto status =
607+
connection_.GetLoadingProgress(request.DatabaseName(), request.CollectionName(),
608+
request.PartitionNames(), progress, refresh_progress, rpc_timeout_ms);
592609
if (!status.IsOk()) {
593610
return status;
594611
}
@@ -830,8 +847,15 @@ MilvusClientV2Impl::LoadPartitions(const LoadPartitionsRequest& request) {
830847
[&request, this](Progress& progress) -> Status {
831848
progress.total_ = 100;
832849
auto db_name = connection_.CurrentDbName(request.DatabaseName());
833-
return connection_.GetLoadingProgress(db_name, request.CollectionName(), request.PartitionNames(),
834-
progress.finished_);
850+
uint32_t loading_progress = 0;
851+
uint32_t refresh_progress = 0;
852+
auto status = connection_.GetLoadingProgress(db_name, request.CollectionName(), request.PartitionNames(),
853+
loading_progress, refresh_progress);
854+
if (!status.IsOk()) {
855+
return status;
856+
}
857+
progress.finished_ = request.Refresh() ? refresh_progress : loading_progress;
858+
return Status::OK();
835859
},
836860
progress_monitor);
837861
};
@@ -1919,8 +1943,8 @@ MilvusClientV2Impl::FlushAll(const FlushAllRequest& request, FlushAllResponse& r
19191943

19201944
auto wait_for_status = [this, &request, &progress_monitor](const proto::milvus::FlushAllResponse& rpc_response) {
19211945
GetFlushAllStateRequest state_request = GetFlushAllStateRequest()
1922-
.WithDatabaseName(request.DatabaseName())
1923-
.WithFlushAllTs(rpc_response.flush_all_ts());
1946+
.WithDatabaseName(request.DatabaseName())
1947+
.WithFlushAllTs(rpc_response.flush_all_ts());
19241948
return ConnectionHandler::WaitForStatus(
19251949
[this, &state_request](Progress& p) -> Status {
19261950
p.total_ = 1;
@@ -1951,7 +1975,7 @@ MilvusClientV2Impl::GetFlushAllState(const GetFlushAllStateRequest& request, Get
19511975

19521976
Status
19531977
MilvusClientV2Impl::getFlushAllState(const GetFlushAllStateRequest& request, GetFlushAllStateResponse& response,
1954-
uint64_t rpc_timeout_ms) {
1978+
uint64_t rpc_timeout_ms) {
19551979
auto pre = [&request](proto::milvus::GetFlushAllStateRequest& rpc_request) {
19561980
rpc_request.set_db_name(request.DatabaseName());
19571981
rpc_request.set_flush_all_ts(request.FlushAllTs());
@@ -1977,12 +2001,13 @@ MilvusClientV2Impl::ListPersistentSegments(const ListPersistentSegmentsRequest&
19772001
return Status::OK();
19782002
};
19792003

1980-
auto post = [&response](const proto::milvus::GetPersistentSegmentInfoResponse& rpc_response) {
2004+
auto post = [&request, &response](const proto::milvus::GetPersistentSegmentInfoResponse& rpc_response) {
19812005
SegmentsInfo segments_info;
19822006
segments_info.reserve(rpc_response.infos_size());
19832007
for (const auto& info : rpc_response.infos()) {
19842008
segments_info.emplace_back(info.collectionid(), info.partitionid(), info.segmentid(), info.num_rows(),
1985-
SegmentStateCast(info.state()));
2009+
SegmentStateCast(info.state()), request.CollectionName(),
2010+
SegmentLevelCast(info.level()), info.storage_version(), info.is_sorted());
19862011
}
19872012
response.SetResult(std::move(segments_info));
19882013
return Status::OK();
@@ -2001,16 +2026,19 @@ MilvusClientV2Impl::ListQuerySegments(const ListQuerySegmentsRequest& request, L
20012026
return Status::OK();
20022027
};
20032028

2004-
auto post = [&response](const proto::milvus::GetQuerySegmentInfoResponse& rpc_response) {
2029+
auto post = [&request, &response](const proto::milvus::GetQuerySegmentInfoResponse& rpc_response) {
20052030
QuerySegmentsInfo segments_info;
20062031
segments_info.reserve(rpc_response.infos_size());
20072032
for (const auto& info : rpc_response.infos()) {
20082033
std::vector<int64_t> ids;
2034+
ids.reserve(info.nodeids_size());
20092035
for (auto id : info.nodeids()) {
20102036
ids.push_back(id);
20112037
}
20122038
segments_info.emplace_back(info.collectionid(), info.partitionid(), info.segmentid(), info.num_rows(),
2013-
milvus::SegmentStateCast(info.state()), info.index_name(), info.indexid(), ids);
2039+
milvus::SegmentStateCast(info.state()), info.index_name(), info.indexid(), ids,
2040+
request.CollectionName(), info.mem_size(), SegmentLevelCast(info.level()),
2041+
info.storage_version(), info.is_sorted());
20142042
}
20152043
response.SetResult(std::move(segments_info));
20162044
return Status::OK();
@@ -2433,7 +2461,7 @@ MilvusClientV2Impl::GetCompactionPlans(const GetCompactionPlansRequest& request,
24332461

24342462
Status
24352463
MilvusClientV2Impl::GetReplicateConfiguration(const GetReplicateConfigurationRequest& request,
2436-
GetReplicateConfigurationResponse& response) {
2464+
GetReplicateConfigurationResponse& response) {
24372465
auto post = [&response](const proto::milvus::GetReplicateConfigurationResponse& rpc_response) {
24382466
ReplicateConfiguration configuration;
24392467
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/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: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -109,8 +109,8 @@ ConnectionHandler::CurrentDbName(const std::string& overwrite_db_name) const {
109109

110110
Status
111111
ConnectionHandler::GetLoadingProgress(const std::string& db_name, const std::string& collection_name,
112-
const std::set<std::string> partition_names, uint32_t& progress,
113-
uint64_t rpc_timeout_ms) {
112+
const std::set<std::string>& partition_names, uint32_t& progress,
113+
uint32_t& refresh_progress, uint64_t rpc_timeout_ms) {
114114
MilvusConnectionPtr connection;
115115
uint64_t timeout = 0;
116116
{
@@ -138,6 +138,7 @@ ConnectionHandler::GetLoadingProgress(const std::string& db_name, const std::str
138138
return status;
139139
}
140140
progress = static_cast<uint32_t>(progress_resp.progress());
141+
refresh_progress = static_cast<uint32_t>(progress_resp.refresh_progress());
141142
return Status::OK();
142143
}
143144

‎src/impl/utils/ConnectionHandler.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,8 @@ class ConnectionHandler {
6565
// This interface is not exposed to users
6666
Status
6767
GetLoadingProgress(const std::string& db_name, const std::string& collection_name,
68-
const std::set<std::string> partition_names, uint32_t& progress, uint64_t rpc_timeout_ms = 0);
68+
const std::set<std::string>& partition_names, uint32_t& progress,
69+
uint32_t& refresh_progress, uint64_t rpc_timeout_ms = 0);
6970

7071
/**
7172
* Internal wait for status query done.

0 commit comments

Comments
 (0)