diff --git a/internal/core/src/index/BitmapIndex.cpp b/internal/core/src/index/BitmapIndex.cpp index 3ead065cc2..e5dc68cbbc 100644 --- a/internal/core/src/index/BitmapIndex.cpp +++ b/internal/core/src/index/BitmapIndex.cpp @@ -112,6 +112,7 @@ BitmapIndex::Build(size_t n, const T* data, const bool* valid_data) { } is_built_ = true; + ComputeByteSize(); } template @@ -168,6 +169,7 @@ BitmapIndex::BuildWithFieldData( proto::schema::DataType_Name(schema_.data_type()))); } is_built_ = true; + ComputeByteSize(); } template @@ -568,6 +570,7 @@ BitmapIndex::LoadWithoutAssemble(const BinarySet& binary_set, is_mmap_); is_built_ = true; + ComputeByteSize(); } template diff --git a/internal/core/src/index/BitmapIndex.h b/internal/core/src/index/BitmapIndex.h index b391958ec8..a6e2b6415e 100644 --- a/internal/core/src/index/BitmapIndex.h +++ b/internal/core/src/index/BitmapIndex.h @@ -114,6 +114,69 @@ class BitmapIndex : public ScalarIndex { return Count(); } + void + ComputeByteSize() override { + ScalarIndex::ComputeByteSize(); + int64_t total = this->cached_byte_size_; + + // valid_bitset_ + total += valid_bitset_.size_in_bytes(); + + if (is_mmap_) { + // mmap mode + total += mmap_size_; + // bitmap_info_map_ overhead (keys and roaring metadata in memory) + size_t num_entries = bitmap_info_map_.size(); + if constexpr (std::is_same_v) { + for (const auto& [key, bitmap] : bitmap_info_map_) { + total += key.capacity(); + } + } else { + total += num_entries * sizeof(T); + } + // roaring metadata + map node overhead per entry + total += num_entries * (sizeof(roaring::Roaring) + 40); + } else if (build_mode_ == BitmapIndexBuildMode::ROARING) { + // data_: map + for (const auto& [key, bitmap] : data_) { + if constexpr (std::is_same_v) { + total += key.capacity(); + } else { + total += sizeof(T); + } + total += bitmap.getSizeInBytes(); + // std::map red-black tree node overhead (~40 bytes per node) + total += 40; + } + } else { + // bitsets_: map + size_t num_entries = bitsets_.size(); + if (num_entries > 0) { + size_t bitset_bytes = bitsets_.begin()->second.size_in_bytes(); + total += num_entries * bitset_bytes; + if constexpr (std::is_same_v) { + for (const auto& [key, bitset] : bitsets_) { + total += key.capacity(); + } + } else { + total += num_entries * sizeof(T); + } + // std::map red-black tree node overhead (~40 bytes per node) + total += num_entries * 40; + } + } + + // offset cache + total += data_offsets_cache_.capacity() * + sizeof(typename decltype(data_offsets_cache_)::value_type); + total += bitsets_offsets_cache_.capacity() * + sizeof(typename decltype(bitsets_offsets_cache_)::value_type); + total += mmap_offsets_cache_.capacity() * + sizeof(typename decltype(mmap_offsets_cache_)::value_type); + + this->cached_byte_size_ = total; + } + IndexStatsPtr Upload(const Config& config = {}) override; diff --git a/internal/core/src/index/HybridScalarIndex.cpp b/internal/core/src/index/HybridScalarIndex.cpp index e541c6c08b..3defa2dc3d 100644 --- a/internal/core/src/index/HybridScalarIndex.cpp +++ b/internal/core/src/index/HybridScalarIndex.cpp @@ -351,6 +351,7 @@ HybridScalarIndex::Load(const BinarySet& binary_set, const Config& config) { index->Load(binary_set, config); is_built_ = true; + ComputeByteSize(); } template @@ -382,6 +383,7 @@ HybridScalarIndex::Load(milvus::tracer::TraceContext ctx, index->Load(ctx, config); is_built_ = true; + ComputeByteSize(); } template class HybridScalarIndex; diff --git a/internal/core/src/index/HybridScalarIndex.h b/internal/core/src/index/HybridScalarIndex.h index 79a7abce90..7b83c4a6e5 100644 --- a/internal/core/src/index/HybridScalarIndex.h +++ b/internal/core/src/index/HybridScalarIndex.h @@ -153,6 +153,16 @@ class HybridScalarIndex : public ScalarIndex { return internal_index_->Size(); } + void + ComputeByteSize() override { + ScalarIndex::ComputeByteSize(); + int64_t total = this->cached_byte_size_; + if (internal_index_) { + total += internal_index_->ByteSize(); + } + this->cached_byte_size_ = total; + } + const bool HasRawData() const override { if (field_type_ == proto::schema::DataType::Array) { diff --git a/internal/core/src/index/Index.h b/internal/core/src/index/Index.h index b045de7e1e..7e5c2fe876 100644 --- a/internal/core/src/index/Index.h +++ b/internal/core/src/index/Index.h @@ -91,6 +91,29 @@ class IndexBase { cell_size_ = cell_size; } + // Returns the memory usage in bytes that scales with data size (O(n)). + // Fixed overhead are minimal and thus not included. + // + // NOTE: This method returns a cached value computed by ComputeByteSize(). + // It is designed for SEALED SEGMENTS only, where the index is fully built + // or loaded and no more data will be added. For GROWING SEGMENTS with + // ongoing inserts, the cached value will NOT be updated automatically. + // ComputeByteSize() should only be called after Build() or Load() completes + // on sealed segments. + int64_t + ByteSize() const { + return cached_byte_size_; + } + + // Computes and caches the memory usage in bytes. + // Subclasses should override this method to calculate their specific memory usage + // and store the result in cached_byte_size_. + // This method should be called at the end of Build() or Load() for sealed segments. + virtual void + ComputeByteSize() { + cached_byte_size_ = 0; + } + protected: explicit IndexBase(IndexType index_type) : index_type_(std::move(index_type)) { @@ -98,6 +121,7 @@ class IndexBase { IndexType index_type_ = ""; cachinglayer::ResourceUsage cell_size_ = {0, 0}; + mutable int64_t cached_byte_size_ = 0; std::unique_ptr mmap_file_raii_; }; diff --git a/internal/core/src/index/InvertedIndexTantivy.cpp b/internal/core/src/index/InvertedIndexTantivy.cpp index 46d50cdb10..c1e99cf545 100644 --- a/internal/core/src/index/InvertedIndexTantivy.cpp +++ b/internal/core/src/index/InvertedIndexTantivy.cpp @@ -218,6 +218,7 @@ InvertedIndexTantivy::Load(milvus::tracer::TraceContext ctx, // the index is loaded in ram, so we can remove files in advance disk_file_manager_->RemoveIndexFiles(); } + ComputeByteSize(); } template @@ -564,6 +565,7 @@ InvertedIndexTantivy::BuildWithRawDataForUT(size_t n, wrapper_->create_reader(milvus::index::SetBitsetSealed); finish(); wrapper_->reload(); + ComputeByteSize(); } template diff --git a/internal/core/src/index/InvertedIndexTantivy.h b/internal/core/src/index/InvertedIndexTantivy.h index 750ac767e8..bdc0271b9d 100644 --- a/internal/core/src/index/InvertedIndexTantivy.h +++ b/internal/core/src/index/InvertedIndexTantivy.h @@ -190,9 +190,18 @@ class InvertedIndexTantivy : public ScalarIndex { return Count(); } - int64_t - ByteSize() const { - return wrapper_->index_size_bytes(); + void + ComputeByteSize() override { + ScalarIndex::ComputeByteSize(); + int64_t total = this->cached_byte_size_; + + // Tantivy index size + total += wrapper_->index_size_bytes(); + + // null_offset_: vector + total += null_offset_.capacity() * sizeof(size_t); + + this->cached_byte_size_ = total; } virtual const TargetBitmap diff --git a/internal/core/src/index/RTreeIndex.cpp b/internal/core/src/index/RTreeIndex.cpp index 159977aec4..10c1a220ae 100644 --- a/internal/core/src/index/RTreeIndex.cpp +++ b/internal/core/src/index/RTreeIndex.cpp @@ -221,6 +221,7 @@ RTreeIndex::Load(milvus::tracer::TraceContext ctx, const Config& config) { total_num_rows_ = wrapper_->count() + static_cast(null_offset_.size()); is_built_ = true; + ComputeByteSize(); LOG_INFO( "Loaded R-Tree index from {} with {} rows", path_, total_num_rows_); @@ -238,6 +239,7 @@ RTreeIndex::Build(const Config& config) { total_num_rows_ = wrapper_->count() + static_cast(null_offset_.size()); is_built_ = true; + ComputeByteSize(); } template @@ -279,6 +281,7 @@ RTreeIndex::BuildWithFieldData( wrapper_->bulk_load_from_field_data(field_datas, schema_.nullable()); total_num_rows_ = total_rows; is_built_ = true; + ComputeByteSize(); return; } } diff --git a/internal/core/src/index/RTreeIndex.h b/internal/core/src/index/RTreeIndex.h index 8c88d40475..0773b3dccc 100644 --- a/internal/core/src/index/RTreeIndex.h +++ b/internal/core/src/index/RTreeIndex.h @@ -143,6 +143,22 @@ class RTreeIndex : public ScalarIndex { return Count(); } + void + ComputeByteSize() override { + ScalarIndex::ComputeByteSize(); + int64_t total = this->cached_byte_size_; + + // null_offset_ vector + total += null_offset_.capacity() * sizeof(size_t); + + // wrapper_ (RTreeIndexWrapper) + if (wrapper_) { + total += wrapper_->ByteSize(); + } + + this->cached_byte_size_ = total; + } + // GIS-specific query methods /** * @brief Query candidates based on spatial operation diff --git a/internal/core/src/index/RTreeIndexWrapper.cpp b/internal/core/src/index/RTreeIndexWrapper.cpp index e1e0fe337b..9061836225 100644 --- a/internal/core/src/index/RTreeIndexWrapper.cpp +++ b/internal/core/src/index/RTreeIndexWrapper.cpp @@ -285,5 +285,22 @@ RTreeIndexWrapper::count() const { return static_cast(rtree_.size()); } +int64_t +RTreeIndexWrapper::ByteSize() const { + int64_t total = 0; + + // values_: vector where Value = std::pair + // Box = bg::model::box = 2 Points = 2 * 2 * sizeof(double) = 32 bytes + // Value = Box + int64_t = 32 + 8 = 40 bytes + total += values_.capacity() * sizeof(Value); + + // rtree_ internal structure (nodes, pointers, MBRs) + // R*-tree with max 16 entries per node has overhead per entry + // Estimated ~18 bytes per entry for internal tree structure + total += rtree_.size() * 18; + + return total; +} + // index/leaf capacity setters removed; not applicable for Boost rtree } // namespace milvus::index \ No newline at end of file diff --git a/internal/core/src/index/RTreeIndexWrapper.h b/internal/core/src/index/RTreeIndexWrapper.h index bac2454676..030e284e1e 100644 --- a/internal/core/src/index/RTreeIndexWrapper.h +++ b/internal/core/src/index/RTreeIndexWrapper.h @@ -97,6 +97,13 @@ class RTreeIndexWrapper { int64_t count() const; + /** + * @brief Get the estimated memory usage of the R-tree index + * @return Memory usage in bytes + */ + int64_t + ByteSize() const; + // Boost rtree does not use index/leaf capacities; keep only fill factor for // compatibility (no-op currently) diff --git a/internal/core/src/index/ScalarIndex.h b/internal/core/src/index/ScalarIndex.h index d56bddbf29..e549245ca3 100644 --- a/internal/core/src/index/ScalarIndex.h +++ b/internal/core/src/index/ScalarIndex.h @@ -147,7 +147,8 @@ class ScalarIndex : public IndexBase { index_type_ == milvus::index::HYBRID_INDEX_TYPE || index_type_ == milvus::index::INVERTED_INDEX_TYPE || index_type_ == milvus::index::MARISA_TRIE || - index_type_ == milvus::index::MARISA_TRIE_UPPER; + index_type_ == milvus::index::MARISA_TRIE_UPPER || + index_type_ == milvus::index::ASCENDING_SORT; } virtual int64_t diff --git a/internal/core/src/index/ScalarIndexSort.cpp b/internal/core/src/index/ScalarIndexSort.cpp index c9035555d2..644190cd24 100644 --- a/internal/core/src/index/ScalarIndexSort.cpp +++ b/internal/core/src/index/ScalarIndexSort.cpp @@ -100,6 +100,7 @@ ScalarIndexSort::Build(size_t n, const T* values, const bool* valid_data) { is_built_ = true; setup_data_pointers(); + ComputeByteSize(); } template @@ -143,6 +144,7 @@ ScalarIndexSort::BuildWithFieldData( is_built_ = true; setup_data_pointers(); + ComputeByteSize(); } template @@ -281,6 +283,7 @@ ScalarIndexSort::LoadWithoutAssemble(const BinarySet& index_binary, } is_built_ = true; + ComputeByteSize(); LOG_INFO("load ScalarIndexSort done, field_id: {}, is_mmap:{}", field_id_, diff --git a/internal/core/src/index/ScalarIndexSort.h b/internal/core/src/index/ScalarIndexSort.h index c2ab8ed3e1..d4fa899a18 100644 --- a/internal/core/src/index/ScalarIndexSort.h +++ b/internal/core/src/index/ScalarIndexSort.h @@ -116,6 +116,28 @@ class ScalarIndexSort : public ScalarIndex { return size_ == 0; } + void + ComputeByteSize() override { + ScalarIndex::ComputeByteSize(); + int64_t total = this->cached_byte_size_; + + // idx_to_offsets_: vector + total += idx_to_offsets_.capacity() * sizeof(int32_t); + + // valid_bitset_: TargetBitmap + total += valid_bitset_.size_in_bytes(); + + if (is_mmap_) { + // mmap mode: add mmap size and filepath + total += mmap_size_; + } else { + // memory mode: add data vector + total += data_.capacity() * sizeof(IndexStructure); + } + + this->cached_byte_size_ = total; + } + IndexStatsPtr Upload(const Config& config = {}) override; diff --git a/internal/core/src/index/StringIndexMarisa.cpp b/internal/core/src/index/StringIndexMarisa.cpp index 594f9d48ad..2d56f61420 100644 --- a/internal/core/src/index/StringIndexMarisa.cpp +++ b/internal/core/src/index/StringIndexMarisa.cpp @@ -57,6 +57,29 @@ StringIndexMarisa::Size() { return total_size_; } +void +StringIndexMarisa::ComputeByteSize() { + StringIndex::ComputeByteSize(); + int64_t total = cached_byte_size_; + + // Size of the trie structure (marisa trie uses io_size() for serialized/memory size) + total += trie_.io_size(); + + // str_ids_: vector + total += str_ids_.capacity() * sizeof(int64_t); + + // str_ids_to_offsets_: map> + for (const auto& [key, vec] : str_ids_to_offsets_) { + total += sizeof(size_t); // key + total += vec.capacity() * sizeof(size_t); // vector capacity + total += sizeof(std::vector); // vector object overhead + } + // Map node overhead (rough estimate: ~40 bytes per node for std::map) + total += str_ids_to_offsets_.size() * 40; + + cached_byte_size_ = total; +} + int64_t StringIndexMarisa::CalculateTotalSize() const { int64_t size = 0; @@ -135,6 +158,7 @@ StringIndexMarisa::BuildWithFieldData( built_ = true; total_size_ = CalculateTotalSize(); + ComputeByteSize(); } void @@ -161,6 +185,7 @@ StringIndexMarisa::Build(size_t n, built_ = true; total_size_ = CalculateTotalSize(); + ComputeByteSize(); } BinarySet @@ -247,6 +272,7 @@ StringIndexMarisa::LoadWithoutAssemble(const BinarySet& set, fill_offsets(); built_ = true; total_size_ = CalculateTotalSize(); + ComputeByteSize(); } void diff --git a/internal/core/src/index/StringIndexMarisa.h b/internal/core/src/index/StringIndexMarisa.h index 18bf0545b5..c07f87ca5a 100644 --- a/internal/core/src/index/StringIndexMarisa.h +++ b/internal/core/src/index/StringIndexMarisa.h @@ -35,6 +35,9 @@ class StringIndexMarisa : public StringIndex { int64_t Size() override; + void + ComputeByteSize() override; + BinarySet Serialize(const Config& config) override; diff --git a/internal/core/src/index/StringIndexSort.cpp b/internal/core/src/index/StringIndexSort.cpp index d546e28f04..00b32dd3f1 100644 --- a/internal/core/src/index/StringIndexSort.cpp +++ b/internal/core/src/index/StringIndexSort.cpp @@ -139,6 +139,7 @@ StringIndexSort::Build(size_t n, is_built_ = true; total_size_ = CalculateTotalSize(); + ComputeByteSize(); } void @@ -182,6 +183,7 @@ StringIndexSort::BuildWithFieldData( is_built_ = true; total_size_ = CalculateTotalSize(); + ComputeByteSize(); } BinarySet @@ -335,6 +337,7 @@ StringIndexSort::LoadWithoutAssemble(const BinarySet& binary_set, is_built_ = true; total_size_ = CalculateTotalSize(); + ComputeByteSize(); } const TargetBitmap @@ -413,6 +416,26 @@ StringIndexSort::CalculateTotalSize() const { return size; } +void +StringIndexSort::ComputeByteSize() { + StringIndex::ComputeByteSize(); + int64_t total = cached_byte_size_; + + // Common structures (always in memory) + // idx_to_offsets_: vector + total += idx_to_offsets_.capacity() * sizeof(int32_t); + + // valid_bitset_: TargetBitmap + total += valid_bitset_.size_in_bytes(); + + // Add impl-specific memory usage + if (impl_) { + total += impl_->ByteSize(); + } + + cached_byte_size_ = total; +} + void StringIndexSortMemoryImpl::BuildFromMap( std::map&& map, @@ -855,6 +878,36 @@ StringIndexSortMemoryImpl::Size() { return size; } +int64_t +StringIndexSortMemoryImpl::ByteSize() const { + int64_t total = 0; + + // unique_values_: vector + // sizeof(std::string) includes the SSO buffer + // For heap-allocated strings (capacity > SSO threshold), we need to add external buffer + const size_t sso_threshold = GetStringSSOThreshold(); + total += unique_values_.capacity() * sizeof(std::string); + for (const auto& str : unique_values_) { + // Only add capacity for heap-allocated strings (non-SSO) + if (str.capacity() > sso_threshold) { + total += str.capacity(); + } + } + + // posting_lists_: vector + // PostingList is folly::small_vector + // sizeof(PostingList) includes the inline buffer for 4 elements + total += posting_lists_.capacity() * sizeof(PostingList); + for (const auto& list : posting_lists_) { + // If the capacity exceeds inline capacity (4), it allocates on heap + if (list.capacity() > 4) { + total += list.capacity() * sizeof(uint32_t); + } + } + + return total; +} + StringIndexSortMmapImpl::~StringIndexSortMmapImpl() { if (mmap_data_ != nullptr && mmap_data_ != MAP_FAILED) { munmap(mmap_data_, mmap_size_); @@ -1147,4 +1200,10 @@ StringIndexSortMmapImpl::Size() { return mmap_size_; } +int64_t +StringIndexSortMmapImpl::ByteSize() const { + // mmap size (O(n) - the mapped index data) + return mmap_size_; +} + } // namespace milvus::index diff --git a/internal/core/src/index/StringIndexSort.h b/internal/core/src/index/StringIndexSort.h index a1e2c44ec0..998aadfcd5 100644 --- a/internal/core/src/index/StringIndexSort.h +++ b/internal/core/src/index/StringIndexSort.h @@ -123,6 +123,11 @@ class StringIndexSort : public StringIndex { int64_t Size() override; + // Computes and caches the total memory usage in bytes. + // For mmap mode, this includes both memory-resident structures and mmap size. + void + ComputeByteSize() override; + protected: int64_t CalculateTotalSize() const; @@ -199,6 +204,10 @@ class StringIndexSortImpl { virtual int64_t Size() = 0; + + // Returns the memory usage in bytes for this impl + virtual int64_t + ByteSize() const = 0; }; class StringIndexSortMemoryImpl : public StringIndexSortImpl { @@ -273,6 +282,9 @@ class StringIndexSortMemoryImpl : public StringIndexSortImpl { int64_t Size() override; + int64_t + ByteSize() const override; + private: // Helper method for binary search size_t @@ -387,6 +399,9 @@ class StringIndexSortMmapImpl : public StringIndexSortImpl { int64_t Size() override; + int64_t + ByteSize() const override; + private: // Binary search for a value size_t diff --git a/internal/core/src/index/Utils.h b/internal/core/src/index/Utils.h index cef0ab0f66..ad102bab65 100644 --- a/internal/core/src/index/Utils.h +++ b/internal/core/src/index/Utils.h @@ -234,4 +234,12 @@ void inline SetBitsetGrowing(void* bitset, } } +// Get the SSO (Small String Optimization) threshold for std::string. +// Strings with capacity <= this threshold store data inline (no heap allocation). +inline size_t +GetStringSSOThreshold() { + static const size_t threshold = std::string().capacity(); + return threshold; +} + } // namespace milvus::index