diff --git a/internal/core/src/common/ArrayOffsets.cpp b/internal/core/src/common/ArrayOffsets.cpp index 389362c957..1fc9dfc35e 100644 --- a/internal/core/src/common/ArrayOffsets.cpp +++ b/internal/core/src/common/ArrayOffsets.cpp @@ -38,8 +38,13 @@ std::pair 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( + 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 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 element_offsets; @@ -106,9 +112,7 @@ ArrayOffsetsSealed::RowBitsetToElementOffsets( return element_offsets; } - int64_t avg_elem_per_row = - static_cast(element_row_ids_.size()) / - (static_cast(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(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(std::vector{}, - std::vector{0}); + return std::make_shared(std::vector{0}); } FieldId field_id = field_meta.get_id(); auto data_type = field_meta.get_data_type(); - std::vector element_row_ids; // Size is row_count + 1, last element stores total_element_count std::vector 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(); 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( - 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(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 ArrayOffsetsGrowing::ElementIDToRowID(int32_t elem_id) const { std::shared_lock lock(mutex_); - assert(elem_id >= 0 && - elem_id < static_cast(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( + 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(element_row_ids_.size()) / - (static_cast(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(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(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(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(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(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(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(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_++; diff --git a/internal/core/src/common/ArrayOffsets.h b/internal/core/src/common/ArrayOffsets.h index 792a8a7e8c..3d699efe62 100644 --- a/internal/core/src/common/ArrayOffsets.h +++ b/internal/core/src/common/ArrayOffsets.h @@ -16,6 +16,7 @@ #pragma once +#include #include #include #include @@ -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 element_row_ids, - std::vector 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 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 @@ -144,7 +143,6 @@ class ArrayOffsetsSealed : public IArrayOffsets { BuildFromSegment(const void* segment, const FieldMeta& field_meta); private: - const std::vector element_row_ids_; const std::vector 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 @@ -202,8 +200,6 @@ class ArrayOffsetsGrowing : public IArrayOffsets { DrainPendingRows(); private: - std::vector element_row_ids_; - std::vector row_to_element_start_; // Number of rows committed (contiguous from 0) diff --git a/internal/core/src/common/ArrayOffsetsTest.cpp b/internal/core/src/common/ArrayOffsetsTest.cpp index 4e12cf6de8..533853a837 100644 --- a/internal/core/src/common/ArrayOffsetsTest.cpp +++ b/internal/core/src/common/ArrayOffsetsTest.cpp @@ -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 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 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 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 lens = {2, 3, 1}; + offsets.Insert(0, lens.data(), 3); + + FixedVector 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 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 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]); +} diff --git a/internal/core/src/exec/operator/IterativeFilterNode.cpp b/internal/core/src/exec/operator/IterativeFilterNode.cpp index c635f15b3c..f2dd890d64 100644 --- a/internal/core/src/exec/operator/IterativeFilterNode.cpp +++ b/internal/core/src/exec/operator/IterativeFilterNode.cpp @@ -89,11 +89,11 @@ insert_helper(milvus::SearchResult& search_result, int& topk, const bool large_is_better, const FixedVector& distances, - const FixedVector& 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 elem_idx) { auto pos = large_is_better ? find_binsert_position(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 offsets; FixedVector distances; FixedVector doc_offsets; - std::vector element_to_doc_mapping; + // For element-level: cache (doc_id, elem_idx) to avoid duplicate ElementIDToRowID calls + std::vector> element_to_doc_mapping; std::unordered_map doc_eval_cache; std::unordered_set 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; } diff --git a/internal/core/src/segcore/InsertRecordOffsetOrderedArrayTest.cpp b/internal/core/src/segcore/InsertRecordOffsetOrderedArrayTest.cpp index bdd45f24e3..3dcb5c6728 100644 --- a/internal/core/src/segcore/InsertRecordOffsetOrderedArrayTest.cpp +++ b/internal/core/src/segcore/InsertRecordOffsetOrderedArrayTest.cpp @@ -157,17 +157,13 @@ TYPED_TEST_P(TypedOffsetOrderedArrayTest, find_first_n_element) { this->seal(); // Build ArrayOffsets: each doc has array_len elements - std::vector element_row_ids; std::vector 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((doc + 1) * array_len)); } - auto array_offsets = std::make_shared( - std::move(element_row_ids), std::move(row_to_element_start)); + auto array_offsets = + std::make_shared(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 element_row_ids; std::vector 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((doc + 1) * array_len)); } - auto array_offsets = std::make_shared( - std::move(element_row_ids), std::move(row_to_element_start)); + auto array_offsets = + std::make_shared(std::move(row_to_element_start)); int total_elements = num * array_len; // 6 diff --git a/internal/core/src/segcore/InsertRecordOffsetOrderedMapTest.cpp b/internal/core/src/segcore/InsertRecordOffsetOrderedMapTest.cpp index 242c6684a4..537c99d7bf 100644 --- a/internal/core/src/segcore/InsertRecordOffsetOrderedMapTest.cpp +++ b/internal/core/src/segcore/InsertRecordOffsetOrderedMapTest.cpp @@ -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 element_row_ids; std::vector 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((doc + 1) * array_len)); } - auto array_offsets = std::make_shared( - std::move(element_row_ids), std::move(row_to_element_start)); + auto array_offsets = + std::make_shared(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 element_row_ids; std::vector 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((doc + 1) * array_len)); } - auto array_offsets = std::make_shared( - std::move(element_row_ids), std::move(row_to_element_start)); + auto array_offsets = + std::make_shared(std::move(row_to_element_start)); int total_elements = num * array_len; // 6