Skip to content

Commit 0832c5b

Browse files
sparknacksparknack
committed
refactor GetChunkIDsByOffsets in ChunkedColumnGroup
Signed-off-by: Shawn Wang <shawn.wang@zilliz.com>
1 parent 832c609 commit 0832c5b

4 files changed

Lines changed: 109 additions & 40 deletions

File tree

‎internal/core/src/mmap/ChunkedColumn.h‎

Lines changed: 41 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -44,16 +44,22 @@ namespace milvus {
4444
using namespace milvus::cachinglayer;
4545

4646
std::pair<size_t, size_t> inline GetChunkIDByOffset(
47-
int64_t offset, const std::vector<int64_t>& num_rows_until_chunk) {
47+
int64_t offset,
48+
const std::vector<int64_t>& num_rows_until_chunk,
49+
int64_t virt_chunk_order,
50+
const std::vector<int64_t>& vcid_to_cid_arr) {
4851
// optimize for single chunk case
4952
if (num_rows_until_chunk.size() == 2) {
5053
return {0, offset};
5154
}
52-
auto iter = std::lower_bound(
53-
num_rows_until_chunk.begin(), num_rows_until_chunk.end(), offset + 1);
54-
size_t chunk_idx = std::distance(num_rows_until_chunk.begin(), iter) - 1;
55-
size_t offset_in_chunk = offset - num_rows_until_chunk[chunk_idx];
56-
return {chunk_idx, offset_in_chunk};
55+
auto vcid = offset >> virt_chunk_order;
56+
auto scid = vcid_to_cid_arr[vcid];
57+
while (scid < num_rows_until_chunk.size() - 1 &&
58+
offset >= num_rows_until_chunk[scid + 1]) {
59+
scid++;
60+
}
61+
auto offset_in_chunk = offset - num_rows_until_chunk[scid];
62+
return {scid, offset_in_chunk};
5763
}
5864

5965
std::pair<std::vector<milvus::cachinglayer::cid_t>,
@@ -104,7 +110,8 @@ std::pair<std::vector<milvus::cachinglayer::cid_t>,
104110
auto offset = offsets[i];
105111
auto vcid = offset >> virt_chunk_order;
106112
auto scid = vcid_to_cid_arr[vcid];
107-
while (scid < num_rows_until_chunk.size() - 1 && offset >= num_rows_until_chunk[scid + 1]) {
113+
while (scid < num_rows_until_chunk.size() - 1 &&
114+
offset >= num_rows_until_chunk[scid + 1]) {
108115
scid++;
109116
}
110117
auto offset_in_chunk = offset - num_rows_until_chunk[scid];
@@ -232,9 +239,7 @@ class ChunkedColumnBase : public ChunkedColumnInterface {
232239
}
233240

234241
void
235-
BulkValueAt(void* dst,
236-
const int64_t* offsets,
237-
int64_t count) override {
242+
BulkValueAt(void* dst, const int64_t* offsets, int64_t count) override {
238243
PanicInfo(ErrorCode::Unsupported,
239244
"BulkValueAt only supported for ChunkedColumn");
240245
}
@@ -284,18 +289,27 @@ class ChunkedColumnBase : public ChunkedColumnInterface {
284289
// "offset {} is out of range, num_rows: {}",
285290
// offset,
286291
// num_rows_);
287-
auto& num_rows_until_chunk = GetNumRowsUntilChunk();
288-
return ::milvus::GetChunkIDByOffset(offset, num_rows_until_chunk);
292+
auto meta = static_cast<milvus::segcore::storagev1translator::CTMeta*>(
293+
slot_->meta());
294+
auto& num_rows_until_chunk = meta->num_rows_until_chunk_;
295+
auto& virt_chunk_order = meta->virt_chunk_order_;
296+
auto& vcid_to_cid_arr = meta->vcid_to_cid_arr_;
297+
return ::milvus::GetChunkIDByOffset(
298+
offset, num_rows_until_chunk, virt_chunk_order, vcid_to_cid_arr);
289299
}
290300

291301
std::pair<std::vector<milvus::cachinglayer::cid_t>, std::vector<int64_t>>
292302
GetChunkIDsByOffsets(const int64_t* offsets, int64_t count) const override {
293-
auto& num_rows_until_chunk = GetNumRowsUntilChunk();
294-
auto meta = static_cast<milvus::segcore::storagev1translator::CTMeta*>(slot_->meta());
303+
auto meta = static_cast<milvus::segcore::storagev1translator::CTMeta*>(
304+
slot_->meta());
305+
auto& num_rows_until_chunk = meta->num_rows_until_chunk_;
295306
auto& virt_chunk_order = meta->virt_chunk_order_;
296307
auto& vcid_to_cid_arr = meta->vcid_to_cid_arr_;
297-
return ::milvus::GetChunkIDsByOffsets(
298-
offsets, count, num_rows_until_chunk, virt_chunk_order, vcid_to_cid_arr);
308+
return ::milvus::GetChunkIDsByOffsets(offsets,
309+
count,
310+
num_rows_until_chunk,
311+
virt_chunk_order,
312+
vcid_to_cid_arr);
299313
}
300314

301315
PinWrapper<Chunk*>
@@ -352,17 +366,15 @@ class ChunkedColumn : public ChunkedColumnBase {
352366
}
353367

354368
void
355-
BulkValueAt(void* dst,
356-
const int64_t* offsets,
357-
int64_t count) override {
369+
BulkValueAt(void* dst, const int64_t* offsets, int64_t count) override {
358370
auto [cids, offsets_in_chunk] = ToChunkIdAndOffset(offsets, count);
359371
auto ca = slot_->PinCells(cids);
360372
auto typed_dst = static_cast<T*>(dst);
361373
for (int64_t i = 0; i < count; i++) {
362374
auto chunk = ca->get_ith_cell(cids[i]);
363375
auto value = chunk->ValueAt(offsets_in_chunk[i]);
364-
typed_dst[i] = *static_cast<const T*>(
365-
static_cast<const void*>(value));
376+
typed_dst[i] =
377+
*static_cast<const T*>(static_cast<const void*>(value));
366378
}
367379
}
368380

@@ -404,9 +416,7 @@ class ChunkedColumn<void> : public ChunkedColumnBase {
404416
}
405417

406418
void
407-
BulkValueAt(void* dst,
408-
const int64_t* offsets,
409-
int64_t count) override {
419+
BulkValueAt(void* dst, const int64_t* offsets, int64_t count) override {
410420
PanicInfo(ErrorCode::Unsupported,
411421
"BulkValueAt is not supported for ChunkedColumn<void>");
412422
}
@@ -636,35 +646,35 @@ MakeChunkedColumnBase(DataType data_type,
636646
case DataType::INT8:
637647
return std::static_pointer_cast<ChunkedColumnInterface>(
638648
std::make_shared<ChunkedColumn<int8_t>>(std::move(translator),
639-
field_meta));
649+
field_meta));
640650
case DataType::INT16:
641651
return std::static_pointer_cast<ChunkedColumnInterface>(
642652
std::make_shared<ChunkedColumn<int16_t>>(std::move(translator),
643-
field_meta));
653+
field_meta));
644654
case DataType::INT32:
645655
return std::static_pointer_cast<ChunkedColumnInterface>(
646656
std::make_shared<ChunkedColumn<int32_t>>(std::move(translator),
647-
field_meta));
657+
field_meta));
648658
case DataType::INT64:
649659
return std::static_pointer_cast<ChunkedColumnInterface>(
650660
std::make_shared<ChunkedColumn<int64_t>>(std::move(translator),
651-
field_meta));
661+
field_meta));
652662
case DataType::FLOAT:
653663
return std::static_pointer_cast<ChunkedColumnInterface>(
654664
std::make_shared<ChunkedColumn<float>>(std::move(translator),
655665
field_meta));
656666
case DataType::DOUBLE:
657667
return std::static_pointer_cast<ChunkedColumnInterface>(
658668
std::make_shared<ChunkedColumn<double>>(std::move(translator),
659-
field_meta));
669+
field_meta));
660670
case DataType::BOOL:
661671
return std::static_pointer_cast<ChunkedColumnInterface>(
662672
std::make_shared<ChunkedColumn<bool>>(std::move(translator),
663-
field_meta));
673+
field_meta));
664674
default:
665675
return std::static_pointer_cast<ChunkedColumnInterface>(
666676
std::make_shared<ChunkedColumn<void>>(std::move(translator),
667-
field_meta));
677+
field_meta));
668678
}
669679
}
670680

