furyhawk/milvus
1
0
Fork 0
mirror of https://github.com/milvus-io/milvus.git synced 2026-07-21 10:15:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

enhance: optimize memory usage for ArrayOffsets (#49093)

Browse Source 
issue: https://github.com/milvus-io/milvus/issues/42148

This PR remove `element_row_ids_` in ArrayOffsets to reduce memory usage
whose memory is O(total element count). Instead, use binary search on
`row_to_element_start_` to calculate the element id.

Signed-off-by: SpadeA <tangchenjie1210@gmail.com>
This commit is contained in:
Spade A
2026-04-23 15:57:44 +08:00
committed by GitHub
parent f79085e5ec
commit 568446dfec
6 changed files with 381 additions and 157 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/core/src/common/ArrayOffsets.cpp
+78 -88
View File
@@ -38,8 +38,13 @@ std::pair<int32_t, int32_t>
ArrayOffsetsSealed::ElementIDToRowID(int32_t elem_id) const {
assert(elem_id >= 0 && elem_id < GetTotalElementCount());
int32_t row_id = element_row_ids_[elem_id];
// Compute elem_idx: elem_idx = elem_id - start_of_this_row
// Binary search: find the row where elem_id belongs
// row_to_element_start_[row_id] <= elem_id < row_to_element_start_[row_id + 1]
auto it = std::upper_bound(
row_to_element_start_.begin(), row_to_element_start_.end(), elem_id);
int32_t row_id = static_cast<int32_t>(
std::distance(row_to_element_start_.begin(), it) - 1);
int32_t elem_idx = elem_id - row_to_element_start_[row_id];
return {row_id, elem_idx};
}
@@ -93,12 +98,13 @@ FixedVector<int32_t>
ArrayOffsetsSealed::RowBitsetToElementOffsets(
const TargetBitmapView& row_bitset, int64_t row_start) const {
int64_t row_count = row_bitset.size();
AssertInfo(row_start >= 0 && row_start + row_count <= GetRowCount(),
int64_t total_rows = GetRowCount();
AssertInfo(row_start >= 0 && row_start + row_count <= total_rows,
"row range out of bounds: row_start={}, row_count={}, "
"total_rows={}",
row_start,
row_count,
GetRowCount());
total_rows);
int64_t selected_rows = row_bitset.count();
FixedVector<int32_t> element_offsets;
@@ -106,9 +112,7 @@ ArrayOffsetsSealed::RowBitsetToElementOffsets(
return element_offsets;
}
int64_t avg_elem_per_row =
static_cast<int64_t>(element_row_ids_.size()) /
(static_cast<int64_t>(row_to_element_start_.size()) - 1);
int64_t avg_elem_per_row = GetTotalElementCount() / total_rows;
element_offsets.reserve(selected_rows * avg_elem_per_row);
@@ -135,10 +139,8 @@ ArrayOffsetsSealed::RowOffsetsToElementOffsets(
}
int32_t row_count = GetRowCount();
int64_t avg_elem_per_row =
(row_count > 0)
? static_cast<int64_t>(element_row_ids_.size()) / row_count
: 1;
int64_t avg_elem_per_row = GetTotalElementCount() / row_count;
element_offsets.reserve(row_offsets.size() * avg_elem_per_row);
for (auto row_id : row_offsets) {
assert(row_id >= 0 && row_id < row_count);
@@ -187,24 +189,21 @@ ArrayOffsetsSealed::BuildFromSegment(const void* segment,
"ArrayOffsetsSealed::BuildFromSegment: empty segment for struct "
"'{}'",
field_meta.get_name().get());
return std::make_shared<ArrayOffsetsSealed>(std::vector<int32_t>{},
std::vector<int32_t>{0});
return std::make_shared<ArrayOffsetsSealed>(std::vector<int32_t>{0});
}
FieldId field_id = field_meta.get_id();
auto data_type = field_meta.get_data_type();
std::vector<int32_t> element_row_ids;
// Size is row_count + 1, last element stores total_element_count
std::vector<int32_t> row_to_element_start(row_count + 1);
// Pre-reserve element_row_ids assuming average 4 elements per row
element_row_ids.reserve(row_count * 4);
auto temp_op_ctx = std::make_unique<OpContext>();
auto op_ctx_ptr = temp_op_ctx.get();
int64_t num_chunks = seg->num_chunk(field_id);
int32_t current_row_id = 0;
int32_t total_elements = 0;
if (data_type == DataType::VECTOR_ARRAY) {
for (int64_t chunk_id = 0; chunk_id < num_chunks; ++chunk_id) {
@@ -218,14 +217,8 @@ ArrayOffsetsSealed::BuildFromSegment(const void* segment,
array_len = vector_array_views[i].length();
}
// Record the start position for this row
row_to_element_start[current_row_id] = element_row_ids.size();
// Add row_id for each element (elem_idx computed on access)
for (int32_t j = 0; j < array_len; ++j) {
element_row_ids.emplace_back(current_row_id);
}
row_to_element_start[current_row_id] = total_elements;
total_elements += array_len;
current_row_id++;
}
}
@@ -241,29 +234,21 @@ ArrayOffsetsSealed::BuildFromSegment(const void* segment,
array_len = array_views[i].length();
}
// Record the start position for this row
row_to_element_start[current_row_id] = element_row_ids.size();
// Add row_id for each element (elem_idx computed on access)
for (int32_t j = 0; j < array_len; ++j) {
element_row_ids.emplace_back(current_row_id);
}
row_to_element_start[current_row_id] = total_elements;
total_elements += array_len;
current_row_id++;
}
}
}
// Store total element count as the last entry
row_to_element_start[row_count] = element_row_ids.size();
row_to_element_start[row_count] = total_elements;
AssertInfo(current_row_id == row_count,
"Row count mismatch: expected {}, got {}",
row_count,
current_row_id);
int64_t total_elements = element_row_ids.size();
LOG_INFO(
"ArrayOffsetsSealed::BuildFromSegment: struct_name='{}', "
"field_id={}, row_count={}, total_elements={}",
@@ -272,9 +257,10 @@ ArrayOffsetsSealed::BuildFromSegment(const void* segment,
row_count,
total_elements);
auto result = std::make_shared<ArrayOffsetsSealed>(
std::move(element_row_ids), std::move(row_to_element_start));
result->resource_size_ = 4 * (row_count + 1) + 4 * total_elements;
auto result =
std::make_shared<ArrayOffsetsSealed>(std::move(row_to_element_start));
// Memory usage: only row_to_element_start_ (4 bytes per entry)
result->resource_size_ = 4 * (row_count + 1);
cachinglayer::Manager::GetInstance().ChargeLoadedResource(
cachinglayer::ResourceUsage{result->resource_size_, 0});
return result;
@@ -283,10 +269,16 @@ ArrayOffsetsSealed::BuildFromSegment(const void* segment,
std::pair<int32_t, int32_t>
ArrayOffsetsGrowing::ElementIDToRowID(int32_t elem_id) const {
std::shared_lock lock(mutex_);
assert(elem_id >= 0 &&
elem_id < static_cast<int32_t>(element_row_ids_.size()));
int32_t row_id = element_row_ids_[elem_id];
// Compute elem_idx: elem_idx = elem_id - start_of_this_row
int64_t total_elements =
row_to_element_start_.empty() ? 0 : row_to_element_start_.back();
assert(elem_id >= 0 && elem_id < total_elements);
// Binary search: find the row where elem_id belongs
auto it = std::upper_bound(
row_to_element_start_.begin(), row_to_element_start_.end(), elem_id);
int32_t row_id = static_cast<int32_t>(
std::distance(row_to_element_start_.begin(), it) - 1);
int32_t elem_idx = elem_id - row_to_element_start_[row_id];
return {row_id, elem_idx};
}
@@ -325,12 +317,15 @@ ArrayOffsetsGrowing::RowBitsetToElementBitset(
TargetBitmap element_bitset(element_count);
TargetBitmap valid_element_bitset(element_count);
for (int64_t elem_id = element_start; elem_id < element_end; ++elem_id) {
auto row_id = element_row_ids_[elem_id];
int64_t bitset_idx = row_id - row_start;
element_bitset[elem_id - element_start] = row_bitset[bitset_idx];
valid_element_bitset[elem_id - element_start] =
valid_row_bitset[bitset_idx];
// Use row-based iteration (more efficient than element-based)
for (int64_t i = 0; i < row_count; ++i) {
int64_t row_id = row_start + i;
int64_t start = row_to_element_start_[row_id] - element_start;
int64_t end = row_to_element_start_[row_id + 1] - element_start;
if (start < end) {
element_bitset.set(start, end - start, row_bitset[i]);
valid_element_bitset.set(start, end - start, valid_row_bitset[i]);
}
}
return {std::move(element_bitset), std::move(valid_element_bitset)};
@@ -354,9 +349,9 @@ ArrayOffsetsGrowing::RowBitsetToElementOffsets(
if (selected_rows == 0) {
return element_offsets;
}
int64_t avg_elem_per_row =
static_cast<int64_t>(element_row_ids_.size()) /
(static_cast<int64_t>(row_to_element_start_.size()) - 1);
int64_t total_elements = row_to_element_start_.back();
int64_t avg_elem_per_row = total_elements / committed_row_count_;
element_offsets.reserve(selected_rows * avg_elem_per_row);
for (int64_t i = 0; i < row_count; ++i) {
@@ -383,11 +378,8 @@ ArrayOffsetsGrowing::RowOffsetsToElementOffsets(
return element_offsets;
}
int64_t avg_elem_per_row =
(committed_row_count_ > 0)
? static_cast<int64_t>(element_row_ids_.size()) /
committed_row_count_
: 1;
int64_t total_elements = row_to_element_start_.back();
int64_t avg_elem_per_row = total_elements / committed_row_count_;
element_offsets.reserve(row_offsets.size() * avg_elem_per_row);
for (auto row_id : row_offsets) {
@@ -436,33 +428,32 @@ ArrayOffsetsGrowing::Insert(int64_t row_id_start,
row_to_element_start_.reserve(row_id_start + count + 1);
// Estimate total elements needed and reserve capacity
int64_t total_elements = 0;
for (int64_t i = 0; i < count; ++i) {
total_elements += array_lengths[i];
}
element_row_ids_.reserve(element_row_ids_.size() + total_elements);
int32_t original_committed_count = committed_row_count_;
for (int64_t i = 0; i < count; ++i) {
int32_t row_id = row_id_start + i;
int32_t array_len = array_lengths[i];
if (row_id == committed_row_count_) {
// Get current total element count (from sentinel or compute)
int32_t current_total = row_to_element_start_.empty()
? 0
: row_to_element_start_.back();
// Record the start position for this row
// If sentinel exists at current position, overwrite it; otherwise push_back
if (row_to_element_start_.size() >
static_cast<size_t>(committed_row_count_)) {
row_to_element_start_[committed_row_count_] =
element_row_ids_.size();
// Sentinel exists, overwrite it with row start
row_to_element_start_[committed_row_count_] = current_total;
} else {
row_to_element_start_.push_back(element_row_ids_.size());
row_to_element_start_.push_back(current_total);
}
// Add row_id for each element (elem_idx computed on access)
for (int32_t j = 0; j < array_len; ++j) {
element_row_ids_.emplace_back(row_id);
// Update sentinel (new total after this row)
int32_t new_total = current_total + array_len;
if (row_to_element_start_.size() >
static_cast<size_t>(committed_row_count_ + 1)) {
row_to_element_start_[committed_row_count_ + 1] = new_total;
} else {
row_to_element_start_.push_back(new_total);
}
committed_row_count_++;
@@ -472,17 +463,6 @@ ArrayOffsetsGrowing::Insert(int64_t row_id_start,
}
DrainPendingRows();
// Update the sentinel (total element count) only if we committed new rows
if (committed_row_count_ > original_committed_count) {
if (row_to_element_start_.size() ==
static_cast<size_t>(committed_row_count_)) {
row_to_element_start_.push_back(element_row_ids_.size());
} else {
row_to_element_start_[committed_row_count_] =
element_row_ids_.size();
}
}
}
void
@@ -495,17 +475,27 @@ ArrayOffsetsGrowing::DrainPendingRows() {
const auto& pending = it->second;
// Get current total element count
int32_t current_total =
(committed_row_count_ > 0)
? row_to_element_start_[committed_row_count_]
: 0;
// If sentinel exists at current position, overwrite it; otherwise push_back
if (row_to_element_start_.size() >
static_cast<size_t>(committed_row_count_)) {
row_to_element_start_[committed_row_count_] =
element_row_ids_.size();
row_to_element_start_[committed_row_count_] = current_total;
} else {
row_to_element_start_.push_back(element_row_ids_.size());
row_to_element_start_.push_back(current_total);
}
for (int32_t j = 0; j < pending.array_len; ++j) {
element_row_ids_.emplace_back(static_cast<int32_t>(pending.row_id));
// Update sentinel for next row
int32_t new_total = current_total + pending.array_len;
if (row_to_element_start_.size() >
static_cast<size_t>(committed_row_count_ + 1)) {
row_to_element_start_[committed_row_count_ + 1] = new_total;
} else {
row_to_element_start_.push_back(new_total);
}
committed_row_count_++;
internal/core/src/common/ArrayOffsets.h
+6 -10
View File
@@ -16,6 +16,7 @@
#pragma once
#include <algorithm>
#include <cstdint>
#include <map>
#include <memory>
@@ -90,13 +91,11 @@ class ArrayOffsetsSealed : public IArrayOffsets {
friend class ArrayOffsetsTest;
public:
ArrayOffsetsSealed() : element_row_ids_(), row_to_element_start_({0}) {
ArrayOffsetsSealed() : row_to_element_start_({0}) {
}
ArrayOffsetsSealed(std::vector<int32_t> element_row_ids,
std::vector<int32_t> row_to_element_start)
: element_row_ids_(std::move(element_row_ids)),
row_to_element_start_(std::move(row_to_element_start)) {
explicit ArrayOffsetsSealed(std::vector<int32_t> row_to_element_start)
: row_to_element_start_(std::move(row_to_element_start)) {
AssertInfo(!row_to_element_start_.empty(),
"row_to_element_start must have at least one element");
}
@@ -113,7 +112,7 @@ class ArrayOffsetsSealed : public IArrayOffsets {
int64_t
GetTotalElementCount() const override {
return element_row_ids_.size();
return row_to_element_start_.empty() ? 0 : row_to_element_start_.back();
}
std::pair<int32_t, int32_t>
@@ -144,7 +143,6 @@ class ArrayOffsetsSealed : public IArrayOffsets {
BuildFromSegment(const void* segment, const FieldMeta& field_meta);
private:
const std::vector<int32_t> element_row_ids_;
const std::vector<int32_t> row_to_element_start_;
int64_t resource_size_{0};
};
@@ -165,7 +163,7 @@ class ArrayOffsetsGrowing : public IArrayOffsets {
int64_t
GetTotalElementCount() const override {
std::shared_lock lock(mutex_);
return element_row_ids_.size();
return row_to_element_start_.empty() ? 0 : row_to_element_start_.back();
}
std::pair<int32_t, int32_t>
@@ -202,8 +200,6 @@ class ArrayOffsetsGrowing : public IArrayOffsets {
DrainPendingRows();
private:
std::vector<int32_t> element_row_ids_;
std::vector<int32_t> row_to_element_start_;
// Number of rows committed (contiguous from 0)
internal/core/src/common/ArrayOffsetsTest.cpp
+270 -8
View File
@@ -40,9 +40,8 @@ TEST_F(ArrayOffsetsTest, SealedBasic) {
// row 0: 2 elements (elem 0, 1)
// row 1: 3 elements (elem 2, 3, 4)
// row 2: 1 element (elem 5)
ArrayOffsetsSealed offsets(
{0, 0, 1, 1, 1, 2}, // element_row_ids
{0, 2, 5, 6} // row_to_element_start (size = row_count + 1)
ArrayOffsetsSealed offsets({0, 2, 5, 6}
// row_to_element_start (size = row_count + 1)
);
// Test GetRowCount
@@ -103,8 +102,7 @@ TEST_F(ArrayOffsetsTest, SealedBasic) {
}
TEST_F(ArrayOffsetsTest, SealedRowBitsetToElementBitset) {
ArrayOffsetsSealed offsets({0, 0, 1, 1, 1, 2}, // element_row_ids
{0, 2, 5, 6} // row_to_element_start
ArrayOffsetsSealed offsets({0, 2, 5, 6} // row_to_element_start
);
// row_bitset: row 0 = true, row 1 = false, row 2 = true
@@ -136,9 +134,8 @@ TEST_F(ArrayOffsetsTest, SealedRowBitsetToElementBitset) {
TEST_F(ArrayOffsetsTest, SealedEmptyArrays) {
// Test with some rows having empty arrays
// row 1 and row 3 are empty
ArrayOffsetsSealed offsets({0, 0, 2, 2, 2}, // element_row_ids
{0, 2, 2, 5, 5} // row_to_element_start
// row 0: 2 elements, row 1: 0 elements, row 2: 3 elements, row 3: 0 elements
ArrayOffsetsSealed offsets({0, 2, 2, 5, 5} // row_to_element_start
);
EXPECT_EQ(offsets.GetRowCount(), 4);
@@ -431,3 +428,268 @@ TEST_F(ArrayOffsetsTest, LargeArrayLength) {
EXPECT_EQ(elem_idx, 9999);
}
}
TEST_F(ArrayOffsetsTest, SealedRowBitsetToElementOffsets) {
// row 0: 2 elements (elem 0, 1)
// row 1: 3 elements (elem 2, 3, 4)
// row 2: 1 element (elem 5)
ArrayOffsetsSealed offsets({0, 2, 5, 6});
// Select row 0 and row 2
TargetBitmap row_bitset(3);
row_bitset[0] = true;
row_bitset[1] = false;
row_bitset[2] = true;
TargetBitmapView view(row_bitset.data(), row_bitset.size());
auto elem_offsets = offsets.RowBitsetToElementOffsets(view, 0);
// Should return elements from row 0 (0, 1) and row 2 (5)
EXPECT_EQ(elem_offsets.size(), 3);
EXPECT_EQ(elem_offsets[0], 0);
EXPECT_EQ(elem_offsets[1], 1);
EXPECT_EQ(elem_offsets[2], 5);
}
TEST_F(ArrayOffsetsTest, SealedRowBitsetToElementOffsetsWithRowStart) {
// row 0: 2 elements (elem 0, 1)
// row 1: 3 elements (elem 2, 3, 4)
// row 2: 1 element (elem 5)
// row 3: 2 elements (elem 6, 7)
ArrayOffsetsSealed offsets({0, 2, 5, 6, 8});
// Select from row 1 onwards, select row 1 and row 3 (relative indices 0 and 2)
TargetBitmap row_bitset(3);
row_bitset[0] = true; // row 1
row_bitset[1] = false; // row 2
row_bitset[2] = true; // row 3
TargetBitmapView view(row_bitset.data(), row_bitset.size());
auto elem_offsets = offsets.RowBitsetToElementOffsets(view, 1);
// Should return elements from row 1 (2, 3, 4) and row 3 (6, 7)
EXPECT_EQ(elem_offsets.size(), 5);
EXPECT_EQ(elem_offsets[0], 2);
EXPECT_EQ(elem_offsets[1], 3);
EXPECT_EQ(elem_offsets[2], 4);
EXPECT_EQ(elem_offsets[3], 6);
EXPECT_EQ(elem_offsets[4], 7);
}
TEST_F(ArrayOffsetsTest, SealedRowBitsetToElementOffsetsEmpty) {
ArrayOffsetsSealed offsets({0, 2, 5, 6});
// No rows selected
TargetBitmap row_bitset(3, false);
TargetBitmapView view(row_bitset.data(), row_bitset.size());
auto elem_offsets = offsets.RowBitsetToElementOffsets(view, 0);
EXPECT_EQ(elem_offsets.size(), 0);
}
TEST_F(ArrayOffsetsTest, SealedRowOffsetsToElementOffsets) {
// row 0: 2 elements (elem 0, 1)
// row 1: 3 elements (elem 2, 3, 4)
// row 2: 1 element (elem 5)
ArrayOffsetsSealed offsets({0, 2, 5, 6});
// Select row 0 and row 2
FixedVector<int32_t> row_offsets = {0, 2};
auto elem_offsets = offsets.RowOffsetsToElementOffsets(row_offsets);
EXPECT_EQ(elem_offsets.size(), 3);
EXPECT_EQ(elem_offsets[0], 0);
EXPECT_EQ(elem_offsets[1], 1);
EXPECT_EQ(elem_offsets[2], 5);
}
TEST_F(ArrayOffsetsTest, SealedRowOffsetsToElementOffsetsEmpty) {
ArrayOffsetsSealed offsets({0, 2, 5, 6});
FixedVector<int32_t> row_offsets;
auto elem_offsets = offsets.RowOffsetsToElementOffsets(row_offsets);
EXPECT_EQ(elem_offsets.size(), 0);
}
TEST_F(ArrayOffsetsTest, SealedForEachRowElementRange) {
// row 0: 2 elements (elem 0, 1)
// row 1: 3 elements (elem 2, 3, 4)
// row 2: 1 element (elem 5)
ArrayOffsetsSealed offsets({0, 2, 5, 6});
// Predicate: return true if row has more than 1 element
auto predicate = [](int32_t elem_start, int32_t elem_end) {
return (elem_end - elem_start) > 1;
};
auto result = offsets.ForEachRowElementRange(predicate, 0, 3);
EXPECT_EQ(result.size(), 3);
EXPECT_TRUE(result[0]); // row 0: 2 elements > 1
EXPECT_TRUE(result[1]); // row 1: 3 elements > 1
EXPECT_FALSE(result[2]); // row 2: 1 element <= 1
}
TEST_F(ArrayOffsetsTest, SealedForEachRowElementRangeWithRowStart) {
// row 0: 2 elements, row 1: 3 elements, row 2: 1 element, row 3: 4 elements
ArrayOffsetsSealed offsets({0, 2, 5, 6, 10});
auto predicate = [](int32_t elem_start, int32_t elem_end) {
return (elem_end - elem_start) >= 3;
};
// Start from row 1, check 3 rows
auto result = offsets.ForEachRowElementRange(predicate, 1, 3);
EXPECT_EQ(result.size(), 3);
EXPECT_TRUE(result[0]); // row 1: 3 elements >= 3
EXPECT_FALSE(result[1]); // row 2: 1 element < 3
EXPECT_TRUE(result[2]); // row 3: 4 elements >= 3
}
TEST_F(ArrayOffsetsTest, SealedRowBitsetToElementBitsetWithRowStart) {
// row 0: 2 elements, row 1: 3 elements, row 2: 1 element, row 3: 2 elements
ArrayOffsetsSealed offsets({0, 2, 5, 6, 8});
// Start from row 1, select row 1 and row 3
TargetBitmap row_bitset(3);
row_bitset[0] = true; // row 1
row_bitset[1] = false; // row 2
row_bitset[2] = true; // row 3
TargetBitmap valid_row_bitset(3, true);
TargetBitmapView row_view(row_bitset.data(), row_bitset.size());
TargetBitmapView valid_view(valid_row_bitset.data(),
valid_row_bitset.size());
auto [elem_bitset, valid_elem_bitset] =
offsets.RowBitsetToElementBitset(row_view, valid_view, 1);
// Elements from row 1-3: elem 2,3,4,5,6,7 (6 elements)
EXPECT_EQ(elem_bitset.size(), 6);
// Row 1 (elem 2,3,4 -> indices 0,1,2) = true
EXPECT_TRUE(elem_bitset[0]);
EXPECT_TRUE(elem_bitset[1]);
EXPECT_TRUE(elem_bitset[2]);
// Row 2 (elem 5 -> index 3) = false
EXPECT_FALSE(elem_bitset[3]);
// Row 3 (elem 6,7 -> indices 4,5) = true
EXPECT_TRUE(elem_bitset[4]);
EXPECT_TRUE(elem_bitset[5]);
}
TEST_F(ArrayOffsetsTest, SealedRowBitsetToElementBitsetWithInvalidRows) {
// row 0: 2 elements, row 1: 3 elements, row 2: 1 element
ArrayOffsetsSealed offsets({0, 2, 5, 6});
TargetBitmap row_bitset(3);
row_bitset[0] = true;
row_bitset[1] = true;
row_bitset[2] = true;
// Row 1 is invalid (e.g., NULL)
TargetBitmap valid_row_bitset(3);
valid_row_bitset[0] = true;
valid_row_bitset[1] = false; // invalid
valid_row_bitset[2] = true;
TargetBitmapView row_view(row_bitset.data(), row_bitset.size());
TargetBitmapView valid_view(valid_row_bitset.data(),
valid_row_bitset.size());
auto [elem_bitset, valid_elem_bitset] =
offsets.RowBitsetToElementBitset(row_view, valid_view, 0);
EXPECT_EQ(elem_bitset.size(), 6);
EXPECT_EQ(valid_elem_bitset.size(), 6);
// Row 0 elements: valid
EXPECT_TRUE(valid_elem_bitset[0]);
EXPECT_TRUE(valid_elem_bitset[1]);
// Row 1 elements: invalid
EXPECT_FALSE(valid_elem_bitset[2]);
EXPECT_FALSE(valid_elem_bitset[3]);
EXPECT_FALSE(valid_elem_bitset[4]);
// Row 2 elements: valid
EXPECT_TRUE(valid_elem_bitset[5]);
}
TEST_F(ArrayOffsetsTest, GrowingRowBitsetToElementOffsets) {
ArrayOffsetsGrowing offsets;
std::vector<int32_t> lens = {2, 3, 1};
offsets.Insert(0, lens.data(), 3);
TargetBitmap row_bitset(3);
row_bitset[0] = true;
row_bitset[1] = false;
row_bitset[2] = true;
TargetBitmapView view(row_bitset.data(), row_bitset.size());
auto elem_offsets = offsets.RowBitsetToElementOffsets(view, 0);
EXPECT_EQ(elem_offsets.size(), 3);
EXPECT_EQ(elem_offsets[0], 0);
EXPECT_EQ(elem_offsets[1], 1);
EXPECT_EQ(elem_offsets[2], 5);
}
TEST_F(ArrayOffsetsTest, GrowingRowOffsetsToElementOffsets) {
ArrayOffsetsGrowing offsets;
std::vector<int32_t> lens = {2, 3, 1};
offsets.Insert(0, lens.data(), 3);
FixedVector<int32_t> row_offsets = {0, 2};
auto elem_offsets = offsets.RowOffsetsToElementOffsets(row_offsets);
EXPECT_EQ(elem_offsets.size(), 3);
EXPECT_EQ(elem_offsets[0], 0);
EXPECT_EQ(elem_offsets[1], 1);
EXPECT_EQ(elem_offsets[2], 5);
}
TEST_F(ArrayOffsetsTest, GrowingForEachRowElementRange) {
ArrayOffsetsGrowing offsets;
std::vector<int32_t> lens = {2, 3, 1};
offsets.Insert(0, lens.data(), 3);
auto predicate = [](int32_t elem_start, int32_t elem_end) {
return (elem_end - elem_start) > 1;
};
auto result = offsets.ForEachRowElementRange(predicate, 0, 3);
EXPECT_EQ(result.size(), 3);
EXPECT_TRUE(result[0]); // 2 > 1
EXPECT_TRUE(result[1]); // 3 > 1
EXPECT_FALSE(result[2]); // 1 <= 1
}
TEST_F(ArrayOffsetsTest, GrowingRowBitsetToElementBitsetWithRowStart) {
ArrayOffsetsGrowing offsets;
std::vector<int32_t> lens = {2, 3, 1, 2};
offsets.Insert(0, lens.data(), 4);
TargetBitmap row_bitset(3);
row_bitset[0] = true; // row 1
row_bitset[1] = false; // row 2
row_bitset[2] = true; // row 3
TargetBitmap valid_row_bitset(3, true);
TargetBitmapView row_view(row_bitset.data(), row_bitset.size());
TargetBitmapView valid_view(valid_row_bitset.data(),
valid_row_bitset.size());
auto [elem_bitset, valid_elem_bitset] =
offsets.RowBitsetToElementBitset(row_view, valid_view, 1);
EXPECT_EQ(elem_bitset.size(), 6);
EXPECT_TRUE(elem_bitset[0]);
EXPECT_TRUE(elem_bitset[1]);
EXPECT_TRUE(elem_bitset[2]);
EXPECT_FALSE(elem_bitset[3]);
EXPECT_TRUE(elem_bitset[4]);
EXPECT_TRUE(elem_bitset[5]);
}
internal/core/src/exec/operator/IterativeFilterNode.cpp
+19 -26
View File
@@ -89,11 +89,11 @@ insert_helper(milvus::SearchResult& search_result,
int& topk,
const bool large_is_better,
const FixedVector<float>& distances,
const FixedVector<int32_t>& offsets,
const int64_t nq_index,
const int64_t unity_topk,
const int i,
const IArrayOffsets* array_offsets = nullptr) {
int64_t doc_id,
std::optional<int32_t> elem_idx) {
auto pos = large_is_better
? find_binsert_position<true>(search_result.distances_,
nq_index * unity_topk,
@@ -104,17 +104,6 @@ insert_helper(milvus::SearchResult& search_result,
nq_index * unity_topk + topk,
distances[i]);
// For element-level: convert element_id to (doc_id, element_index)
int64_t doc_id;
int32_t elem_idx = -1;
if (array_offsets != nullptr) {
auto [doc, idx] = array_offsets->ElementIDToRowID(offsets[i]);
doc_id = doc;
elem_idx = idx;
} else {
doc_id = offsets[i];
}
if (topk > pos) {
std::memmove(&search_result.distances_[pos + 1],
&search_result.distances_[pos],
@@ -122,15 +111,15 @@ insert_helper(milvus::SearchResult& search_result,
std::memmove(&search_result.seg_offsets_[pos + 1],
&search_result.seg_offsets_[pos],
(topk - pos) * sizeof(int64_t));
if (array_offsets != nullptr) {
if (elem_idx.has_value()) {
std::memmove(&search_result.element_indices_[pos + 1],
&search_result.element_indices_[pos],
(topk - pos) * sizeof(int32_t));
}
}
search_result.seg_offsets_[pos] = doc_id;
if (array_offsets != nullptr) {
search_result.element_indices_[pos] = elem_idx;
if (elem_idx.has_value()) {
search_result.element_indices_[pos] = elem_idx.value();
}
search_result.distances_[pos] = distances[i];
++topk;
@@ -226,7 +215,8 @@ PhyIterativeFilterNode::GetOutput() {
FixedVector<int32_t> offsets;
FixedVector<float> distances;
FixedVector<int32_t> doc_offsets;
std::vector<int64_t> element_to_doc_mapping;
// For element-level: cache (doc_id, elem_idx) to avoid duplicate ElementIDToRowID calls
std::vector<std::pair<int32_t, int32_t>> element_to_doc_mapping;
std::unordered_map<int64_t, bool> doc_eval_cache;
std::unordered_set<int64_t> unique_doc_ids;
@@ -270,12 +260,13 @@ PhyIterativeFilterNode::GetOutput() {
if (element_level) {
// 1. Convert element_ids to doc_ids and do filter on those doc_ids
// 2. element_ids with doc_ids that pass the filter are what we interested in
// Cache both doc_id and elem_idx to avoid duplicate ElementIDToRowID calls
element_to_doc_mapping.reserve(offsets.size());
for (auto element_id : offsets) {
auto [doc_id, elem_index] =
auto [doc_id, elem_idx] =
array_offsets->ElementIDToRowID(element_id);
element_to_doc_mapping.push_back(doc_id);
element_to_doc_mapping.push_back({doc_id, elem_idx});
unique_doc_ids.insert(doc_id);
}
@@ -312,17 +303,17 @@ PhyIterativeFilterNode::GetOutput() {
}
for (size_t i = 0; i < offsets.size(); ++i) {
int64_t doc_id = element_to_doc_mapping[i];
auto [doc_id, elem_idx] = element_to_doc_mapping[i];
if (doc_eval_cache[doc_id]) {
insert_helper(search_result,
topk,
large_is_better,
distances,
offsets,
nq_index,
unity_topk,
i,
array_offsets.get());
doc_id,
elem_idx);
if (topk == unity_topk) {
break;
}
@@ -337,10 +328,11 @@ PhyIterativeFilterNode::GetOutput() {
topk,
large_is_better,
distances,
offsets,
nq_index,
unity_topk,
i);
i,
offsets[i],
std::nullopt);
if (topk == unity_topk) {
break;
}
@@ -355,10 +347,11 @@ PhyIterativeFilterNode::GetOutput() {
topk,
large_is_better,
distances,
offsets,
nq_index,
unity_topk,
i);
i,
offsets[i],
std::nullopt);
if (topk == unity_topk) {
break;
}
internal/core/src/segcore/InsertRecordOffsetOrderedArrayTest.cpp
+4 -12
View File
@@ -157,17 +157,13 @@ TYPED_TEST_P(TypedOffsetOrderedArrayTest, find_first_n_element) {
this->seal();
// Build ArrayOffsets: each doc has array_len elements
std::vector<int32_t> element_row_ids;
std::vector<int32_t> row_to_element_start = {0};
for (int doc = 0; doc < num; doc++) {
for (int e = 0; e < array_len; e++) {
element_row_ids.push_back(doc);
}
row_to_element_start.push_back(
static_cast<int32_t>((doc + 1) * array_len));
}
auto array_offsets = std::make_shared<ArrayOffsetsSealed>(
std::move(element_row_ids), std::move(row_to_element_start));
auto array_offsets =
std::make_shared<ArrayOffsetsSealed>(std::move(row_to_element_start));
int total_elements = num * array_len; // 15
@@ -275,17 +271,13 @@ TYPED_TEST_P(TypedOffsetOrderedArrayTest, find_first_n_element_has_more) {
// Build ArrayOffsets: each doc has array_len elements
// doc0: [elem0, elem1], doc1: [elem2, elem3], doc2: [elem4, elem5]
std::vector<int32_t> element_row_ids;
std::vector<int32_t> row_to_element_start = {0};
for (int doc = 0; doc < num; doc++) {
for (int e = 0; e < array_len; e++) {
element_row_ids.push_back(doc);
}
row_to_element_start.push_back(
static_cast<int32_t>((doc + 1) * array_len));
}
auto array_offsets = std::make_shared<ArrayOffsetsSealed>(
std::move(element_row_ids), std::move(row_to_element_start));
auto array_offsets =
std::make_shared<ArrayOffsetsSealed>(std::move(row_to_element_start));
int total_elements = num * array_len; // 6
internal/core/src/segcore/InsertRecordOffsetOrderedMapTest.cpp
+4 -13
View File
@@ -156,19 +156,14 @@ TYPED_TEST_P(TypedOffsetOrderedMapTest, find_first_n_element) {
}
// Build ArrayOffsets: each doc has array_len elements
// element_row_ids: [0,0,0, 1,1,1, 2,2,2, 3,3,3, 4,4,4]
// row_to_element_start: [0, 3, 6, 9, 12, 15]
std::vector<int32_t> element_row_ids;
std::vector<int32_t> row_to_element_start = {0};
for (int doc = 0; doc < num; doc++) {
for (int e = 0; e < array_len; e++) {
element_row_ids.push_back(doc);
}
row_to_element_start.push_back(
static_cast<int32_t>((doc + 1) * array_len));
}
auto array_offsets = std::make_shared<ArrayOffsetsSealed>(
std::move(element_row_ids), std::move(row_to_element_start));
auto array_offsets =
std::make_shared<ArrayOffsetsSealed>(std::move(row_to_element_start));
int total_elements = num * array_len; // 15
@@ -276,17 +271,13 @@ TYPED_TEST_P(TypedOffsetOrderedMapTest, find_first_n_element_has_more) {
}
// Build ArrayOffsets: each doc has array_len elements
std::vector<int32_t> element_row_ids;
std::vector<int32_t> row_to_element_start = {0};
for (int doc = 0; doc < num; doc++) {
for (int e = 0; e < array_len; e++) {
element_row_ids.push_back(doc);
}
row_to_element_start.push_back(
static_cast<int32_t>((doc + 1) * array_len));
}
auto array_offsets = std::make_shared<ArrayOffsetsSealed>(
std::move(element_row_ids), std::move(row_to_element_start));
auto array_offsets =
std::make_shared<ArrayOffsetsSealed>(std::move(row_to_element_start));
int total_elements = num * array_len; // 6