‎internal/core/src/mmap/ChunkedColumnGroup.h‎

Lines changed: 61 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,60 @@ class ChunkedColumnGroup {
9999
return meta->num_rows_until_chunk_;
100100
}
101101

102+
std::pair<size_t, size_t>
103+
GetChunkIDByOffset(int64_t offset) const {
104+
auto meta =
105+
static_cast<milvus::segcore::storagev2translator::GroupCTMeta*>(
106+
slot_->meta());
107+
auto& num_rows_until_chunk = meta->num_rows_until_chunk_;
108+
// optimize for single chunk case
109+
if (num_rows_until_chunk.size() == 2) {
110+
return {0, offset};
111+
}
112+
auto& virt_chunk_order = meta->virt_chunk_order_;
113+
auto& vcid_to_cid_arr = meta->vcid_to_cid_arr_;
114+
auto vcid = offset >> virt_chunk_order;
115+
auto scid = vcid_to_cid_arr[vcid];
116+
while (scid < num_rows_until_chunk.size() - 1 &&
117+
offset >= num_rows_until_chunk[scid + 1]) {
118+
scid++;
119+
}
120+
auto offset_in_chunk = offset - num_rows_until_chunk[scid];
121+
return {scid, offset_in_chunk};
122+
}
123+
124+
std::pair<std::vector<milvus::cachinglayer::cid_t>, std::vector<int64_t>>
125+
GetChunkIDsByOffsets(const int64_t* offsets, int64_t count) const {
126+
std::vector<milvus::cachinglayer::cid_t> cids(count);
127+
std::vector<int64_t> offsets_in_chunk(count);
128+
auto meta =
129+
static_cast<milvus::segcore::storagev2translator::GroupCTMeta*>(
130+
slot_->meta());
131+
auto& num_rows_until_chunk = meta->num_rows_until_chunk_;
132+
// optimize for single chunk case
133+
if (num_rows_until_chunk.size() == 2) {
134+
for (int64_t i = 0; i < count; i++) {
135+
offsets_in_chunk[i] = offsets[i];
136+
}
137+
return std::make_pair(std::move(cids), std::move(offsets_in_chunk));
138+
}
139+
auto& virt_chunk_order = meta->virt_chunk_order_;
140+
auto& vcid_to_cid_arr = meta->vcid_to_cid_arr_;
141+
for (int64_t i = 0; i < count; i++) {
142+
auto offset = offsets[i];
143+
auto vcid = offset >> virt_chunk_order;
144+
auto scid = vcid_to_cid_arr[vcid];
145+
while (scid < num_rows_until_chunk.size() - 1 &&
146+
offset >= num_rows_until_chunk[scid + 1]) {
147+
scid++;
148+
}
149+
auto offset_in_chunk = offset - num_rows_until_chunk[scid];
150+
cids[i] = scid;
151+
offsets_in_chunk[i] = offset_in_chunk;
152+
}
153+
return std::make_pair(std::move(cids), std::move(offsets_in_chunk));
154+
}
155+
102156
size_t
103157
NumFieldsInGroup() const {
104158
auto meta =
@@ -293,15 +347,13 @@ class ProxyChunkColumn : public ChunkedColumnInterface {
293347

294348
std::pair<size_t, size_t>
295349
GetChunkIDByOffset(int64_t offset) const override {
296-
int64_t current_offset = 0;
297-
for (int64_t i = 0; i < num_chunks(); ++i) {
298-
auto rows = chunk_row_nums(i);
299-
if (current_offset + rows > offset) {
300-
return {i, offset - current_offset};
301-
}
302-
current_offset += rows;
303-
}
304-
return {num_chunks() - 1, chunk_row_nums(num_chunks() - 1) - 1};
350+
auto meta = static_cast<milvus::segcore::storagev2translator::GroupCTMeta*>(
351+
slot_->meta());
352+
auto& num_rows_until_chunk = meta->num_rows_until_chunk_;
353+
auto& virt_chunk_order = meta->virt_chunk_order_;
354+
auto& vcid_to_cid_arr = meta->vcid_to_cid_arr_;
355+
return GetChunkIDByOffset(offset, num_rows_until_chunk,
356+
virt_chunk_order, vcid_to_cid_arr);
305357
}
306358

307359
std::pair<std::vector<milvus::cachinglayer::cid_t>, std::vector<int64_t>>

‎internal/core/src/segcore/storagev2translator/GroupCTMeta.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@ namespace milvus::segcore::storagev2translator {
2222
struct GroupCTMeta : public milvus::cachinglayer::Meta {
2323
std::vector<int64_t> num_rows_until_chunk_;
2424
std::vector<int64_t> chunk_memory_size_;
25+
int64_t virt_chunk_order_;
26+
std::vector<int64_t> vcid_to_cid_arr_;
2527
size_t num_fields_;
2628
GroupCTMeta(size_t num_fields,
2729
milvus::cachinglayer::StorageType storage_type,

‎internal/core/src/segcore/storagev2translator/GroupChunkTranslator.cpp‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,11 @@ GroupChunkTranslator::GroupChunkTranslator(
8282
column_group_info_.field_id,
8383
meta_.num_rows_until_chunk_.back(),
8484
column_group_info_.row_count));
85+
virtual_chunk_config(column_group_info_.row_count,
86+
row_group_meta_list_.size(),
87+
meta_.num_rows_until_chunk_,
88+
meta_.virt_chunk_order_,
89+
meta_.vcid_to_cid_arr_);
8590
}
8691

8792
GroupChunkTranslator::~GroupChunkTranslator() {

0 commit comments

Comments
 (0)